[better_errors] Continue adding debug info to Jaxprs (step 7)

This follows in a series, starting with jax-ml#26078 and jax-ml#26313, adding debug_info to more calls to lu.wrap_init.

Fixes in jet, stateful code, key_reuse, ode, tests.

Author: gnecula

jax/_src/api_util.py

@@ -610,10 +610,17 @@ def fun_signature(fun: Callable) -> inspect.Signature | None:
   except (ValueError, TypeError):
     return None
 
-def save_wrapped_fun_sourceinfo(wrapper: Callable, wrapped: Callable) -> None:
+def save_wrapped_fun_sourceinfo(wrapper: Callable,
+                                wrapped: Callable | core.DebugInfo | None) -> None:
   # Prefer this to functools.wraps because it does not create a reference to
   # the wrapped function.
-  setattr(wrapper, "__fun_sourceinfo__", fun_sourceinfo(wrapped))
+  if isinstance(wrapped, core.DebugInfo):
+    func_src_info = wrapped.func_src_info
+  elif callable(wrapped):
+    func_src_info = fun_sourceinfo(wrapped)
+  else:
+    return
+  setattr(wrapper, "__fun_sourceinfo__", func_src_info)
 
 _fun_name_re = re.compile(r"(?:<built-in function (\S+)>)")
 

jax/_src/core.py

