[JAX] inspect_array: make name positional so custom_partitioning can resolve it

The prior commit that threaded a probe `name` through the
InspectPrimitive declared it as a keyword-only argument (`*, name`)
on `impl` / `partition` and left `impl_static_args = ()`. With
JAX's `custom_partitioning`, that breaks at trace time:

    TypeError: keyword arguments could not be resolved to positions

`register_primitive` wraps `cls.impl` as
`custom_partitioning(cls.impl, static_argnums=cls.impl_static_args)`,
and the wrapper's `__call__` runs `_resolve_kwargs(self.fun, args,
kwargs)` to push bind-time kwargs back into positional slots. A
keyword-only parameter has no positional slot to resolve into, so any
`outer_primitive.bind(x, ..., name=name)` call from inside a
`jax.jit` aborts before the FFI ever runs.

Follow the established TE pattern (e.g. `ActLuPrimitive`,
`FusedMoEAuxLossBwdPrimitive`):
  * Set `impl_static_args = (5,)` to declare position of `name`.
  * Drop the `*` separator on `impl` so `name` is
    positional-or-keyword; `_resolve_kwargs` can now push the
    bind kwarg to position 5.
  * Move static args to the head of `partition(name, mesh, ...)`
    and `shardy_sharding_rule(name, mesh, ...)` per
    `custom_partitioning`'s convention when `static_argnums`
    is set on the wrapped impl.
  * `abstract` and `lowering` keep `*, name` because they are
    called by JAX directly with bind kwargs and never go through
    `_resolve_kwargs`.

The user-facing API (`inspect_array(x, name)`) and the
`outer_primitive.bind(x, ..., name=name)` call site are unchanged.

Author: tdophung

transformer_engine/jax/debug/experimental/inspect.py

@@ -22,7 +22,13 @@ class InspectPrimitive(BasePrimitive):
 
     name = "te_inspect_ffi"
     multiple_results = False
-    impl_static_args = ()
+    # ``name`` is at positional index 5 in ``impl``; declaring it static
+    # here lets ``custom_partitioning`` resolve ``bind(..., name=...)``
+    # kwargs back to that position via ``_resolve_kwargs``. Keyword-only
+    # parameters (``*, name``) would raise
+    #     TypeError: keyword arguments could not be resolved to positions
+    # at trace time, so ``impl`` accepts ``name`` positionally below.
+    impl_static_args = (5,)
     inner_primitive = None
     outer_primitive = None
 
@@ -89,14 +95,15 @@ def impl(
         x_max,
         x_mean,
         x_std,
-        *,
         name,
     ):
         """
-        inspect implementation
+        inspect implementation. ``name`` is positional (not keyword-only)
+        so ``custom_partitioning(static_argnums=(5,))`` can resolve
+        ``bind(..., name=...)`` kwargs to position 5.
         """
         assert InspectPrimitive.inner_primitive is not None
-        (x) = InspectPrimitive.inner_primitive.bind(
+        x = InspectPrimitive.inner_primitive.bind(
             x,
             x_min,
             x_max,
@@ -107,13 +114,16 @@ def impl(
         return x
 
     @staticmethod
-    def partition(mesh, arg_infos, result_infos, *, name):
+    def partition(name, mesh, arg_infos, result_infos):
         """
         Identity in sharding: the output carries the same sharding as ``x``;
         the four scalar stats (x_min, x_max, x_mean, x_std) are fully
         replicated. Without this override the primitive falls back to
         ``BasePrimitive``'s abstract partition and any multi-device JIT
         rejects the call.
+
+        Static args precede ``mesh`` per the ``custom_partitioning``
+        convention when ``static_argnums`` is set on the wrapped impl.
         """
         del result_infos
         x_sharding = arg_infos[0].sharding
@@ -128,19 +138,20 @@ def partition(mesh, arg_infos, result_infos, *, name):
         out_sharding = x_sharding
 
         def sharded_impl(x, x_min, x_max, x_mean, x_std):
-            return InspectPrimitive.impl(x, x_min, x_max, x_mean, x_std, name=name)
+            return InspectPrimitive.impl(x, x_min, x_max, x_mean, x_std, name)
 
         return mesh, sharded_impl, out_sharding, arg_shardings
 
     @staticmethod
-    def shardy_sharding_rule(*args, **kwargs):
+    def shardy_sharding_rule(*args):
         """
         Five operands, one output. ``x`` and the output carry the same
         wildcard rank; the four scalar stats are rank-0 (empty operand
-        entries between commas). The ``name`` keyword attribute does not
-        participate in the rule.
+        entries between commas). ``name`` is a static arg (precedes
+        ``mesh``/``value_types``/``result_types`` in ``args``) and does
+        not participate in the rule.
         """
-        del args, kwargs
+        del args
         return "..., , , , -> ..."