Repository: jax-ml/jax
Pull request #29539 (2025-06-18 11:59:27)

[debug_info] Improve debug info for loop constructs.
Several loop constructs (`fori_loop`, `while`) are implemented in terms of scan and we transform internally the user-provided loop body to match the API of scan. We want to make sure that the debug info matches the user-provided function, not the manufactured scan body.

To achieve this we construct the debug info for the user provided function and we attach it to the attribute `__fun_debug_info__` of the manufactured scan body.

Also, refactor for_loop.scan to trace its argument only once.

Previously, the `f` argument for `for_loop.scap` was traced twice, once to process the initial values and then again in the scan loop. Not only is this wastefull, but in the second tracing the `arg_names` part of the debug info is messed up. See, e.g., the change in the escaped tracer debug info in the `debug_info_test::test_grad_scan`.


Patch:
diff --git a/jax/_src/api_util.py b/jax/_src/api_util.py
index f8f124519..54ef461da 100644
--- a/jax/_src/api_util.py
+++ b/jax/_src/api_util.py
@@ -612,6 +612,9 @@ def debug_info(
 
   See docstring for linear_util.DebugInfo.
   """
+  res = getattr(fun, "__fun_debug_info__", None)
+  if res is not None:
+    return res
   if sourceinfo is None:
     sourceinfo = fun_sourceinfo(fun)
   if signature is None:
@@ -627,25 +630,15 @@ def fun_signature(fun: Callable) -> inspect.Signature | None:
   except (ValueError, TypeError):
     return None
 
-def save_wrapped_fun_sourceinfo(wrapper: Callable,
-                                wrapped: Callable | core.DebugInfo) -> None:
-  # Prefer this to functools.wraps because it does not create a reference to
-  # the wrapped function.
-  if isinstance(wrapped, core.DebugInfo):
-    func_src_info = wrapped.func_src_info
-  elif callable(wrapped):
-    func_src_info = fun_sourceinfo(wrapped)
-  else:
-    assert False, wrapped  # Unreachable
-  setattr(wrapper, "__fun_sourceinfo__", func_src_info)
+def save_wrapped_fun_debug_info(wrapper: Callable,
+                                dbg: core.DebugInfo) -> None:
+  setattr(wrapper, "__fun_debug_info__", dbg)
 
 _fun_name_re = re.compile(r"(?:<built-in function (\S+)>)")
 
 # TODO(mattjj): make this function internal to this module
 def fun_sourceinfo(fun: Callable) -> str:
   # See DebugInfo.fun_src_info
-  res = getattr(fun, "__fun_sourceinfo__", None)
-  if res is not None: return res
   while isinstance(fun, partial):
     fun = fun.func
   fun = inspect.unwrap(fun)
diff --git a/jax/_src/custom_batching.py b/jax/_src/custom_batching.py
index d1dca7d34..5e072a30f 100644
--- a/jax/_src/custom_batching.py
+++ b/jax/_src/custom_batching.py
@@ -295,7 +295,7 @@ def custom_vmap_jvp(primals, tangents, *,
       out_mutually_batched.store(out_batched)
       return out
 
-    api_util.save_wrapped_fun_sourceinfo(to_jvp, call.jaxpr.debug_info)
+    api_util.save_wrapped_fun_debug_info(to_jvp, call.jaxpr.debug_info)
     def to_vmap_over_extra_batched_dims(primals, tangents):
       return api.jvp(to_jvp, primals, tangents)
 
diff --git a/jax/_src/lax/control_flow/common.py b/jax/_src/lax/control_flow/common.py
index 3df82867e..9d790cca6 100644
--- a/jax/_src/lax/control_flow/common.py
+++ b/jax/_src/lax/control_flow/common.py
@@ -63,7 +63,7 @@ def _initial_style_open_jaxpr(fun: Callable,
 def _initial_style_jaxpr(fun: Callable,
                          in_tree: PyTreeDef,
                          in_avals: Sequence[core.AbstractValue],
-                         debug_info: core.DebugInfo):
+                         debug_info: core.DebugInfo) -> tuple[core.ClosedJaxpr, Sequence[Any], PyTreeDef]:
   jaxpr, consts, out_tree = _initial_style_open_jaxpr(
       fun, in_tree, in_avals, debug_info)
   closed_jaxpr = pe.close_jaxpr(pe.convert_constvars_jaxpr(jaxpr))
diff --git a/jax/_src/lax/control_flow/for_loop.py b/jax/_src/lax/control_flow/for_loop.py
index 5e6e19aef..7270d49a6 100644
--- a/jax/_src/lax/control_flow/for_loop.py
+++ b/jax/_src/lax/control_flow/for_loop.py
@@ -165,7 +165,7 @@ def scan(f: Callable[[Carry, X], tuple[Carry, Y]],
     raise TypeError("scan: f argument should be a callable.")
   if unroll < 1:
     raise ValueError("`unroll` must be a positive integer.")
-  xs_flat, xs_tree = tree_flatten(xs)
+  xs_flat, _ = tree_flatten(xs)
 
   try:
     lengths = [x.shape[0] for x in xs_flat]
@@ -196,27 +196,32 @@ def scan(f: Callable[[Carry, X], tuple[Carry, Y]],
   x_dtypes = [dtypes.canonicalize_dtype(x.dtype) for x in xs_flat]
   x_avals = tuple(map(core.ShapedArray, x_shapes, x_dtypes))
 
-  def _create_jaxpr(init):
-    init_flat = tree_leaves(init)
-    _, in_tree = tree_flatten((init, xs))
-    dbg = api_util.debug_info("scan", f, (init, xs), {})
-    carry_avals = tuple(map(core.get_aval, init_flat))
-    jaxpr, _, out_tree = _initial_style_jaxpr(
-        f, in_tree, carry_avals + x_avals, dbg)
-    return jaxpr, out_tree
-  jaxpr, out_tree = _create_jaxpr(init)
+  init_flat = tree_leaves(init)
+  _, in_tree = tree_flatten((init, xs))
+  dbg = api_util.debug_info("scan", f, (init, xs), {})
+  carry_avals = tuple(map(core.get_aval, init_flat))
+  jaxpr, consts, out_tree = _initial_style_jaxpr(
+      f, in_tree, carry_avals + x_avals, dbg)
+  del carry_avals
+
   _, ys_avals = tree_unflatten(out_tree, jaxpr.out_avals)
   ys = tree_map(lambda aval: lax.full([length, *aval.shape], 0, aval.dtype),
                 ys_avals)
   def for_body(i, refs):
     carry_refs, xs_refs, ys_refs = refs
-    carry = tree_map(lambda x: x[()], carry_refs)
-    x = tree_map(lambda x: x[i], xs_refs)
-    carry, y = f(carry, x)
-    tree_map(lambda c_ref, c: ref_set(c_ref, (), c), carry_refs, carry)
-    tree_map(lambda y_ref, y: ref_set(y_ref, (i,), y), ys_refs, y)
+    flat_carry_refs = tree_leaves(carry_refs)
+    flat_xs_refs = tree_leaves(xs_refs)
+    flat_ys_refs = tree_leaves(ys_refs)
+    flat_outs = core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts,
+                                * consts,
+                                * map(lambda x: x[()], flat_carry_refs),
+                                * map(lambda x: x[i], flat_xs_refs))
+    flat_carry_out, flat_ys_out = split_list(flat_outs, [len(flat_carry_refs)])
+    map(lambda c_ref, c: ref_set(c_ref, (), c), flat_carry_refs, flat_carry_out)
+    map(lambda y_ref, y: ref_set(y_ref, (i,), y), flat_ys_refs, flat_ys_out)
+
   assert isinstance(length, int)
-  api_util.save_wrapped_fun_sourceinfo(for_body, f)
+  api_util.save_wrapped_fun_debug_info(for_body, dbg)
   init, _, ys = for_loop(length, for_body, (init, xs, ys), reverse=reverse,
                          unroll=unroll)
   return init, ys
diff --git a/jax/_src/lax/control_flow/loops.py b/jax/_src/lax/control_flow/loops.py
index b68401ddf..3b4920dce 100644
--- a/jax/_src/lax/control_flow/loops.py
+++ b/jax/_src/lax/control_flow/loops.py
@@ -2292,20 +2292,28 @@ def _fori_cond_fun(loop_carry):
   return lax.lt(i, upper)
 
 @weakref_lru_cache
-def _fori_body_fun(body_fun):
+def _fori_body_fun(body_fun: Callable, body_fun_dbg: core.DebugInfo) -> Callable:
   body_fun_ref = weakref.ref(body_fun)
 
   def while_body_fun(loop_carry):
     i, upper, x = loop_carry
     return lax.add(i, lax._const(i, 1)), upper, body_fun_ref()(i, x)
+  api_util.save_wrapped_fun_debug_info(
+      while_body_fun,
+      body_fun_dbg._replace(arg_names=(body_fun_dbg.arg_names[0],
+                                       "",  # upper,
+                                       * body_fun_dbg.arg_names[1:])))
   return while_body_fun
 
 @weakref_lru_cache
-def _fori_scan_body_fun(body_fun):
+def _fori_scan_body_fun(body_fun: Callable, body_fun_dbg: core.DebugInfo) -> Callable:
   body_fun_ref = weakref.ref(body_fun)
   def scanned_fun(loop_carry, _):
     i, x = loop_carry
     return (i + 1, body_fun_ref()(i, x)), None
+  api_util.save_wrapped_fun_debug_info(
+      scanned_fun,
+      body_fun_dbg._replace(arg_names=body_fun_dbg.arg_names + ("",)))
   return scanned_fun
 
 @api_boundary
@@ -2406,6 +2414,9 @@ def fori_loop(lower, upper, body_fun, init_val,
   else:
     use_scan = False
 
+  body_fun_dbg = api_util.debug_info("fori_loop", body_fun,
+                                     (0, init_val), {})
+
   if use_scan:
     if unroll is None:
       unroll = False
@@ -2413,9 +2424,7 @@ def fori_loop(lower, upper, body_fun, init_val,
     if config.disable_jit.value and length == 0:
       # non-jit implementation of scan does not support length=0
       return init_val
-
-    scan_body = _fori_scan_body_fun(body_fun)
-    api_util.save_wrapped_fun_sourceinfo(scan_body, body_fun)
+    scan_body = _fori_scan_body_fun(body_fun, body_fun_dbg)
     (_, result), _ = scan(
         scan_body,
         (lower_, init_val),
@@ -2432,8 +2441,7 @@ def fori_loop(lower, upper, body_fun, init_val,
     lower = lax.convert_element_type(lower, dtype)  # type: ignore
   if upper_dtype != dtype:
     upper = lax.convert_element_type(upper, dtype)  # type: ignore
-  while_body_fun = _fori_body_fun(body_fun)
-  api_util.save_wrapped_fun_sourceinfo(while_body_fun, body_fun)
+  while_body_fun = _fori_body_fun(body_fun, body_fun_dbg)
   _, _, result = while_loop(_fori_cond_fun, while_body_fun,
                             (lower, upper, init_val))
   return result
diff --git a/jax/_src/pallas/core.py b/jax/_src/pallas/core.py
index 9b0a8810b..0900d36bf 100644
--- a/jax/_src/pallas/core.py
+++ b/jax/_src/pallas/core.py
@@ -502,7 +502,10 @@ class BlockSpec:
   ) -> BlockMapping:
     if self.index_map is None:
       index_map_func = default_index_map(len(array_aval.shape))
-      api_util.save_wrapped_fun_sourceinfo(index_map_func, default_index_map)
+      index_map_dbg = api_util.debug_info("pallas_call index_map",
+                                          default_index_map, (),{}
+                                          )._replace(arg_names=("",) * len(index_map_avals))
+      api_util.save_wrapped_fun_debug_info(index_map_func, index_map_dbg)
     else:
       index_map_func = self.index_map
     if self.block_shape is None:


Test patch:
diff --git a/tests/debug_info_test.py b/tests/debug_info_test.py
index 303b7e935..443e0f734 100644
--- a/tests/debug_info_test.py
+++ b/tests/debug_info_test.py
@@ -310,24 +310,19 @@ class DebugInfoTest(jtu.JaxTestCase):
     self.assertIsNotNone(dbg.func_lineno)
     self.assertEqual(dbg.arg_names, ("my_arg",))
 
-  def test_debug_info_save_wrapped_fun_source_info(self):
+  def test_debug_info_save_wrapped_fun_debug_info(self):
     def wrapper(x, y):
       return x
 
-    dbg = api_util.debug_info("test", wrapper, (1, 2), {})
-    self.assertEqual("wrapper", dbg.func_name)
-
-    api_util.save_wrapped_fun_sourceinfo(wrapper, lambda x, y: x)
-    dbg = api_util.debug_info("test", wrapper, (1, 2), {})
-    self.assertEqual("<lambda>", dbg.func_name)
-
     def other_f():
       pass
-    dbg_other = api_util.debug_info("test other", other_f, (), {})
-    api_util.save_wrapped_fun_sourceinfo(wrapper, dbg_other)
+
     dbg = api_util.debug_info("test", wrapper, (1, 2), {})
-    self.assertEqual("other_f", dbg.func_name)
-    self.assertEqual("test", dbg.traced_for)
+    self.assertEqual("wrapper", dbg.func_name)
+
+    api_util.save_wrapped_fun_debug_info(other_f, dbg)
+    dbg = api_util.debug_info("other", other_f, (1, 2), {})
+    self.assertEqual("wrapper", dbg.func_name)
 
   def test_debug_info_no_source_info_not_callable(self):
     # built-in function "int" does not have an inspect.Signature
@@ -1234,6 +1229,7 @@ class DebugInfoTest(jtu.JaxTestCase):
             "traced_for=jit, fun=<lambda>, arg_names=x, from None",
             "traced_for=linear_call fun, fun=fn, arg_names=r,x['c'], from r",
             "traced_for=linear_call fun, fun=fn, arg_names=r,x['c'], from x['c']",
+            # TODO(necula): should be from t['b']
             "traced_for=linear_call fun_transpose, fun=fn_tp, arg_names=r,t['c'], from t['c']",
         ]),
 
@@ -1369,10 +1365,11 @@ class DebugInfoTest(jtu.JaxTestCase):
         ])
 
   def test_grad_scan(self):
-      # Based on control_flow_test:testScanHigherOrderDifferentiation
+    # Based on control_flow_test:testScanHigherOrderDifferentiation
     tracer_spy = TracerSpy()
-    def f(c, a):
+    def f(c, a):  # c: f32, a: f32[2]
       tracer_spy.append(c)
+      tracer_spy.append(a)
       d = 0.75
       b = jnp.sin(c * jnp.sum(jnp.cos(d * a)))
       c = 0.9 * jnp.cos(d * jnp.sum(jnp.sin(c * a)))
@@ -1382,13 +1379,13 @@ class DebugInfoTest(jtu.JaxTestCase):
     c = jnp.array(1, dtype=as_.dtype)
 
     @jax.jit
-    def my_f(x, as_):
+    def my_f(x, as_):  # x: f32, as_: f32[3, 2]
       tracer_spy.append(x)
       def to_remat(a, b):
         return for_loop.scan(f, a, b)
       return jax.remat(to_remat)(c, as_)
 
-    def the_grad(c, as_):
+    def the_grad(c, as_):  # c: f32[], as_: f32[3, 2],
       tracer_spy.append(c)
       _, pullback = jax.vjp(my_f, c, as_)
       return pullback((c, np.arange(3, dtype=c.dtype)))
@@ -1396,14 +1393,12 @@ class DebugInfoTest(jtu.JaxTestCase):
     if config.use_direct_linearize.value:
       expected_jaxpr_debug_infos = [
           "traced_for=jit, fun=the_grad, arg_names=c,as_, result_paths=result[0],result[1]",
-
-          "traced_for=for_loop, fun=f, arg_names=,,, result_paths=,",
-          "traced_for=for_loop, fun=f, arg_names=i,refs[0],refs[1],refs[2], result_paths=",
+          "traced_for=scan, fun=f, arg_names=,,, result_paths=,",
+          "traced_for=scan, fun=f, arg_names=c,a, result_paths=",
           "traced_for=checkpoint / remat, fun=to_remat, arg_names=,,, result_paths=,",
-          "traced_for=for_loop, fun=f, arg_names=,,,,,, result_paths=,",
-          "traced_for=for_loop, fun=f, arg_names=i,refs[0],refs[1],refs[2], result_paths=",
-          "traced_for=for_loop, fun=f, arg_names=,,,,,,,,,,,,,,, result_paths=,",
-          "traced_for=for_loop, fun=f, arg_names=,,,,,,,,,,, result_paths=",
+          "traced_for=scan, fun=f, arg_names=,,,,,, result_paths=,",
+          "traced_for=scan, fun=f, arg_names=,,,,,,,,,,,,,,, result_paths=,",
+          "traced_for=scan, fun=f, arg_names=,,,,,,,,,,, result_paths=",
       ]
       if config.use_simplified_jaxpr_constants.value:
         expected_jaxpr_debug_infos.extend([
@@ -1418,13 +1413,12 @@ class DebugInfoTest(jtu.JaxTestCase):
     else:
       expected_jaxpr_debug_infos = [
           "traced_for=jit, fun=the_grad, arg_names=c,as_, result_paths=result[0],result[1]",
-          "traced_for=for_loop, fun=f, arg_names=,,, result_paths=,",
-          "traced_for=for_loop, fun=f, arg_names=i,refs[0],refs[1],refs[2], result_paths=",
+          "traced_for=scan, fun=f, arg_names=,,, result_paths=,",
           "traced_for=checkpoint / remat, fun=to_remat, arg_names=,,, result_paths=,",
-          "traced_for=for_loop, fun=f, arg_names=,,,,,, result_paths=,",
-          "traced_for=for_loop, fun=f, arg_names=i,refs[0],refs[1],refs[2], result_paths=",
-          "traced_for=for_loop, fun=f, arg_names=,,,,,,,,,,,,,,, result_paths=,",
-          "traced_for=for_loop, fun=f, arg_names=,,,,,,,,,,, result_paths=",
+          "traced_for=scan, fun=f, arg_names=,,,,,, result_paths=,",
+          "traced_for=scan, fun=f, arg_names=c,a, result_paths=",
+          "traced_for=scan, fun=f, arg_names=,,,,,,,,,,,,,,, result_paths=,",
+          "traced_for=scan, fun=f, arg_names=,,,,,,,,,,, result_paths=",
       ]
       if config.use_simplified_jaxpr_constants.value:
         expected_jaxpr_debug_infos.extend([
@@ -1444,17 +1438,14 @@ class DebugInfoTest(jtu.JaxTestCase):
         expected_tracer_debug_infos=[
             "traced_for=jit, fun=the_grad, arg_names=c,as_, from c",
             "traced_for=scan, fun=f, arg_names=c,a, from c",
+            "traced_for=scan, fun=f, arg_names=c,a, from a",
             "traced_for=jit, fun=my_f, arg_names=x,as_, from x",
-            # TODO(necula): arg_names, and "from x"
-            "traced_for=for_loop, fun=f, arg_names=i,refs[0],refs[1],refs[2], from refs[0]",
         ],
         expected_lowering_lines=[
             re.compile(r".*func.func public @main\(%arg0: tensor<f..> loc\(\"c\"\)"),
             re.compile(r".*func.func public @main\(.*, %arg1: tensor<3x2xf..> loc\(\"as_\"\)"),
             re.compile(r".*func.func public @main\(.* -> .*tensor<f..> {jax.result_info = \"result\[0\]\""),
             re.compile(r".*func.func public @main\(.* -> .*tensor<3x2xf..> {jax.result_info = \"result\[1\]\""),
-            # TODO(necula): unnamed function?
-            re.compile(r".*func.func private @None"),
         ])
 
   def test_while_loop(self):
@@ -1497,13 +1488,11 @@ class DebugInfoTest(jtu.JaxTestCase):
         tracer_spy=tracer_spy,
         expected_jaxpr_debug_infos=[
             "traced_for=jit, fun=<lambda>, arg_names=x, result_paths=result",
-            # TODO(necula): bad arg_names, result_paths
-            "traced_for=scan, fun=my_body, arg_names=loop_carry[0],loop_carry[1], result_paths=result[0][0],result[0][1]",
+            "traced_for=fori_loop, fun=my_body, arg_names=_,c,, result_paths=result[0][0],result[0][1]",
 
         ],
         expected_tracer_debug_infos=[
-            # TODO(necula): the arg_names are not right
-            "traced_for=scan, fun=my_body, arg_names=loop_carry[0],loop_carry[1], from loop_carry[1]",
+            "traced_for=fori_loop, fun=my_body, arg_names=_,c,, from c",
         ]
     )
 
@@ -1516,13 +1505,12 @@ class DebugInfoTest(jtu.JaxTestCase):
         tracer_spy=tracer_spy,
         expected_jaxpr_debug_infos=[
             "traced_for=jit, fun=<lambda>, arg_names=ub,x, result_paths=result",
+            # The fori_cond fun is entire manufactured internally
             re.compile(r"traced_for=while_cond, fun=_fori_cond_fun at .*loops.py:.*, arg_names=loop_carry\[0\],loop_carry\[1\],loop_carry\[2\], result_paths="),
-            # TODO(necula): arg_names and result_paths are not right
-            "traced_for=while_body, fun=my_body, arg_names=loop_carry[0],loop_carry[1],loop_carry[2], result_paths=result[0],result[1],result[2]",
+            "traced_for=fori_loop, fun=my_body, arg_names=_,,c, result_paths=result[0],result[1],result[2]",
         ],
         expected_tracer_debug_infos=[
-            # TODO(necula): the arg_names are not right
-            "traced_for=while_body, fun=my_body, arg_names=loop_carry[0],loop_carry[1],loop_carry[2], from loop_carry[2]",
+            "traced_for=fori_loop, fun=my_body, arg_names=_,,c, from c",
         ])
 
   def test_scan(self):
@@ -1756,7 +1744,7 @@ class DebugInfoTest(jtu.JaxTestCase):
         ],
     )
 
-    (x).block_until_ready()
+    x.block_until_ready()
 
   def test_remat(self):
     tracer_spy = TracerSpy()