Repository: jax-ml/jax
Pull request #27463 (2025-03-26 10:43:59)

[better_errors] Fix the handling of kwargs for debug_info.
kwargs are passed sorted by the actual kwarg keyword. This order must be accounted for when we construct the `debug_info.arg_names`.

Extended the tests to be more precise about not mixing up kwargs, e.g., use different shapes and look for the shape in the HLO.


Patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
index 55e2b2126..fb10245c3 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -192,7 +192,7 @@ def jit(
       constant).
 
       Static arguments should be hashable, meaning both ``__hash__`` and
-      ``__eq__`` are implemented, and immutable. Otherwise they can be arbitrary
+      ``__eq__`` are implemented, and immutable. Otherwise, they can be arbitrary
       Python objects. Calling the jitted function with different values for
       these constants will trigger recompilation. Arguments that are not
       array-like or containers thereof must be marked as static.
diff --git a/jax/_src/api_util.py b/jax/_src/api_util.py
index a42141b96..451d2e490 100644
--- a/jax/_src/api_util.py
+++ b/jax/_src/api_util.py
@@ -28,12 +28,12 @@ from jax._src.state.types import AbstractRef
 from jax._src.tree_util import (
     PyTreeDef, tree_flatten, tree_unflatten, tree_map,
     treedef_children, generate_key_paths, broadcast_prefix,
-    prefix_errors)
-from jax._src.tree_util import _replace_nones
+    prefix_errors, _replace_nones)
 from jax._src import linear_util as lu
 from jax._src.util import (safe_map, WrapKwArgs, Hashable, HashableFunction,
-                           Unhashable, safe_zip)
+                           Unhashable, safe_zip as zip)
 from jax._src import traceback_util
+
 traceback_util.register_exclusion(__file__)
 
 map = safe_map
@@ -201,9 +201,11 @@ def _validate_argnames(
                      f"in {argnames_name}. Function does not take these args.")
 
 
-def argnums_partial(f, dyn_argnums, args, require_static_args_hashable=True):
+def argnums_partial(f: lu.WrappedFun, dyn_argnums: int | Sequence[int],
+                    args: Sequence, require_static_args_hashable=True):
   dyn_argnums = _ensure_index_tuple(dyn_argnums)
   dyn_argnums = _ensure_inbounds(False, len(args), dyn_argnums)
+  fixed_args: list
   if require_static_args_hashable:
     fixed_args = []
     for i, arg in enumerate(args):