@@ -141,6 +141,8 @@ def __init__(self, constvars: Sequence[Var], invars: Sequence[Var],
     self._eqns = list(eqns)
     self._effects = effects
     self._debug_info = debug_info and debug_info.resolve_result_paths()
+    if debug_info is None:
+      assert False  # DO_NOT_SUBMIT
     # TODO(necula): re-enable these safety checks
     # assert (not debug_info or len(debug_info.arg_names) == len(invars)), (debug_info, invars)
     # assert (not debug_info or len(debug_info.result_paths) == len(outvars)), (debug_info, outvars)

jax/_src/interpreters/pxla.py

@@ -1824,7 +1824,7 @@ def _move_mutable_consts(
   invars = (*jaxpr.invars, *mutvars)
   effects = pe.make_jaxpr_effects(constvars, invars, jaxpr.outvars, jaxpr.eqns)
   jaxpr = core.Jaxpr(constvars, invars, jaxpr.outvars, jaxpr.eqns,
-                     effects, None)
+                     effects, closed_jaxpr.jaxpr.debug_info)
   return core.ClosedJaxpr(jaxpr, consts), in_mut
 
 @weakref_lru_cache

jax/_src/linear_util.py

@@ -166,6 +166,8 @@ def __init__(self, f: Callable,
     self.params = params
     self.in_type = in_type
     self.debug_info = debug_info
+    if debug_info is None:
+      assert False  # DO_NOT_SUBMIT
 
   @property
   def __name__(self):

jax/_src/state/discharge.py

@@ -460,12 +460,13 @@ def _addupdate_discharge(x, val, idx, tree):
   return _prepend_scatter(x, indexer, val, add=True)
 
 @weakref_lru_cache
-def _cached_closed_jaxpr_discharge(closed_jaxpr):
+def _cached_closed_jaxpr_discharge(closed_jaxpr: core.ClosedJaxpr):
   jaxpr, consts = closed_jaxpr.jaxpr, closed_jaxpr.consts
   num_outs = len(jaxpr.outvars)
   discharged_jaxpr, discharged_consts = discharge_state(jaxpr, consts)
   discharged_closed_jaxpr = core.ClosedJaxpr(discharged_jaxpr, discharged_consts)
-  fun = lu.wrap_init(core.jaxpr_as_fun(discharged_closed_jaxpr))
+  fun = lu.wrap_init(core.jaxpr_as_fun(discharged_closed_jaxpr),
+                     debug_info=discharged_jaxpr.debug_info)
   return discharged_closed_jaxpr, num_outs, fun
 
 @register_discharge_rule(core.closed_call_p)
@@ -598,7 +599,6 @@ def _convert_outputs_to_writes(
   assert not jaxpr.constvars, "Jaxpr shouldn't have constvars."
 
   in_avals = [v.aval for v in jaxpr.invars]
-  @lu.wrap_init
   def eval_jaxpr(*refs):
     # We split the refs into the original input refs and the dummy residual
     # refs.
@@ -610,14 +610,15 @@ def eval_jaxpr(*refs):
   res_ref_avals = [AbstractRef(v.aval) if not isinstance(v.aval, AbstractRef)
                    else v.aval for v in jaxpr.outvars]
   jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(
-      eval_jaxpr, [*in_avals, *res_ref_avals])
+      lu.wrap_init(eval_jaxpr,
+                   debug_info=jaxpr.debug_info),
+      [*in_avals, *res_ref_avals])
   assert not consts
   return jaxpr, [core.ShapedArray(a.shape, a.dtype) for a in res_ref_avals]
 
 def _convert_inputs_to_reads(num_res: int, jaxpr: core.Jaxpr) -> core.Jaxpr:
   assert not jaxpr.constvars, "Jaxpr should not have constvars"
 
-  @lu.wrap_init
   def eval_jaxpr(*refs):
     residual_refs, orig_refs = split_list(refs, [num_res])
     residual_vals = [r[...] for r in residual_refs]
@@ -629,7 +630,9 @@ def eval_jaxpr(*refs):
   res_ref_avals = [AbstractRef(aval) if not isinstance(aval, AbstractRef) else
                    aval for aval in res_val_avals]
   jaxpr, _, (), () = pe.trace_to_jaxpr_dynamic(
-      eval_jaxpr, [*res_ref_avals, *orig_ref_avals])
+      lu.wrap_init(eval_jaxpr,
+                   debug_info=jaxpr.debug_info),
+      [*res_ref_avals, *orig_ref_avals])
   return jaxpr
 
 def _run_state_partial_eval(trace: pe.JaxprTrace, *tracers: pe.JaxprTracer,
@@ -845,12 +848,13 @@ def _run_state_partial_eval_custom(
       *[v.aval for v in res_staged_invars], **staged_params)
   _, staged_outvars = partition_list(in_unknowns, eqn.outvars)
   if num_res:
-    @lu.wrap_init
+
     def staged(*args):
       out = run_state_p.bind(*args, **staged_params)
       return out[num_res:]
-    staged_call_jaxpr, _, (), () = pe.trace_to_jaxpr_dynamic(staged,
-                                                             [v.aval for v in res_staged_invars])
+    staged_call_jaxpr, _, (), () = pe.trace_to_jaxpr_dynamic(
+        lu.wrap_init(staged, debug_info=jaxpr_staged.debug_info),
+        [v.aval for v in res_staged_invars])
     eqn_staged = pe.new_jaxpr_eqn(res_staged_invars,
                                   staged_outvars,
                                   core.closed_call_p,
@@ -918,7 +922,9 @@ def trans(*args):
     ad.backward_pass(tangent_jaxpr, False, (), (*primals_args, *ct_args), ())
     return []
   jaxpr_trans, _, consts, () = pe.trace_to_jaxpr_dynamic(
-      lu.wrap_init(trans), [v.aval for v in jaxpr.invars])
+      lu.wrap_init(trans,
+                   debug_info=jaxpr.debug_info),
+      [v.aval for v in jaxpr.invars])
   return jaxpr_trans, consts
 
 def _run_state_transpose(in_cts, *args, jaxpr: core.Jaxpr,

jax/experimental/jet.py

@@ -60,6 +60,7 @@
 import numpy as np
 
 from jax import lax
+from jax import api_util
 import jax.numpy as jnp
 from jax.experimental import pjit
 from jax.tree_util import (register_pytree_node, tree_structure,
@@ -147,7 +148,9 @@ def flatten_fun_output(f, store, *args):
     store.store(tree)
     return ans
 
-  f, out_tree = flatten_fun_output(lu.wrap_init(fun))
+  f, out_tree = flatten_fun_output(
+      lu.wrap_init(fun,
+                   debug_info=api_util.debug_info("jet", fun, primals, {})))
   out_primals, out_terms = jet_fun(jet_subtrace(f), order).call_wrapped(primals, series)
   return tree_unflatten(out_tree(), out_primals), tree_unflatten(out_tree(), out_terms)
 
@@ -723,7 +726,7 @@ def _scatter_add_rule(primals_in, series_in, *, update_jaxpr, update_consts,
 def _jet_jaxpr(
     jaxpr: core.ClosedJaxpr, order: int, primals_and_series_avals, in_tree_def
 ) -> tuple[core.ClosedJaxpr, Any]:
-  f = lu.wrap_init(core.jaxpr_as_fun(jaxpr))
+  f = lu.wrap_init(core.jaxpr_as_fun(jaxpr), debug_info=jaxpr.jaxpr.debug_info)
   f_jet, out_tree_def = traceable(jet_fun(jet_subtrace(f), order), in_tree_def)
   jaxpr_jet, _, consts, () = pe.trace_to_jaxpr_dynamic(
       f_jet, primals_and_series_avals)

jax/experimental/key_reuse/_core.py

@@ -397,7 +397,10 @@ def is_consumed(var: core.Atom):
 def function_type_signature(fun: Callable[..., Any], *args: Any) -> KeyReuseSignature:
   args_flat, in_tree = tree_util.tree_flatten(args)
   in_avals_flat = [core.get_aval(arg) for arg in args_flat]
-  wrapped_fun, _ = api_util.flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
+  wrapped_fun, _ = api_util.flatten_fun_nokwargs(
+      lu.wrap_init(fun,
+                   debug_info=api_util.debug_info("key_reuse", fun, args, {})),
+      in_tree)
   jaxpr, _, _, () = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals_flat)
   return jaxpr_type_signature(jaxpr)
 

jax/experimental/ode.py

@@ -28,8 +28,10 @@
 
 from functools import partial
 import operator as op
+from typing import Callable
 
 import jax
+from jax import api_util
 import jax.numpy as jnp
 from jax._src import core
 from jax import custom_derivatives
@@ -44,8 +46,9 @@
 zip = safe_zip
 
 
-def ravel_first_arg(f, unravel):
-  return ravel_first_arg_(lu.wrap_init(f), unravel).call_wrapped
+def ravel_first_arg(f: Callable, unravel, debug_info: core.DebugInfo):
+  return ravel_first_arg_(lu.wrap_init(f, debug_info=debug_info),
+                          unravel).call_wrapped
 
 @lu.transformation2
 def ravel_first_arg_(f, unravel, y_flat, *args):
@@ -179,9 +182,10 @@ def odeint(func, y0, t, *args, rtol=1.4e-8, atol=1.4e-8, mxstep=jnp.inf, hmax=jn
   return _odeint_wrapper(converted, rtol, atol, mxstep, hmax, y0, t, *args, *consts)
 
 @partial(jax.jit, static_argnums=(0, 1, 2, 3, 4))
-def _odeint_wrapper(func, rtol, atol, mxstep, hmax, y0, ts, *args):
+def _odeint_wrapper(func: Callable, rtol, atol, mxstep, hmax, y0, ts, *args):
   y0, unravel = ravel_pytree(y0)
-  func = ravel_first_arg(func, unravel)
+  debug = api_util.debug_info("odeint", func, args, {})
+  func = ravel_first_arg(func, unravel, debug)
   out = _odeint(func, rtol, atol, mxstep, hmax, y0, ts, *args)
   return jax.vmap(unravel)(out)
 

tests/debug_info_test.py

@@ -881,7 +881,7 @@ def to_remat(x):
         tracer_spy=tracer_spy,
         expected_jaxpr_debug_infos=[
             # TODO(necula): what are these flat_index components?
-            "traced_for=jit, fun=apply_fn, arg_names=inp, result_paths=[0],[1][0][0][0][0][0]",
+            "traced_for=jit, fun=apply_fn, arg_names=inp, result_paths=[0],[1][<flat index 0>][0][<flat index 0>][0][0]",
             re.compile(r"traced_for=custom_jvp fun, fun=relu at .*nn.functions.py:.*, arg_names=x, result_paths="),
             re.compile(r"traced_for=jit, fun=relu at .*nn.functions.py:.*, arg_names=x, result_paths="),
         ],

tests/jaxpr_effects_test.py

@@ -17,6 +17,7 @@
 
 from absl.testing import absltest
 import jax
+from jax import api_util
 import jax.numpy as jnp
 from jax import lax
 from jax.experimental import pjit
@@ -178,12 +179,13 @@ class HigherOrderPrimitiveTest(jtu.JaxTestCase):
   def test_core_call_primitive_inherits_effects(self):
 
     def f(x):
-      @lu.wrap_init
       def f_(x):
         effect_p.bind(effect=foo_effect)
         effect_p.bind(effect=bar_effect)
         return [x]
-      return core.call(f_, x)[0]
+      dbg = api_util.debug_info("test", f_, (2.,), {})
+      return core.call(
+          lu.wrap_init(f_, debug_info=dbg), x)[0]
     jaxpr = jax.make_jaxpr(f)(2.)
     self.assertIn(foo_effect, jaxpr.jaxpr.effects)
     self.assertIn(bar_effect, jaxpr.jaxpr.effects)

tests/name_stack_test.py

@@ -15,6 +15,7 @@
 
 from absl.testing import absltest
 import jax
+from jax import api_util
 import jax.numpy as jnp
 from jax._src import core
 from jax import lax
@@ -85,11 +86,12 @@ def f(x):
   def test_call_primitive_jaxpr_should_not_store_outer_name_stack(self):
     @jax.named_scope('foo')
     def f(x):
-      @lu.wrap_init
       @jax.named_scope('bar')
       def _f(x):
         return [x + 1]
-      return core.call(_f, x)[0]
+      return core.call(lu.wrap_init(
+          _f,
+          debug_info=api_util.debug_info("test", _f, (0,), {})), x)[0]
 
     jaxpr = jax.make_jaxpr(f)(2).jaxpr
     self.assertEqual(str(jaxpr.eqns[0].params['call_jaxpr'].eqns[0].source_info.name_stack), 'bar')

tests/state_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 
 import jax
+from jax import api_util
 from jax._src import linear_util as lu
 from jax._src import test_util as jtu
 from jax._src import util
@@ -62,7 +63,10 @@ def kw_to_positional(f, store, factor, *args, **kwargs):
       store.store(aux_output)
       return (results[0:len(args)], dict(zip(kwargs_keys, results[len(args):])))
 
-    wf = lu.wrap_init(f)  # Wraps `f` as a `WrappedFun`.
+    # Wraps `f` as a `WrappedFun`.
+    wf = lu.wrap_init(
+        f,
+        debug_info=api_util.debug_info("test", f, (1, 2), dict(three=3, four=4)))
     wf, out_thunk = kw_to_positional(wf, 2)
     # Call the transformed function.
     scaled_positional, scaled_kwargs = wf.call_wrapped(1, 2, three=3, four=4)