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

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

Here I changed the `custom_jvp_call` to replace the parameter
`jvp_jaxpr_thunk` (a callable) with `jvp_jaxpr_fun` (a `lu.WrappedFun`
that can carry debug info).

Also fixed uses in shard_map, checkify, sparse, attrs, and jax2tf.

Author: gnecula

jax/_src/api_util.py

@@ -27,7 +27,7 @@
 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, keystr, broadcast_prefix,
+    treedef_children, generate_key_paths, broadcast_prefix,
     prefix_errors)
 from jax._src.tree_util import _replace_nones
 from jax._src import linear_util as lu
@@ -664,12 +664,13 @@ def _non_static_arg_names(fn_signature: inspect.Signature | None,
     except (ValueError, TypeError):
       pass
     else:
-      return tuple(f'{name}{keystr(path)}' for name, x in ba.arguments.items()
+      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{keystr(path)}'
+  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{keystr(path)}'
+  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

jax/_src/checkify.py

@@ -833,7 +833,7 @@ def new_body_f(*c_consts_and_vals):
     # This checks if the next cond application will error
     _ = cond_f(*c_consts, *out)
     return out
-  new_body_f_ = lu.wrap_init(new_body_f)
+  new_body_f_ = lu.wrap_init(new_body_f, debug_info=body_jaxpr.jaxpr.debug_info)
   c_consts_avals = cond_jaxpr.in_avals[:c_consts_num]
   jaxpr, _, (), () = pe.trace_to_jaxpr_dynamic(new_body_f_, [*c_consts_avals,
                                                              *body_jaxpr.in_avals])
@@ -952,7 +952,8 @@ def remat_error_check(error, enabled_errors, *vals_in, jaxpr, **params):
 
 
 def shard_map_error_check(
-    error, enabled_errors, *vals_in, jaxpr, in_names, out_names, **kwargs
+    error: Error, enabled_errors, *vals_in,
+    jaxpr: core.Jaxpr, in_names, out_names, **kwargs
 ):
   if (mesh := kwargs.get('mesh')) is None:
     raise ValueError('Mesh must be provided for shard_map with checkify.')
@@ -976,7 +977,6 @@ def shard_map_error_check(
     )
   num_out_error_vals = out_tree.num_leaves - len(out_names)
 
-  @lu.wrap_init
   def expand_errors_leading_dim(*xs):
     outs = core.eval_jaxpr(checked_jaxpr.jaxpr, checked_jaxpr.consts, *xs)
     errs, outs = split_list(outs, [num_out_error_vals])
@@ -985,15 +985,18 @@ def expand_errors_leading_dim(*xs):
 
   with core.extend_axis_env_nd(mesh.shape.items()):
     jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(
-        expand_errors_leading_dim, checked_jaxpr.in_avals
+        lu.wrap_init(expand_errors_leading_dim,
+                     debug_info=checked_jaxpr.jaxpr.debug_info),
+        checked_jaxpr.in_avals
     )
   checked_jaxpr = core.ClosedJaxpr(jaxpr, consts)
 
   # Update shard_map params to account for extra error values.
   # Use fully sharded partitioning for out errors.
   new_out_names = (*([{0: mesh.axis_names}] * num_out_error_vals), *out_names)
   subfun = lu.hashable_partial(
-      lu.wrap_init(core.eval_jaxpr), checked_jaxpr.jaxpr, checked_jaxpr.consts
+      lu.wrap_init(core.eval_jaxpr, debug_info=checked_jaxpr.jaxpr.debug_info),
+      checked_jaxpr.jaxpr, checked_jaxpr.consts
   )
   new_params = dict(
       jaxpr=checked_jaxpr.jaxpr,
@@ -1007,8 +1010,10 @@ def expand_errors_leading_dim(*xs):
   return tree_unflatten(out_tree, err_and_out)
 error_checks[shard_map.shard_map_p] = shard_map_error_check
 
-def custom_jvp_call_rule(in_err, enabled_errors, *in_vals, num_consts,
-                         jvp_jaxpr_thunk, call_jaxpr, **params):
+def custom_jvp_call_rule(in_err: Error,
+                         enabled_errors: set, *in_vals, num_consts,
+                         jvp_jaxpr_fun: lu.WrappedFun,
+                         call_jaxpr: core.ClosedJaxpr, **params):
   # The types to have in mind are:
   #   jvp : (a -> b) -> (a, T a) -> (b, T b)
   #   checkify : (a -> b) -> a -> Err b
@@ -1021,10 +1026,11 @@ def custom_jvp_call_rule(in_err, enabled_errors, *in_vals, num_consts,
   err_vals, err_tree = jtu.tree_flatten(in_err)
   partial_checkify = lu.wrap_init(
       functools.partial(checkify_jaxpr_flat, call_jaxpr.jaxpr,
-                        call_jaxpr.consts, enabled_errors, err_tree))
+                        call_jaxpr.consts, enabled_errors, err_tree),
+      debug_info=call_jaxpr.jaxpr.debug_info)
   partial_checkify, f_metadata = _flatten_and_get_error_metadata_thunk(
       partial_checkify)
-  jvp = lift_jvp(err_tree.num_leaves, num_consts, jvp_jaxpr_thunk)
+  jvp = lift_jvp(err_tree.num_leaves, num_consts, jvp_jaxpr_fun)
   jvp, jvp_out_tree = flatten_fun_output(jvp)
   all_outs = custom_derivatives.custom_jvp_call_p.bind(
       partial_checkify, jvp, *err_vals, *in_vals, **params)
@@ -1041,17 +1047,17 @@ def custom_jvp_call_rule(in_err, enabled_errors, *in_vals, num_consts,
 
 # Compared to custom_derivatives.lift_jvp, we're handling the extra inputs and
 # outputs that checkify adds (just forwarding the error data's primal and
-# tangent components). The jaxpr in jvp_jaxpr_thunk doesn't expect those.
+# tangent components). The jaxpr in jvp_jaxpr_fun doesn't expect those.
 # TODO(mattjj): can we simplify this, or dedup with custom_derivatives.lift_jvp?
 # Adding another layer of lu.transformation was tricky, though maybe doable.
-def lift_jvp(num_errs, num_consts, jvp_jaxpr_thunk):
-  @lu.wrap_init
+def lift_jvp(num_errs: int, num_consts: int,
+             jvp_jaxpr_fun: lu.WrappedFun) -> lu.WrappedFun:
   def jvp(*xs):
     n, ragged = divmod(len(xs), 2)
     assert not ragged
     primals, tangents = xs[num_consts+num_errs:n], xs[n+num_consts+num_errs:]
     zeros = [type(t) is SymbolicZero for t in tangents]
-    jvp_jaxpr, jvp_consts, out_zeros = jvp_jaxpr_thunk(*zeros)
+    jvp_jaxpr, jvp_consts, out_zeros = jvp_jaxpr_fun.call_wrapped(*zeros)
     nonzero_tangents = [t for t in tangents if type(t) is not SymbolicZero]
     out = core.eval_jaxpr(jvp_jaxpr, jvp_consts, *primals, *nonzero_tangents)
     out_primals, nz_out_tangents = split_list(out, [len(out_zeros)])
@@ -1063,7 +1069,7 @@ def jvp(*xs):
     primal_errs = xs[num_consts:num_consts+num_errs]
     tangent_errs = xs[n+num_consts:n+num_consts+num_errs]
     return [*primal_errs, *out_primals, *tangent_errs, *out_tangents]
-  return jvp
+  return lu.wrap_init(jvp, debug_info=jvp_jaxpr_fun.debug_info)
 
 def custom_vjp_call_jaxpr_rule(in_err, enabled_errors, *in_vals,
                                fun_jaxpr: core.ClosedJaxpr,

jax/_src/custom_derivatives.py

@@ -378,20 +378,19 @@ def get_bind_params(self, params):
     new_params = dict(params)
     call_jaxpr: core.ClosedJaxpr = new_params.pop('call_jaxpr')
     num_consts: int = new_params.pop('num_consts')
-    jvp_jaxpr_thunk = new_params.pop('jvp_jaxpr_thunk')
+    jvp_jaxpr_fun = new_params.pop('jvp_jaxpr_fun')
     fun = lu.wrap_init(core.jaxpr_as_fun(call_jaxpr),
                        debug_info=call_jaxpr.jaxpr.debug_info)
-    jvp = lift_jvp(num_consts, jvp_jaxpr_thunk, call_jaxpr.jaxpr.debug_info)
+    jvp = lift_jvp(num_consts, jvp_jaxpr_fun)
     return [fun, jvp], new_params
 
-def lift_jvp(num_consts: int, jvp_jaxpr_thunk: Callable,
-             debug_info: core.DebugInfo | None) -> lu.WrappedFun:
+def lift_jvp(num_consts: int, jvp_jaxpr_fun: lu.WrappedFun) -> lu.WrappedFun:
   def jvp(*xs):
     n, ragged = divmod(len(xs), 2)
     assert not ragged
     primals, tangents = xs[num_consts:n], xs[n+num_consts:]
     zeros = [type(t) is SymbolicZero for t in tangents]
-    jvp_jaxpr, jvp_consts, out_zeros = jvp_jaxpr_thunk(*zeros)
+    jvp_jaxpr, jvp_consts, out_zeros = jvp_jaxpr_fun.call_wrapped(*zeros)
     nonzero_tangents = [t for t in tangents if type(t) is not SymbolicZero]
     out = core.eval_jaxpr(jvp_jaxpr, jvp_consts, *primals, *nonzero_tangents)
     out_primals, nz_out_tangents = split_list(out, [len(out_zeros)])
@@ -401,26 +400,26 @@ def jvp(*xs):
                     for p, z in zip(out_primals, out_zeros)]
     assert next(nz_out_tangents_, None) is None
     return [*out_primals, *out_tangents]
-  return lu.wrap_init(jvp, debug_info=debug_info)
+  return lu.wrap_init(jvp, debug_info=jvp_jaxpr_fun.debug_info)
 
 effects.custom_derivatives_allowed_effects.add_type(lax.InOutFeedEffect)
 
 custom_jvp_call_p = CustomJVPCallPrimitive('custom_jvp_call')
 
-def _custom_jvp_call_typecheck(_, *in_avals, call_jaxpr, jvp_jaxpr_thunk,
+def _custom_jvp_call_typecheck(_, *in_avals, call_jaxpr, jvp_jaxpr_fun,
                                num_consts, symbolic_zeros):
   # TODO(mattjj): could do more checking here...
-  del in_avals, jvp_jaxpr_thunk, num_consts
+  del in_avals, jvp_jaxpr_fun, num_consts
   disallowed_effects = effects.custom_derivatives_allowed_effects.filter_not_in(call_jaxpr.effects)
   if disallowed_effects:
     raise NotImplementedError(
         f'Effects not supported in `custom_jvp`: {disallowed_effects}')
   return call_jaxpr.out_avals, call_jaxpr.effects
 core.custom_typechecks[custom_jvp_call_p] = _custom_jvp_call_typecheck
 
-def _custom_jvp_call_mlir_translation(ctx, *args, call_jaxpr, jvp_jaxpr_thunk,
+def _custom_jvp_call_mlir_translation(ctx, *args, call_jaxpr, jvp_jaxpr_fun,
                                       num_consts, symbolic_zeros):
-  del jvp_jaxpr_thunk, num_consts, symbolic_zeros
+  del jvp_jaxpr_fun, num_consts, symbolic_zeros
   consts = mlir._ir_consts(call_jaxpr.consts)
   out, tokens = mlir.jaxpr_subcomp(ctx.module_context, call_jaxpr.jaxpr,
                                    ctx.name_stack, ctx.tokens_in, consts,
@@ -452,7 +451,7 @@ def _custom_jvp_call_dce(
     return [False] * len(eqn.invars), None
 
   call_jaxpr = eqn.params["call_jaxpr"]
-  jvp_jaxpr_thunk = eqn.params["jvp_jaxpr_thunk"]
+  jvp_jaxpr_fun = eqn.params["jvp_jaxpr_fun"]
   # We must set instantiate=True because some inputs that are unused by the
   # DCE'ed primal might be used in the JVP rule.
   dce_call_jaxpr, used_ins = _cached_closed_call_dce_instantiate(
@@ -461,7 +460,7 @@ def _custom_jvp_call_dce(
 
   @pe._memoize
   def dce_jvp_jaxpr_thunk(*in_zeros):
-    jvp_jaxpr, consts, out_zeros = jvp_jaxpr_thunk(*in_zeros)
+    jvp_jaxpr, consts, out_zeros = jvp_jaxpr_fun.call_wrapped(*in_zeros)
     dce_jvp_jaxpr, _ = pe.dce_jaxpr(jvp_jaxpr, [*used_outs, *used_outs], True)
     dce_out_zeros = [v for used, v in zip(used_outs, out_zeros) if used]
     return dce_jvp_jaxpr, consts, dce_out_zeros
@@ -470,7 +469,8 @@ def dce_jvp_jaxpr_thunk(*in_zeros):
   new_params = dict(
       eqn.params,
       call_jaxpr=dce_call_jaxpr,
-      jvp_jaxpr_thunk=dce_jvp_jaxpr_thunk,
+      jvp_jaxpr_fun=lu.wrap_init(dce_jvp_jaxpr_thunk,
+                                 debug_info=jvp_jaxpr_fun.debug_info)
   )
   new_eqn = pe.new_jaxpr_eqn(
       eqn.invars, outvars, eqn.primitive, new_params, dce_call_jaxpr.effects,

jax/_src/interpreters/ad.py

@@ -725,7 +725,9 @@ def f_tangent(*args):
 
     nz_tangents_in = [t for (t, nz) in zip(tangents, nzs_in) if nz]
     nz_tangents_out = call_primitive.bind_with_trace(
-        self.tangent_trace, (lu.wrap_init(f_tangent), *residuals, *nz_tangents_in), new_params)
+        self.tangent_trace, (lu.wrap_init(f_tangent,
+                                          debug_info=lin_jaxpr.debug_info),
+                             *residuals, *nz_tangents_in), new_params)
     nz_tangents_out_iter = iter(nz_tangents_out)
     tangents_out = [next(nz_tangents_out_iter) if nz else Zero.from_primal_value(primal)
                     for nz, primal in zip(nzs_out, primals_out)]

jax/_src/interpreters/partial_eval.py

@@ -502,7 +502,7 @@ def _closed_call_param_updater(params, _, __):
   return dict(params, call_jaxpr=core.ClosedJaxpr(jaxpr, ()))
 call_param_updaters[core.closed_call_p] = _closed_call_param_updater
 
-def abstract_eval_fun(fun, *avals, debug_info=None, **params):
+def abstract_eval_fun(fun: Callable, *avals, debug_info=None, **params):
   _, avals_out, _, () = trace_to_jaxpr_dynamic(
       lu.wrap_init(fun, params, debug_info=debug_info), avals)
   assert all(isinstance(aval, AbstractValue) for aval in avals_out)
@@ -1992,7 +1992,9 @@ def process_map(self, map_primitive, f: lu.WrappedFun,
       self.frame.add_eqn(eqn)
     return out_tracers
 
-  def process_custom_jvp_call(self, prim, fun, jvp, tracers, symbolic_zeros):
+  def process_custom_jvp_call(self, prim, fun: lu.WrappedFun,
+                              jvp: lu.WrappedFun, tracers,
+                              symbolic_zeros: bool):
     tracers = map(self.to_jaxpr_tracer, tracers)
     in_avals = [t.aval for t in tracers]
     in_tangent_avals = [t.to_tangent_aval() for t in in_avals]
@@ -2014,7 +2016,8 @@ def jvp_jaxpr_thunk(*in_zeros):
     outvars = map(self.makevar, out_tracers)
     eqn = new_jaxpr_eqn([*constvars, *invars], outvars, prim,
                         dict(call_jaxpr=closed_fun_jaxpr,
-                             jvp_jaxpr_thunk=jvp_jaxpr_thunk,
+                             jvp_jaxpr_fun=lu.wrap_init(jvp_jaxpr_thunk,
+                                                        debug_info=jvp.debug_info),
                              num_consts=len(consts),
                              symbolic_zeros=symbolic_zeros),
                         fun_jaxpr.effects,

jax/_src/lax/control_flow/loops.py

@@ -1255,7 +1255,9 @@ def arrange_jaxpr_args_for_wrapped(args):
   )
   # TODO(cperivol): avoid tracing the jaxpr twice. When doing so don't
   # forget to manage the effects.
-  new_jaxpr, _, (), () = pe.trace_to_jaxpr_dynamic(lu.wrap_init(wrapped), avals_for_wrapped_no_refs)
+  new_jaxpr, _, (), () = pe.trace_to_jaxpr_dynamic(
+      lu.wrap_init(wrapped, debug_info=discharged_jaxpr.debug_info),
+      avals_for_wrapped_no_refs)
   all_out = scan_p.bind(*args_for_wrapped,
                         jaxpr=core.ClosedJaxpr(new_jaxpr, ()),
                         length=length,
@@ -1922,9 +1924,9 @@ def new_body(*consts_refs_carry):
     carry, refs_out = split_list(carry_refs, [num_carry])
     return [*refs_out, *carry]
   new_body_jaxpr, _, new_body_consts, () = pe.trace_to_jaxpr_dynamic(
-      lu.wrap_init(new_body), [*remaining_body_const_avals, *[a.inner_aval for a
-                                                              in ref_avals],
-                               *carry_avals])
+      lu.wrap_init(new_body, debug_info=discharged_body_jaxpr.debug_info),
+      [*remaining_body_const_avals, *[a.inner_aval for a in ref_avals],
+      *carry_avals])
   if new_body_consts: raise NotImplementedError
 
   # Since some `Ref`s that were previously consts are now carries, we need to
@@ -1936,9 +1938,8 @@ def new_cond(*consts_refs_carry):
     del refs  # We don't use them here!
     return core.eval_jaxpr(cond_jaxpr, cond_jaxpr_consts, *consts, *carry)
   new_cond_jaxpr, _, new_cond_consts, () = pe.trace_to_jaxpr_dynamic(
-      lu.wrap_init(new_cond), [*cond_consts_avals,
-                               *[a.inner_aval for a in ref_avals],
-                               *carry_avals])
+      lu.wrap_init(new_cond, debug_info=cond_jaxpr.debug_info),
+      [*cond_consts_avals, *[a.inner_aval for a in ref_avals], *carry_avals])
   if new_cond_consts: raise NotImplementedError
 
   out = while_p.bind(*cond_consts, *remaining_body_consts, *refs, *carry,

jax/_src/lax/control_flow/solves.py

@@ -159,7 +159,8 @@ def _root_jvp(const_lengths, jaxprs, primals, tangents):
   linearize_and_solve = partial(
       core.jaxpr_as_fun(jaxprs.l_and_s), *params.l_and_s)
   f_at_solution = lambda *params: f(*params, *solution)
-  _, rhs = ad.jvp(lu.wrap_init(f_at_solution)).call_wrapped(
+  _, rhs = ad.jvp(lu.wrap_init(f_at_solution,
+                               debug_info=jaxprs.f.jaxpr.debug_info)).call_wrapped(
       params.f, params_dot.f)
   solution_dot = _map(
       operator.neg, linearize_and_solve(*solution, *rhs))

jax/_src/linear_util.py

@@ -65,13 +65,14 @@ def trans1(static_arg, *dynamic_args, **kwargs):
 
 from collections.abc import Callable, Sequence
 from functools import partial
+import re
 from typing import Any, NamedTuple
 import weakref
 
 from jax._src import config
 from jax._src import core
 from jax._src import traceback_util
-from jax._src.tree_util import keystr, generate_key_paths
+from jax._src.tree_util import keystr, KeyPath, generate_key_paths
 from jax._src.util import curry, cache_clearing_funs, HashableFunction
 
 
@@ -329,10 +330,16 @@ def wrap_init(f: Callable, params=None, *,
   return fun
 
 
+# We replace <flat index 0> with 0
+_re_clean_keystr_arg_names = re.compile(r"<flat index ([^>]+)>")
+def _clean_keystr_arg_names(k: KeyPath) -> str:
+  res = keystr(k)
+  return _re_clean_keystr_arg_names.sub(r"\1", res)
+
 @transformation_with_aux2
 def _get_result_paths_thunk(_fun: Callable, _store: Store, *args, **kwargs):
   ans = _fun(*args, **kwargs)
-  result_paths = [keystr(path) for path, _ in generate_key_paths(ans)]
+  result_paths = [_clean_keystr_arg_names(path) for path, _ in generate_key_paths(ans)]
   if _store:
     # In some instances a lu.WrappedFun is called multiple times, e.g.,
     # the bwd function in a custom_vjp

jax/_src/pallas/mosaic/lowering.py

@@ -3026,11 +3026,11 @@ def _custom_jvp_call_lowering_rule(
     ctx: LoweringRuleContext,
     *args,
     call_jaxpr: jax_core.Jaxpr,
-    jvp_jaxpr_thunk: Callable,
+    jvp_jaxpr_fun: lu.WrappedFun,
     num_consts: int,
     symbolic_zeros: bool,
 ):
-  del jvp_jaxpr_thunk
+  del jvp_jaxpr_fun
   if symbolic_zeros: raise NotImplementedError
   if num_consts: raise NotImplementedError
   if call_jaxpr.consts: raise NotImplementedError

jax/_src/pallas/pallas_call.py

@@ -242,7 +242,9 @@ def _block_map_function(new_idx, *args):
 
   with grid_mapping.trace_env():
     block_mapping_jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(
-        lu.wrap_init(_block_map_function), idx_avals)
+        lu.wrap_init(_block_map_function,
+                     debug_info=block_mapping.index_map_jaxpr.jaxpr.debug_info),
+        idx_avals)
   shape = block_mapping.block_shape
   if dim is batching.not_mapped:
     new_block_shape = shape

jax/_src/state/utils.py

@@ -73,7 +73,7 @@ def _hoist(*consts_args):
     return core.eval_jaxpr(jaxpr, all_consts, *args0, *args1)
 
   hoisted_jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(
-      lu.wrap_init(_hoist), in_avals)
+      lu.wrap_init(_hoist, debug_info=jaxpr.debug_info), in_avals)
   assert not consts, "All consts should have been converted to refs"
   return hoisted_jaxpr
 

jax/experimental/attrs.py

@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import Any, Callable
 
 from jax._src import core
 from jax._src import source_info_util
@@ -90,7 +90,9 @@ def jvp(f, primals, tangents, attr_tangents):
   primals_flat, in_tree = tree_flatten((attr_primals, *primals))
   tangents_flat, in_tree_ = tree_flatten((attr_tangents, *tangents))
   if in_tree != in_tree_: raise Exception
-  f_, out_tree = flatten_fun_nokwargs(_set_attrs(lu.wrap_init(f), attrs), in_tree)
+  dbg = api_util.debug_info("attrs_jvp", f, primals, {})
+  f_, out_tree = flatten_fun_nokwargs(
+      _set_attrs(lu.wrap_init(f, debug_info=dbg), attrs), in_tree)
   out_primals_flat, out_tangents_flat, tangent_attrs_out = _jvp(f_).call_wrapped(
       primals_flat, tangents_flat)
   out_primals = tree_unflatten(out_tree(), out_primals_flat)
@@ -151,12 +153,14 @@ def _getattr_jvp(trace, obj, attr):
 ad.LinearizeTrace.process_setattr = _setattr_jvp
 ad.LinearizeTrace.process_getattr = _getattr_jvp
 
-def linearize(f, *primals, attrs: list[tuple[Any, str]] = []):
+def linearize(f: Callable, *primals, attrs: list[tuple[Any, str]] = []):
   attr_primals = [jax_getattr(o, a) for o, a in attrs]
   attr_avals = [core.get_aval(p) for p in attr_primals]
   primals_flat, in_tree = tree_flatten(primals)
   tree = treedef_tuple((tree_structure(attr_primals), *in_tree.children()))
-  f_, out_tree = flatten_fun_nokwargs(_set_attrs(lu.wrap_init(f), attrs), tree)
+  dbg = api_util.debug_info("attrs linearize", f, primals, {})
+  f_, out_tree = flatten_fun_nokwargs(
+      _set_attrs(lu.wrap_init(f, debug_info=dbg), attrs), tree)
   primal_out, out_pvals, jaxpr, consts, attrs_out = _linearize(
       f_, *attr_primals, *primals_flat)
   f_lin = _lin_wrap(jaxpr, consts, out_pvals, attr_avals, (in_tree, out_tree()),
@@ -206,7 +210,9 @@ def vjp(f, *primals, attrs: list[tuple[Any, str]] = []):
   attr_primals = [jax_getattr(o, a) for o, a in attrs]
   primals_flat, in_tree = tree_flatten(primals)
   tree = treedef_tuple((tree_structure(attr_primals), *in_tree.children()))
-  f_, out_tree = flatten_fun_nokwargs(_set_attrs(lu.wrap_init(f), attrs), tree)
+  dbg = api_util.debug_info("attrs vjp", f, primals, {})
+  f_, out_tree = flatten_fun_nokwargs(
+      _set_attrs(lu.wrap_init(f, debug_info=dbg), attrs), tree)
   primal_out, out_pvals, jaxpr, consts, attrs_out = _linearize(
       f_, *attr_primals, *primals_flat)
   attr_avals = [core.get_aval(jax_getattr(o, a)).to_tangent_aval()

jax/experimental/jax2tf/jax2tf.py

@@ -599,6 +599,8 @@ def __init__(self, fun_jax, *,
     name_stack = util.wrap_name(fun_name, "jax2tf")
     self.name_stack = name_stack
     self.args_flat_tf = args_flat_tf
+    self.debug = api_util.debug_info("jax2tf", fun_jax,
+                                      args_specs, kwargs_specs)
 
   def before_conversion(self):
     prev_enable_xla = _thread_local_state.enable_xla
@@ -623,7 +625,10 @@ def _restore_context():
     dim_values, _ = _interpret_fun_jax(
         partial(shape_poly.compute_dim_vars_from_arg_shapes,
                 self.args_avals_flat, args_kwargs_tree=self.in_tree),
-        self.args_flat_tf, self.args_avals_flat, self.name_stack)
+        self.args_flat_tf, self.args_avals_flat, self.name_stack,
+        debug_info=api_util.debug_info("jax2tf dim_vars",
+                                       shape_poly.compute_dim_vars_from_arg_shapes,
+                                       self.args_specs, self.kwargs_specs))
 
     _thread_local_state.shape_env = zip(dim_vars, dim_values)
 
@@ -639,7 +644,8 @@ def run_fun_tf(self,
         fun_flat_jax,
         args_flat_tf, self.args_avals_flat,
         self.name_stack,
-        fresh_constant_cache=True)
+        fresh_constant_cache=True,
+        debug_info=self.debug)
     return outs_tf, self.outs_avals, out_tree_thunk()
 
   def get_vjp_fun(self) -> tuple[Callable,
@@ -849,10 +855,12 @@ def _interpret_fun_jax(
     fun_jax: Callable,
     args_tf: Sequence[TfVal],
     args_avals: Sequence[core.ShapedArray],
-    extra_name_stack: str | None,
+    extra_name_stack: str | None, *,
     fresh_constant_cache: bool = False,
+    debug_info: core.DebugInfo,
 ) -> tuple[tuple[TfVal, ...], tuple[core.ShapedArray, ...]]:
-  subtrace_fun = _interpret_subtrace(lu.wrap_init(fun_jax), args_avals)
+  subtrace_fun = _interpret_subtrace(
+      lu.wrap_init(fun_jax, debug_info=debug_info), args_avals)
   with _extended_name_stack(extra_name_stack):
     out_vals: Sequence[tuple[TfVal, core.ShapedArray]] = \
         _call_wrapped_with_new_constant_cache(subtrace_fun, args_tf,
@@ -1033,7 +1041,9 @@ def impl_multiple_results_jax(*args_jax):
 
     results_tf, _ = _interpret_fun_jax(
         impl_multiple_results_jax, args_tf, _in_avals,
-        extra_name_stack)
+        extra_name_stack,
+        debug_info=api_util.debug_info("jax2tf", impl_jax,
+                                       args_tf, kwargs))
     return results_tf if multiple_results else results_tf[0]
 
   return wrapped_tf
@@ -1066,7 +1076,8 @@ def _interpret_jaxpr(jaxpr: core.ClosedJaxpr, *args_tf: TfVal,
   """
   outs_tf, _ = _interpret_fun_jax(core.jaxpr_as_fun(jaxpr),
                                   args_tf, jaxpr.in_avals, extra_name_stack,
-                                  fresh_constant_cache=fresh_constant_cache)
+                                  fresh_constant_cache=fresh_constant_cache,
+                                  debug_info=jaxpr.jaxpr.debug_info)
   return outs_tf
 
 
@@ -1197,7 +1208,9 @@ def _eval_shape(shape: Sequence[shape_poly.DimSize], dtype=None) -> Sequence[TfV
   dim_vars, dim_values = util.unzip2(_thread_local_state.shape_env)
   shape_values_tf, _ = _interpret_fun_jax(
       partial(core.evaluate_shape, shape, dim_vars),
-      dim_values, [core.dim_value_aval()] * len(dim_values), "")  # type: ignore
+      dim_values, [core.dim_value_aval()] * len(dim_values), "",  # type: ignore
+      debug_info=api_util.debug_info("jax2tf evaluate_shape", core.evaluate_shape,
+                                     (0, 0, *dim_values), {}))
   # Keep only the non-constant dimensions
   return tuple(operator.index(d) if core.is_constant_dim(d) else d_tf  # type: ignore
                for d, d_tf in zip(shape, shape_values_tf))
@@ -3431,10 +3444,10 @@ def _tridiagonal_solve(*args: TfVal, _in_avals, _out_aval, **params):
 tf_impl_with_avals[lax.linalg.tridiagonal_solve_p] = _tridiagonal_solve
 
 def _custom_jvp_call(*args: TfVal, call_jaxpr: core.ClosedJaxpr,
-                           jvp_jaxpr_thunk: Callable,
+                           jvp_jaxpr_fun: Callable,
                            num_consts: int) -> Sequence[TfVal]:
   # TODO(necula): ensure that there is no AD transformation in scope
-  del jvp_jaxpr_thunk, num_consts
+  del jvp_jaxpr_fun, num_consts
   return _interpret_jaxpr(call_jaxpr, *args, extra_name_stack="custom_jvp",
                           fresh_constant_cache=False)
 

jax/experimental/jax2tf/tests/shape_poly_test.py

@@ -38,6 +38,7 @@
 import jax.numpy as jnp
 from jax import random
 from jax import tree_util
+from jax._src import api_util
 from jax._src import config
 from jax._src import core
 from jax._src import test_util as jtu
@@ -442,7 +443,10 @@ def f_tf(*args_tf):
             partial(shape_poly.compute_dim_vars_from_arg_shapes,
                     avals,
                     args_kwargs_tree=tree_util.tree_flatten((avals, {}))[1]),
-            args_tf, avals, "")
+            args_tf, avals, "",
+        debug_info=api_util.debug_info("jax2tf dim_vars",
+                                       shape_poly.compute_dim_vars_from_arg_shapes,
+                                       avals, {}))
         if expected_shapes is not None:
           expected_avals = tree_util.tree_map(
               lambda shape_str: core.ShapedArray(

jax/experimental/shard_map.py

@@ -1387,7 +1387,7 @@ def _closed_call_check(mesh, *in_rep, call_jaxpr, **kwargs):
 
 
 @register_check(custom_derivatives.custom_jvp_call_p)
-def _custom_jvp_call_check(mesh, *in_rep, call_jaxpr, jvp_jaxpr_thunk,
+def _custom_jvp_call_check(mesh, *in_rep, call_jaxpr, jvp_jaxpr_fun,
                            num_consts, symbolic_zeros):
   return _check_rep(mesh, call_jaxpr.jaxpr, in_rep)
 
@@ -1780,31 +1780,33 @@ def _partial_eval_jaxpr_custom_rule(
 pe.partial_eval_jaxpr_custom_rules[shard_map_p] = \
     _partial_eval_jaxpr_custom_rule
 
-def _add_reshapes(which, jaxpr_known, jaxpr_staged):
+def _add_reshapes(which: Sequence[bool],
+                  jaxpr_known: core.Jaxpr,
+                  jaxpr_staged: core.Jaxpr) -> tuple[core.Jaxpr, core.Jaxpr]:
   # add singleton axes to residuals which are from jaxpr_known and are scalars
   which_ = [w and not v.aval.shape  # pytype: disable=attribute-error
             for w, v in zip(which, jaxpr_staged.invars[:len(which)])]
   if not any(which_): return jaxpr_known, jaxpr_staged
   assert not jaxpr_known.constvars and not jaxpr_staged.constvars
 
-  @lu.wrap_init
   def known(*args):
     out = core.eval_jaxpr(jaxpr_known, (), *args)
     out_known, res = split_list(out, [len(out) - sum(which)])
     res = [_add_singleton(x) if not x.shape else x for x in res]
     return [*out_known, *res]
   avals_in = [v.aval for v in jaxpr_known.invars]
-  jaxpr_known, _, (), () = pe.trace_to_jaxpr_dynamic(known, avals_in)
+  jaxpr_known, _, (), () = pe.trace_to_jaxpr_dynamic(
+      lu.wrap_init(known, debug_info=jaxpr_known.debug_info), avals_in)
 
-  @lu.wrap_init
   def staged(*args):
     res_, ins = split_list(args, [len(which)])
     res = [_rem_singleton(x) if w else x for x, w in zip(res_, which_)]
     return core.eval_jaxpr(jaxpr_staged, (), *res, *ins)
   res_avals = [core.unmapped_aval(1, 0, v.aval) if w else v.aval
                for w, v in zip(which_, jaxpr_staged.invars[:len(which)])]
   avals_in = [*res_avals, *[v.aval for v in jaxpr_staged.invars[len(which):]]]
-  jaxpr_staged, _, (), () = pe.trace_to_jaxpr_dynamic(staged, avals_in)
+  jaxpr_staged, _, (), () = pe.trace_to_jaxpr_dynamic(
+      lu.wrap_init(staged, debug_info=jaxpr_staged.debug_info), avals_in)
 
   return jaxpr_known, jaxpr_staged
 
@@ -2070,7 +2072,8 @@ def _replication_rewrite_match(
     in_rep: Sequence[set[AxisName]],
     out_rep_dst: Sequence[set[AxisName]],
 ) -> core.ClosedJaxpr:
-  f = lu.wrap_init(partial(core.eval_jaxpr, jaxpr.jaxpr, jaxpr.consts))
+  f = lu.wrap_init(partial(core.eval_jaxpr, jaxpr.jaxpr, jaxpr.consts),
+                   debug_info=jaxpr.jaxpr.debug_info)
   f, out_rep = _efficient_transpose_rewrite_nomatch(f, mesh, in_rep)
   f = _match_rep(f, mesh, out_rep, out_rep_dst)
   jaxpr_, _, consts, () = pe.trace_to_jaxpr_dynamic(f, jaxpr.in_avals)

jax/experimental/sparse/bcoo.py

@@ -1042,8 +1042,14 @@ def _bcoo_dot_general_sampled_impl(A, B, indices, *, dimension_numbers):
 
 @bcoo_dot_general_sampled_p.def_abstract_eval
 def _bcoo_dot_general_sampled_abstract_eval(A, B, indices, *, dimension_numbers):
-  dense_result, = pe.abstract_eval_fun(lambda *args: [lax.dot_general(*args, dimension_numbers=dimension_numbers)], A, B)
-  sparse_result, = pe.abstract_eval_fun(lambda *args: [_bcoo_extract(*args)], indices, dense_result)
+  dbg = api_util.debug_info("bcoo_dot_general_sampled_abstract_eval",
+                            lax.dot_general, (A, B), dict(dimension_numbers=dimension_numbers))
+  dense_result, = pe.abstract_eval_fun(lambda *args: [lax.dot_general(*args, dimension_numbers=dimension_numbers)], A, B,
+                                       debug_info=dbg)
+  dbg = api_util.debug_info("bcoo_dot_general_sampled_abstract_eval",
+                            _bcoo_extract, (indices, dense_result), {})
+  sparse_result, = pe.abstract_eval_fun(lambda *args: [_bcoo_extract(*args)], indices, dense_result,
+                                        debug_info=dbg)
   return sparse_result
 
 def _bcoo_dot_general_sampled_transpose(ct, A, B, indices, *, dimension_numbers):

jax/experimental/sparse/transform.py

@@ -55,6 +55,7 @@
 
 import jax
 from jax import lax
+from jax._src import api_util
 from jax._src import config
 from jax._src import core
 from jax._src.custom_derivatives import lift_jvp
@@ -365,12 +366,15 @@ def sparsify_fun(wrapped_fun, args: list[ArrayOrSparse]):
   spenv = SparsifyEnv(out_bufs)
   return spvalues_to_arrays(spenv, out_spvalues())
 
-def _sparsify_with_tracer(fun):
+def _sparsify_with_tracer(fun: Callable):
   """Implementation of sparsify() using tracers."""
   @functools.wraps(fun)
   def _wrapped(*args):
     args_flat, in_tree = tree_flatten(args, is_leaf=_is_sparse_obj)
-    wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
+    wrapped_fun, out_tree = flatten_fun_nokwargs(
+        lu.wrap_init(fun,
+                     debug_info=api_util.debug_info("sparsify", fun, args, {})),
+        in_tree)
     out = sparsify_fun(wrapped_fun, args_flat)
     return tree_unflatten(out_tree(), out)
   return _wrapped
@@ -439,7 +443,12 @@ def wrapped(
   ) -> tuple[Sequence[SparsifyValue], pytree.PyTreeDef]:
     spvalues_flat, in_tree = tree_flatten(spvalues, is_leaf=_is_spvalue)
     in_avals_flat = spvalues_to_avals(spenv, spvalues_flat)
-    wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(f, params), in_tree)
+    wrapped_fun, out_tree = flatten_fun_nokwargs(
+        lu.wrap_init(
+            f, params,
+            debug_info=api_util.debug_info("sparsify", f,
+                                           spvalues_to_arrays(spenv, spvalues), {})),
+        in_tree)
     jaxpr, out_avals_flat, consts, () = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals_flat)
     result = eval_sparse(jaxpr, consts, spvalues_flat, spenv)
     if len(out_avals_flat) != len(result):
@@ -716,14 +725,14 @@ def _gather_sparse_rule(spenv, *args, dimension_numbers, slice_sizes, unique_ind
 
 sparse_rules_bcoo[lax.gather_p] = _gather_sparse_rule
 
-def _sparsify_jaxpr(spenv, jaxpr, *spvalues):
+def _sparsify_jaxpr(spenv: SparsifyEnv,
+                    jaxpr: core.ClosedJaxpr, *spvalues):
   # TODO(jakevdp): currently this approach discards all information about
   #   shared data & indices when generating the sparsified jaxpr. The
   #   current approach produces valid sparsified while loops, but they
   #   don't work in corner cases (see associated TODO in sparsify_test.py)
   out_tree: pytree.PyTreeDef | None = None
 
-  @lu.wrap_init
   def wrapped(*args_flat):
     # TODO(frostig,jakevdp): This closes over `spenv`, which can bring
     # in buffers from the "outer scope" as constants. Is this a
@@ -740,7 +749,8 @@ def wrapped(*args_flat):
   args = spvalues_to_arrays(spenv, spvalues)
   args_flat, in_tree = tree_flatten(args)
   avals_flat = [core.get_aval(arg) for arg in args_flat]
-  sp_jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(wrapped, avals_flat)
+  sp_jaxpr, _, consts, () = pe.trace_to_jaxpr_dynamic(
+      lu.wrap_init(wrapped, debug_info=jaxpr.jaxpr.debug_info), avals_flat)
   sp_jaxpr = pe.ClosedJaxpr(sp_jaxpr, consts)
   assert out_tree is not None
   return sp_jaxpr, out_tree
@@ -866,17 +876,18 @@ def _todense_sparse_rule(spenv, spvalue, *, tree):
 
 def _custom_jvp_sparse_rule(spenv, *spvalues, **params):
   call_jaxpr: core.ClosedJaxpr = params.pop('call_jaxpr')
-  jvp_jaxpr_thunk = params.pop('jvp_jaxpr_thunk')
-  num_consts = params.pop('num_consts')
+  jvp_jaxpr_fun: lu.WrappedFun = params.pop('jvp_jaxpr_fun')
+  num_consts: int = params.pop('num_consts')
   sp_call_jaxpr, out_tree = _sparsify_jaxpr(spenv, call_jaxpr, *spvalues)
-  @lu.wrap_init
   def fun(*arrs):
     sparrs = arrays_to_spvalues(spenv, arrs)
     out = eval_sparse(call_jaxpr.jaxpr, call_jaxpr.consts, sparrs, spenv)
     return spvalues_to_arrays(spenv, out)
-  jvp = lift_jvp(num_consts, jvp_jaxpr_thunk, call_jaxpr.jaxpr.debug_info)
+  jvp = lift_jvp(num_consts, jvp_jaxpr_fun)
   invals = spvalues_to_arrays(spenv, spvalues)
-  outvals = jax.custom_derivatives.custom_jvp_call_p.bind(fun, jvp, *invals, **params)
+  outvals = jax.custom_derivatives.custom_jvp_call_p.bind(
+      lu.wrap_init(fun, debug_info=call_jaxpr.jaxpr.debug_info),
+      jvp, *invals, **params)
   return arrays_to_spvalues(spenv, outvals)
 
 sparse_rules_bcoo[jax.custom_derivatives.custom_jvp_call_p] = _custom_jvp_sparse_rule

tests/core_test.py

@@ -26,9 +26,8 @@
 from jax import lax
 from jax import numpy as jnp
 from jax import jvp, linearize, vjp, jit, make_jaxpr
-from jax.api_util import flatten_fun_nokwargs
+from jax.api_util import flatten_fun_nokwargs, debug_info
 from jax._src import config
-
 from jax._src import core
 from jax._src import linear_util as lu
 from jax._src import util
@@ -48,14 +47,16 @@ def call(f, *args):
 @util.curry
 def core_call(f, *args):
   args, in_tree = jax.tree.flatten(args)
-  f, out_tree = flatten_fun_nokwargs(lu.wrap_init(f), in_tree)
+  dbg = debug_info("core_call_test", f, args, {})
+  f, out_tree = flatten_fun_nokwargs(lu.wrap_init(f, debug_info=dbg), in_tree)
   out = core.call_p.bind(f, *args)
   return jax.tree.unflatten(out_tree(), out)
 
 @util.curry
 def core_closed_call(f, *args):
   args, in_tree = jax.tree.flatten(args)
-  f, out_tree = flatten_fun_nokwargs(lu.wrap_init(f), in_tree)
+  dbg = debug_info("core_closed_call_test", f, args, {})
+  f, out_tree = flatten_fun_nokwargs(lu.wrap_init(f, debug_info=dbg), in_tree)
   out = core.closed_call_p.bind(f, *args)
   return jax.tree.unflatten(out_tree(), out)
 
@@ -362,7 +363,10 @@ def body(c, _):
 
     aval = core.ShapedArray((), jnp.dtype('int32'))
     pval = pe.PartialVal.unknown(aval)
-    jaxpr, _, _ = pe.trace_to_jaxpr_nounits(lu.wrap_init(f), [pval], False)
+    jaxpr, _, _ = pe.trace_to_jaxpr_nounits(
+        lu.wrap_init(f,
+                     debug_info=debug_info("test", f, (0,), {})),
+        [pval], False)
     dropvar, b = jaxpr.eqns[0].outvars
     self.assertEqual(dropvar.aval, aval)
 
@@ -548,12 +552,13 @@ def test_staging_basic(self):
     a = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
     b = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
 
-    @lu.wrap_init
     def f(x, y):
       return x, y
 
     jaxpr, _, _, () = pe.trace_to_jaxpr_dynamic(
-      f, [n, a, b], keep_inputs=[False, True, True])
+      lu.wrap_init(f,
+                   debug_info=debug_info("test", f, (1, 2), {})),
+      [n, a, b], keep_inputs=[False, True, True])
 
     self.assertLen(jaxpr.invars, 3)
     self.assertEqual((jaxpr.invars[0],), jaxpr.invars[1].aval.shape)
@@ -569,15 +574,16 @@ def test_staging_nested(self):
     a = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
     b = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
 
-    @lu.wrap_init
     def f(x, y):
       @jax.jit
       def g(x, y, z, w):
         return (x, w)
       return g(x, y, x, y)
 
     jaxpr, _, _, () = pe.trace_to_jaxpr_dynamic(
-      f, [n, a, b], keep_inputs=[False, True, True])
+      lu.wrap_init(f,
+                   debug_info=debug_info("test", f, (0, 1), {})),
+        [n, a, b], keep_inputs=[False, True, True])
 
     self.assertLen(jaxpr.invars, 1 + 2)  # one axis size var, two other inputs
     self.assertEqual((jaxpr.invars[0],), jaxpr.invars[1].aval.shape)
@@ -605,15 +611,16 @@ def test_staging_nested_including_shape_arg(self):
     a = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
     b = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
 
-    @lu.wrap_init
     def f(x, y):
       @jax.jit
       def g(_, x, y, z, w):
         return (x, w)
       return g(x.shape[0], x, y, x, y)
 
     jaxpr, _, _, () = pe.trace_to_jaxpr_dynamic(
-      f, [n, a, b], keep_inputs=[False, True, True])
+        lu.wrap_init(f,
+                     debug_info=debug_info("test", f, (1, 2), {})),
+        [n, a, b], keep_inputs=[False, True, True])
 
     # { lambda ; a:i32[] b:f32[a] c:f32[a]. let
     #     d:f32[a] e:f32[a] = xla_call[
@@ -641,15 +648,16 @@ def test_staging_primitive_applications(self):
     a = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
     b = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
 
-    @lu.wrap_init
     def f(x, y):
       z = lax.mul(x, y)
       w = lax.sin(z)
       u = lax_internal._reduce_sum(w, [0])
       return (u,)
 
     jaxpr, _, _, () = pe.trace_to_jaxpr_dynamic(
-      f, [n, a, b], keep_inputs=[False, True, True])
+        lu.wrap_init(f,
+                     debug_info=debug_info("test", f, (1, 2), {})),
+        [n, a, b], keep_inputs=[False, True, True])
 
     self.assertLen(jaxpr.invars, 1 + 2)  # one axis size var, two other inputs
     self.assertLen(jaxpr.eqns, 3)
@@ -667,14 +675,15 @@ def test_typecheck_staging_nested(self):
     a = core.DShapedArray((DBIdx(0),), jnp.dtype('float32'), weak_type=False)
     b = core.DShapedArray((DBIdx(1),), jnp.dtype('float32'), weak_type=False)
 
-    @lu.wrap_init
     def f(a, b):
       @jax.jit
       def g(x): return x
       return g(a),
 
     jaxpr, _, _, () = pe.trace_to_jaxpr_dynamic(
-      f, [n, m, a, b], keep_inputs=[False, False, True, True])
+        lu.wrap_init(f,
+                     debug_info=debug_info("test", f, (1, 2), {})),
+        [n, m, a, b], keep_inputs=[False, False, True, True])
     # { lambda ; a:i32[] b:i32[] c:f32[a] d:f32[b]. let
     #     e:f32[a] = xla_call[
     #       call_jaxpr={ lambda ; f:i32[] g:f32[f]. let  in (g,) }

tests/debug_info_test.py

@@ -881,9 +881,9 @@ 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][<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="),
+            "traced_for=jit, fun=apply_fn, arg_names=inp, result_paths=[0],[1][0][0][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="),
         ],
         check_tracer_arg_name=True,
         expected_tracer_debug_infos=[
@@ -1071,8 +1071,7 @@ def fn_tp(r, t):
         expected_jaxpr_debug_infos=[
             "traced_for=cond, fun=my_f, arg_names=x['c'], result_paths=",
             "traced_for=cond, fun=<lambda>, arg_names=x['c'], result_paths=",
-            # TODO(necula): flat_index?
-            "traced_for=jit, fun=<lambda>, arg_names=x, result_paths=[<flat index 0>][0][0],[<flat index 0>][0][1]",
+            "traced_for=jit, fun=<lambda>, arg_names=x, result_paths=[0][0][0],[0][0][1]",
         ],
         check_tracer_arg_name=True,
         expected_tracer_debug_infos=[

tests/shard_map_test.py

@@ -29,6 +29,7 @@
 
 import jax
 import jax.ad_checkpoint
+from jax import api_util
 from jax import lax
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
@@ -1329,7 +1330,11 @@ def foo(x):
 
   def test_rewrite_process_call(self):
     def f(x):
-      return core.call_p.bind(lu.wrap_init(lambda x: [2. * x]), x)[0] * x
+      return core.call_p.bind(
+          lu.wrap_init(lambda x: [2. * x],
+                       debug_info=api_util.debug_info("test", lambda x: [2. * x],
+                                                      (x,), {})),
+          x)[0] * x
 
     mesh = jtu.create_mesh((4,), ('x',))
     g = shard_map(f, mesh, in_specs=(P('x'),), out_specs=P('x'))
@@ -1345,7 +1350,10 @@ def test_rewrite_post_process_call(self):
     @jax.jit
     @partial(shard_map, mesh=mesh, in_specs=(P('x'),), out_specs=P('x'))
     def f(x):
-      return core.call_p.bind(lu.wrap_init(lambda: [2. * x]))[0] * x
+      return core.call_p.bind(
+          lu.wrap_init(lambda: [2. * x],
+                       debug_info=api_util.debug_info("test", lambda: [2. * x],
+                                                      (), {})))[0] * x
 
     x = jnp.arange(4.)
     y = f(x)