@@ -273,7 +275,9 @@ def argnums_partial_except(f: lu.WrappedFun, static_argnums: tuple[int, ...],
   return _argnums_partial(f, dyn_argnums, tuple(fixed_args)), dyn_args
 
 @lu.transformation2
-def _argnums_partial(_fun, _dyn_argnums, _fixed_args, *dyn_args, **kwargs):
+def _argnums_partial(_fun: Callable,
+                     _dyn_argnums: Sequence[int],
+                     _fixed_args: Sequence, *dyn_args, **kwargs):
   sentinel = object()
   args = [sentinel] * (len(_fixed_args) + len(dyn_args))
   for i, arg in zip(_dyn_argnums, dyn_args):
@@ -334,7 +338,7 @@ def donation_vector(donate_argnums, donate_argnames, in_tree,
     donate = bool(i in donate_argnums)
     res.extend((donate,) * arg.num_leaves)
   if kwargs_tree is not None:
-    for key, val in safe_zip(kwargs_tree.node_data()[1], kwargs_tree.children()):  # type: ignore
+    for key, val in zip(kwargs_tree.node_data()[1], kwargs_tree.children()):  # type: ignore
       donate = key in donate_argnames
       res.extend((donate,) * val.num_leaves)
   return tuple(res)
@@ -673,28 +677,45 @@ def _non_static_arg_names(fn_signature: inspect.Signature | None,
   top-level arguments. In other cases, including when the `args` and `kwargs`
   do not match the signature, we use names like `args[0[]`, `args[1]`, etc.
   """
+  # Use the same argument parsing as jit: positional followed by kwargs
+  # sorted by keys.
   static = object()
   static_argnums_ = _ensure_inbounds(True, len(args), static_argnums)
   static_argnames_ = set(static_argnames)
   args_ = [static if i in static_argnums_ else x for i, x in enumerate(args)]
-  kwargs_ = {k:static if k in static_argnames_ else x for k, x in kwargs.items()}
+  kwargs_ = {k: static if k in static_argnames_ else x for k, x in kwargs.items()}
+  ordered_args: Sequence[tuple[str, Any]] | None = None
   if fn_signature is not None:
     try:
       ba = fn_signature.bind(*args_, **kwargs_)
     except (ValueError, TypeError):
       pass
     else:
-      return tuple(f'{name}{lu._clean_keystr_arg_names(path)}'
-                   for name, x in ba.arguments.items()
-                   for path, l in generate_key_paths(x) if l is not static)
-  args_arg_names = tuple(f'args{lu._clean_keystr_arg_names(path)}'
-                         for path, l in generate_key_paths(args_)
-                         if l is not static)
-  kwargs_arg_names = tuple(f'kwargs{lu._clean_keystr_arg_names(path)}'
-                           for path, l in generate_key_paths(kwargs_)
-                           if l is not static)
-  arg_names = args_arg_names + kwargs_arg_names
-  return arg_names
+      # Do we have a **kwargs
+      kwargs_name = next((name for name, p in fn_signature.parameters.items()
+                          if p.kind == inspect.Parameter.VAR_KEYWORD), None)
+      # Positional argument are those not passed by keyword and not passed
+      # by **kwargs.
+      positional = [(name, x) for name, x in ba.arguments.items()
+                    if name not in kwargs and name != kwargs_name]
+      # Keyword arguments are passed sorted by actual kwarg keyword
+      sorted_kwargs = sorted(((name, x) for name, x in kwargs_.items()),
+                              key=lambda name_x: name_x[0])
+      sorted_kwargs = [(name if name in ba.arguments else f"{kwargs_name}['{name}']",
+                        x)
+                       for name, x in sorted_kwargs]
+      ordered_args = positional + sorted_kwargs
+
+  if ordered_args is None:
+    positional = [("args", args_)]
+    keyword = sorted([(f"kwargs['{name}']", x) for name, x in kwargs_.items() if x is not static],
+                     key=lambda name_x: name_x[0])
+    ordered_args = positional + keyword
+
+  return tuple(f'{name}{lu._clean_keystr_arg_names(path)}'
+               for name, x in ordered_args
+               for path, l in generate_key_paths(x) if l is not static)
+
 
 def hoist_obj_attrs(f, flat_args):
   idxs, objs, flat_args_ = [], [], []
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 5727c36a6..af744ae5d 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -683,7 +683,7 @@ class InferParamsCacheEntry:
 
 # We use an outer cache that is keyed on the signature of the arguments, but
 # when populating a cache entry using _infer_params_impl, we need to provide
-# actual arguments. In principle we could refactor _infer_params_impl to look
+# actual arguments. In principle, we could refactor _infer_params_impl to look
 # only at an argument signature instead of args/kwargs in those cases that we
 # cache, but this was a more minimal change.
 @util.weakref_lru_cache
@@ -730,7 +730,7 @@ def _infer_params_internal(
   if entry.pjit_params is None:
     p, args_flat = _infer_params_impl(
         fun, ji, ctx_mesh, dbg, args, kwargs, in_avals=avals)
-    if p.attrs_tracked:  # if attrs, don't popoulate the cache
+    if p.attrs_tracked:  # if attrs, don't populate the cache
       return p, p.consts + args_flat
     entry.pjit_params = p
   return entry.pjit_params, entry.pjit_params.consts + dynargs


Test patch:
diff --git a/tests/debug_info_test.py b/tests/debug_info_test.py
index 1d2935ea3..8ec5c42ef 100644
--- a/tests/debug_info_test.py
+++ b/tests/debug_info_test.py
@@ -46,6 +46,7 @@ from jax._src import test_util as jtu
 from jax._src.compilation_cache import is_persistent_cache_enabled
 from jax._src.lax.control_flow import for_loop
 from jax._src.interpreters import mlir
+from jax._src import util as util
 
 
 import numpy as np
@@ -241,7 +242,7 @@ class DebugInfoTest(jtu.JaxTestCase):
     dbg = api_util.debug_info("jit", my_f, (1, 2), dict(z=3, w=4))
     self.assertRegex(dbg.func_src_info, r"^my_f at .*debug_info_test.py:\d+")
     self.assertEqual(dbg.func_name, "my_f")
-    self.assertEqual(dbg.arg_names, ("x", "y", "z", "w"))
+    self.assertEqual(dbg.arg_names, ("x", "y", "w", "z"))
     self.assertIsNone(dbg.result_paths)
 
   def test_debug_info_arg_passed_as_kwarg(self):
@@ -261,23 +262,29 @@ class DebugInfoTest(jtu.JaxTestCase):
                                      "y_tree['w']", "y_tree['z']"))
 
   def test_debug_info_with_statics(self):
-    def my_f(x, y, *, z, w):
+    def my_f(x, z, *, w, y):
       pass
 
-    dbg = api_util.debug_info("jit", my_f, (1, 2), dict(z=3, w=4),
+    dbg = api_util.debug_info("jit", my_f, (1,), dict(y=2, z=3, w=4),
                               static_argnums=(1,),
                               static_argnames=("w",))
-    self.assertEqual(dbg.arg_names, ("x", "z"))
+    self.assertEqual(dbg.arg_names, ("x", "y", "z"))
 
   def test_debug_info_with_pytrees_and_statics(self):
-    def my_f(x, y, *, z, w):
+    def my_f(x, y, *, z, w, t):
       pass
 
     dbg = api_util.debug_info("jit", my_f, ((1, 2), (2, 3)),
-                              dict(z=(3, 4), w=(5, 6)),
+                              dict(z=(3, 4), w=(5, 6), t=7),
                               static_argnums=(1,),
                               static_argnames=("w",))
-    self.assertEqual(dbg.arg_names, ("x[0]", "x[1]", "z[0]", "z[1]"))
+    self.assertEqual(dbg.arg_names, ("x[0]", "x[1]", "t", "z[0]", "z[1]"))
+
+    dbg = api_util.debug_info("jit", my_f, ((1, 2),),
+                              dict(z=(3, 4), w=(5, 6), t=7, y=3),
+                              static_argnums=(1,),
+                              static_argnames=("w",))
+    self.assertEqual(dbg.arg_names, ("x[0]", "x[1]", "t", "y", "z[0]", "z[1]"))
 
   def test_debug_info_too_many_args(self):
     def my_f(x):
@@ -287,7 +294,7 @@ class DebugInfoTest(jtu.JaxTestCase):
     self.assertEqual(dbg.arg_names, ('args[0]', 'args[1]', 'args[2]', "kwargs['z']"))
 
   def test_debug_info_no_source_info_built_in(self):
-    # built-in function "int" does not have an inspect.Signature
+    # built-in function "max" does not have an inspect.Signature
     dbg = api_util.debug_info("jit", max, (1,), {})
     self.assertEqual(dbg.func_src_info, "max")
     self.assertEqual(dbg.arg_names, ("args[0]",))
@@ -761,6 +768,122 @@ class DebugInfoTest(jtu.JaxTestCase):
             re.compile(r".*func.func public @main\(.*\{jax.result_info = \"result\"\}"),
         ])
 
+  def test_jit_arg_names_with_out_of_order_kwargs(self):
+    tracer_spy = TracerSpy()
+
+    # The shapes are different, to differentiate them easily
+    a1 = (np.float32(0),)  # a hashable tuple, can be static
+    b2 = np.arange(2, dtype=np.float32)  # b2
+    z3 = np.arange(3, dtype=np.float32)
+    y4 = (np.float32(0.), np.float32(1.), np.float32(2.), np.float32(3.))
+    x5 = np.arange(5, dtype=np.float32)
+    u6 = np.arange(6, dtype=np.float32)
+    t7 = np.arange(7, dtype=np.float32)
+
+    def my_f(a1, b2, z3, y4, x5, *, u6, t7):
+      assert np.shape(a1[0]) == ()
+      assert np.shape(b2) == (2,)
+      assert np.shape(z3) == (3,)
+      assert np.shape(y4) == (4,)
+      assert np.shape(x5) == (5,)
+      assert np.shape(u6) == (6,)
+      assert np.shape(t7) == (7,)
+      tracer_spy.append(b2)
+      tracer_spy.append(x5)
+      return a1[0] + b2[0] + z3[0] + y4[0] + x5[0] + u6[0] + t7[0]
+
+    self._check_tracers_and_jaxprs(
+        jax.jit(my_f, static_argnums=(0,), static_argnames=("y4",)),
+        # Some positional args passed as keyword
+        a1, b2, x5=x5, y4=y4, z3=z3, t7=t7, u6=u6,
+        expected_jaxpr_debug_infos=[
+            "traced_for=jit, fun=my_f, arg_names=b2,t7,u6,x5,z3, result_paths=result",
+        ],
+        tracer_spy=tracer_spy,
+        expected_tracer_debug_infos=[
+            "traced_for=jit, fun=my_f, arg_names=b2,t7,u6,x5,z3, from b2",
+            "traced_for=jit, fun=my_f, arg_names=b2,t7,u6,x5,z3, from x5",
+        ],
+        expected_lowering_lines=[
+            re.compile(r".*func.func public @main\(%arg0: tensor<2xf..> loc\(\"b2\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg1: tensor<7xf..> loc\(\"t7\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg2: tensor<6xf..> loc\(\"u6\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg3: tensor<5xf..> loc\(\"x5\"\)"),
+        ]
+    )
+
+    tracer_spy.tracers = []
+    util.clear_all_caches()
+    self._check_tracers_and_jaxprs(
+        jax.jit(my_f, static_argnames=("y4",)),
+        # Positional argument y4 is static and passed by kwarg
+        a1, b2, z3, x5=x5, y4=y4, t7=t7, u6=u6,
+        expected_jaxpr_debug_infos=[
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, result_paths=result",
+        ],
+        tracer_spy=tracer_spy,
+        expected_tracer_debug_infos=[
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, from b2",
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, from x5",
+        ],
+        expected_lowering_lines=[
+            re.compile(r".*func.func public @main\(%arg0: tensor<f..> loc\(\"a1\[0\]\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg1: tensor<2xf..> loc\(\"b2\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg2: tensor<3xf..> loc\(\"z3\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg3: tensor<7xf..> loc\(\"t7\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg4: tensor<6xf..> loc\(\"u6\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg5: tensor<5xf..> loc\(\"x5\"\)"),
+        ]
+    )
+
+    tracer_spy.tracers = []
+    util.clear_all_caches()
+    self._check_tracers_and_jaxprs(
+        jax.jit(my_f, static_argnames=("y4",)),
+        # Positional argument y4 is static (declared as static_argnames)
+        a1, b2, z3, y4, x5=x5, t7=t7, u6=u6,
+        expected_jaxpr_debug_infos=[
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, result_paths=result",
+        ],
+        tracer_spy=tracer_spy,
+        expected_tracer_debug_infos=[
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, from b2",
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, from x5",
+        ],
+        expected_lowering_lines=[
+            re.compile(r".*func.func public @main\(%arg0: tensor<f..> loc\(\"a1\[0\]\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg1: tensor<2xf..> loc\(\"b2\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg2: tensor<3xf..> loc\(\"z3\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg3: tensor<7xf..> loc\(\"t7\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg4: tensor<6xf..> loc\(\"u6\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg5: tensor<5xf..> loc\(\"x5\"\)"),
+        ]
+    )
+
+    tracer_spy.tracers = []
+    util.clear_all_caches()
+    self._check_tracers_and_jaxprs(
+        jax.jit(my_f, static_argnums=(3,)),
+        # Positional argument y4 is static (declared as static_argnums)
+        a1, b2, z3, y4, x5=x5, t7=t7, u6=u6,
+        expected_jaxpr_debug_infos=[
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, result_paths=result",
+        ],
+        tracer_spy=tracer_spy,
+        expected_tracer_debug_infos=[
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, from b2",
+            "traced_for=jit, fun=my_f, arg_names=a1[0],b2,z3,t7,u6,x5, from x5",
+        ],
+        expected_lowering_lines=[
+            re.compile(r".*func.func public @main\(%arg0: tensor<f..> loc\(\"a1\[0\]\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg1: tensor<2xf..> loc\(\"b2\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg2: tensor<3xf..> loc\(\"z3\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg3: tensor<7xf..> loc\(\"t7\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg4: tensor<6xf..> loc\(\"u6\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg5: tensor<5xf..> loc\(\"x5\"\)"),
+        ]
+    )
+
   def test_jit_result_info(self):
     def f(x, y, z):
       return {'a': x, 'b': [y]}
@@ -1493,34 +1616,50 @@ class DebugInfoTest(jtu.JaxTestCase):
 
   def test_pmap_with_arg_and_result_names(self):
     tracer_spy = TracerSpy()
-    x = np.ones((jax.device_count(),), dtype=np.float32)
-    def my_f(x, y, *args, a, **kwargs):
-      # y and kwargs[c] is dead
+
+    # Use different shapes arguments to distinguish them in the HLO
+    def my_f(x0, y1, *args, b4, **kwargs):
+      assert np.shape(x0) == ()
+      assert np.shape(y1) == (1,)
+      assert np.shape(args[0]) == (2,)
+      assert np.shape(args[1]) == (3,)
+      assert np.shape(b4) == (4,)
+      assert np.shape(kwargs["a5"]) == (5,)
+      assert np.shape(kwargs["c6"]) == (6,)
+      # kwargs[b5] is dead
       tracer_spy.append(args[1])
-      s = x + a + args[1] + kwargs["d"]
-      return dict(u=s, v=x)
+      tracer_spy.append(b4)
+      tracer_spy.append(kwargs["c6"])
+      s0 = x0 + y1[0] + b4[0] + args[1][0] + kwargs["c6"][0]
+      return dict(v1=jnp.broadcast_to(s0, (1,)), u0=s0)
 
     self._check_tracers_and_jaxprs(
         jax.pmap(my_f, static_broadcasted_argnums=(0,)),
-        1., x, x, x,  # x, y, args[0], args[1]
-        d=x, a=x, b=x,  # kwargs
+        1.,  # x0
+        np.ones((jax.device_count(), 1), dtype=np.float32),  # y1
+        np.ones((jax.device_count(), 2), dtype=np.float32),  # args[0]
+        np.ones((jax.device_count(), 3), dtype=np.float32),  # args[1]
+        b4=np.ones((jax.device_count(), 4), dtype=np.float32),
+        a5=np.ones((jax.device_count(), 5), dtype=np.float32),
+        c6=np.ones((jax.device_count(), 6), dtype=np.float32),
         expected_jaxpr_debug_infos=[
-            "traced_for=pmap, fun=my_f, arg_names=y,args[0],args[1],a,kwargs['b'],kwargs['d'], result_paths=result['u'],result['v']",
+            "traced_for=pmap, fun=my_f, arg_names=y1,args[0],args[1],kwargs['a5'],b4,kwargs['c6'], result_paths=result['u0'],result['v1']",
         ],
         tracer_spy=tracer_spy,
         expected_tracer_debug_infos=[
-            "traced_for=pmap, fun=my_f, arg_names=y,args[0],args[1],a,kwargs['b'],kwargs['d'], from args[1]",
+            "traced_for=pmap, fun=my_f, arg_names=y1,args[0],args[1],kwargs['a5'],b4,kwargs['c6'], from args[1]",
+            "traced_for=pmap, fun=my_f, arg_names=y1,args[0],args[1],kwargs['a5'],b4,kwargs['c6'], from b4",
+            "traced_for=pmap, fun=my_f, arg_names=y1,args[0],args[1],kwargs['a5'],b4,kwargs['c6'], from kwargs['c6']",
         ],
         expected_lowering_lines=[
-            # TODO(necula): we did not DCE y?
-            re.compile(r".*func.func public @main\(.*%arg0: tensor<1xf..> loc\(\"y\"\)"),
-            re.compile(r".*func.func public @main\(.*%arg1: tensor<1xf..> loc\(\"args\[0\]\"\)"),
-            re.compile(r".*func.func public @main\(.*%arg2: tensor<1xf..> loc\(\"args\[1\]\"\)"),
-            re.compile(r".*func.func public @main\(.*%arg3: tensor<1xf..> loc\(\"a\"\)"),
-            re.compile(r".*func.func public @main\(.*%arg4: tensor<1xf..> loc\(\"kwargs\['b'\]\"\)"),
-            re.compile(r".*func.func public @main\(.*%arg5: tensor<1xf..> loc\(\"kwargs\['d'\]\"\)"),
-            re.compile(r".*func.func public @main\(.* -> .*\{jax.result_info = \"result\['u'\]\"\}"),
-            re.compile(r".*func.func public @main\(.* -> .*\{jax.result_info = \"result\['v'\]\"\}"),
+            re.compile(r".*func.func public @main\(.*%arg0: tensor<1x1xf..> loc\(\"y1\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg1: tensor<1x2xf..> loc\(\"args\[0\]\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg2: tensor<1x3xf..> loc\(\"args\[1\]\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg3: tensor<1x5xf..> loc\(\"kwargs\['a5'\]\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg4: tensor<1x4xf..> loc\(\"b4\"\)"),
+            re.compile(r".*func.func public @main\(.*%arg5: tensor<1x6xf..> loc\(\"kwargs\['c6'\]\"\)"),
+            re.compile(r".*func.func public @main\(.* -> .*\{jax.result_info = \"result\['u0'\]\"\}"),
+            re.compile(r".*func.func public @main\(.* -> .*\{jax.result_info = \"result\['v1'\]\"\}"),
         ]
     )