[debug] Extract reproducers from JAX errors.

This is just a draft of an experiment.
The purpose is to help debug JAX failures in large users examples,
by trying to extract a smaller program that contains only JAX
API calls. E.g., in a large program using Flax, this would get
Flax out of the way.

Author: gnecula

jax/_src/BUILD

@@ -1358,6 +1358,17 @@ pytype_strict_library(
     ],
 )
 
+pytype_strict_library(
+    name = "repro",
+    srcs = ["repro.py", "repro_runtime.py"],
+    visibility = ["//jax:internal"] + jax_visibility("repro"),
+    deps = [
+        ":source_info_util",
+        ":traceback_util",
+        "//jax/_src/lib",
+    ],
+)
+
 pytype_strict_library(
     name = "state_types",
     srcs = [

jax/_src/ad_checkpoint.py

@@ -200,6 +200,7 @@ def policy(prim, *args, **params):
 ### Main API
 
 @api_boundary
+@partial(core.jax_boundary, api_name="jax.checkpoint")
 def checkpoint(fun: Callable, *, prevent_cse: bool = True,
                policy: Callable[..., bool] | None = None,
                static_argnums: int | tuple[int, ...] = (),

jax/_src/api.py

@@ -147,6 +147,7 @@ def _update_debug_special_thread_local(_):
 float0 = dtypes.float0
 
 
+@partial(core.jax_boundary, api_name="jax.jit")
 def jit(
   fun: Callable, /, *,
   in_shardings: Any = sharding_impls.UNSPECIFIED,
@@ -347,6 +348,7 @@ def disable_jit(disable: bool = True):
     yield
 
 
+@partial(core.jax_boundary, api_name="jax.grad")
 def grad(fun: Callable, argnums: int | Sequence[int] = 0,
          has_aux: bool = False, holomorphic: bool = False,
          allow_int: bool = False,
@@ -413,6 +415,7 @@ def grad_f_aux(*args, **kwargs):
 
   return grad_f_aux if has_aux else grad_f
 
+@partial(core.jax_boundary, api_name="jax.value_and_grad")
 def value_and_grad(fun: Callable, argnums: int | Sequence[int] = 0,
                    has_aux: bool = False, holomorphic: bool = False,
                    allow_int: bool = False, reduce_axes: Sequence[AxisName] = ()
@@ -545,6 +548,7 @@ def _check_output_dtype_revderiv(name, holomorphic, x):
                     "jax.vjp directly.")
 _check_output_dtype_grad = partial(_check_output_dtype_revderiv, "grad")
 
+@partial(core.jax_boundary, api_name="jax.fwd_and_bwd")
 def fwd_and_bwd(
     fun: Callable, argnums: int | Sequence[int], has_aux: bool = False,
     jitted: bool = True,
@@ -619,6 +623,7 @@ def bwd(f_vjp, outgrad):
   return fwd, bwd
 
 
+@partial(core.jax_boundary, api_name="jaxfwd")
 def jacfwd(fun: Callable, argnums: int | Sequence[int] = 0,
            has_aux: bool = False, holomorphic: bool = False) -> Callable:
   """Jacobian of ``fun`` evaluated column-by-column using forward-mode AD.
@@ -709,6 +714,7 @@ def _check_output_dtype_jacfwd(holomorphic, x):
       raise TypeError("jacfwd with holomorphic=True requires outputs with complex dtype, "
                       f"but got {aval.dtype.name}.")
 
+@partial(core.jax_boundary, api_name="jax.jacrev")
 def jacrev(fun: Callable, argnums: int | Sequence[int] = 0,
            has_aux: bool = False, holomorphic: bool = False,
            allow_int: bool = False) -> Callable:
@@ -789,7 +795,7 @@ def jacobian(fun: Callable, argnums: int | Sequence[int] = 0,
 _check_input_dtype_jacrev = partial(_check_input_dtype_revderiv, "jacrev")
 _check_output_dtype_jacrev = partial(_check_output_dtype_revderiv, "jacrev")
 
-
+@partial(core.jax_boundary, api_name="jax.hessian")
 def hessian(fun: Callable, argnums: int | Sequence[int] = 0,
             has_aux: bool = False, holomorphic: bool = False) -> Callable:
   """Hessian of ``fun`` as a dense array.
@@ -915,6 +921,7 @@ def _split(x, indices, axis):
     return x._split(indices, axis)
 
 
+@partial(core.jax_boundary, api_name="jax.vmap")
 def vmap(fun: F,
          in_axes: int | None | Sequence[Any] = 0,
          out_axes: Any = 0,
@@ -1234,7 +1241,7 @@ def _all_sizes_index(sz):
       msg.append(f"  * some axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\n")
   raise ValueError(''.join(msg)[:-2])  # remove last semicolon and newline
 
-
+@partial(core.jax_boundary, api_name="jax.pmap")
 def pmap(
     fun: Callable,
     axis_name: AxisName | None = None,
@@ -1789,6 +1796,7 @@ def _cpp_mapped_lower(pmap_f, *args, **kwargs):
 
 
 @api_boundary
+@partial(core.jax_boundary, api_name="jax.jvp")
 def jvp(
     fun: Callable, primals, tangents, has_aux: bool = False
   ) -> tuple[Any, ...]:
@@ -1883,6 +1891,7 @@ def linearize(fun: Callable, *primals, has_aux: Literal[True]
               ) -> tuple[Any, Callable, Any]:
   ...
 
+@partial(core.jax_boundary, api_name="jax.linearize")
 def linearize(fun: Callable, *primals, has_aux: bool = False
               ) -> tuple[Any, Callable] | tuple[Any, Callable, Any]:
   """Produces a linear approximation to ``fun`` using :py:func:`jvp` and partial eval.
@@ -2073,6 +2082,7 @@ def vjp(fun: Callable[..., tuple[T, U]], *primals: Any,
   ...
 
 @api_boundary
+@partial(core.jax_boundary, api_name="jax.vjp")
 def vjp(
     fun: Callable, *primals, has_aux: bool = False, reduce_axes=()
   ) -> tuple[Any, Callable] | tuple[Any, Callable, Any]:
@@ -2146,7 +2156,7 @@ def _vjp(fun: lu.WrappedFun, *primals, has_aux=False):
   else:
     return out_primal_py, vjp_py, tree_unflatten(aux_tree, aux)
 
-
+@partial(core.jax_boundary, api_name="jax.saved_input_vjp")
 def saved_input_vjp(f: Callable, which: Sequence[bool], *primals,
                     allow_unused: bool = True, allow_opaque: bool = True):
   if len(which) != len(primals):
@@ -2490,6 +2500,7 @@ def make_jaxpr(
 ) -> Callable[..., tuple[core.ClosedJaxpr, Any]]:
   ...
 
+@partial(core.jax_boundary, api_name="jax.make_japr")
 def make_jaxpr(
     fun: Callable,
     static_argnums: int | Iterable[int] = (),
@@ -3001,6 +3012,7 @@ def eval_shape(fun, *args, **kwargs):
   return jit(fun).trace(*args, **kwargs).out_info
 
 
+@partial(core.jax_boundary, api_name="jax.named_call")
 def named_call(
     fun: F,
     *,

jax/_src/config.py

@@ -2030,3 +2030,19 @@ def _default_pmap_no_rank_reduction(new_val):
         ' ragged_dot_general_p.'
     ),
 )
+
+repro_dir = string_flag(
+    name='jax_output_dir',
+    default=os.getenv("JAX_REPRO_FLAGS", ""),
+    help=(
+        'Turn on saving of repros. EXPERIMENTAL.'
+    ),
+)
+
+repro_flags = string_flag(
+    name='jax_repro_flags',
+    default=os.getenv("JAX_REPRO_FLAGS", ""),
+    help=(
+        'Comma-separated flags for repros. EXPERIMENTAL.'
+    ),
+)

jax/_src/core.py

@@ -47,6 +47,7 @@
     TracerIntegerConversionError, UnexpectedTracerError)
 from jax._src import linear_util as lu
 from jax._src.tree_util import tree_flatten, tree_unflatten
+from jax._src import repro
 from jax._src import source_info_util
 from jax._src.util import (safe_zip, safe_map, curry, tuple_insert,
                            tuple_delete, cache,
@@ -83,6 +84,8 @@
 
 def identity(x): return x
 
+jax_boundary = repro.jax_boundary
+
 # -------------------- jaxprs --------------------
 
 Effect = effects.Effect
@@ -631,7 +634,10 @@ def __repr__(self):
 
   def bind(self, *args, **params):
     args = args if self.skip_canonicalization else map(canonicalize_value, args)
-    return self._true_bind(*args, **params)
+    if config.repro_dir.value:
+      return repro._true_bind_primitive(self, args, params)
+    else:
+      return self._true_bind(*args, **params)
 
   def _true_bind(self, *args, **params):
     for arg in args:

jax/_src/custom_derivatives.py

@@ -138,6 +138,7 @@ def f_jvp(primals, tangents):
   jvp: Callable[..., tuple[ReturnValue, ReturnValue]] | None = None
   symbolic_zeros: bool = False
 
+  @partial(core.jax_boundary, api_name="jax.custom_jvp", api_constructor=True)
   def __init__(self,
                fun: Callable[..., ReturnValue],
                nondiff_argnums: Sequence[int] = (),
@@ -162,6 +163,7 @@ def __init__(self,
 
   __getattr__ = custom_api_util.forward_attr
 
+  @partial(core.jax_boundary, api_name="jax_defjvp")
   def defjvp(self,
              jvp: Callable[..., tuple[ReturnValue, ReturnValue]],
              symbolic_zeros: bool = False,
@@ -214,6 +216,7 @@ def defjvp(self,
     self.symbolic_zeros = symbolic_zeros
     return jvp
 
+  @partial(core.jax_boundary, api_name="jax_defjvps")
   def defjvps(self, *jvps: Callable[..., ReturnValue] | None) -> None:
     """Convenience wrapper for defining JVPs for each argument separately.
 
@@ -257,6 +260,7 @@ def jvp(primals, tangents):
     self.defjvp(jvp)
 
   @traceback_util.api_boundary
+  @partial(core.jax_boundary, api_name="jax.custom_jvp_call")
   def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:  # pytype: disable=invalid-annotation
     debug = debug_info("custom_jvp fun", self.fun, args, kwargs,
                        static_argnums=self.nondiff_argnums)
@@ -559,7 +563,7 @@ def f_bwd(res, g):
 
   .. _tutorial: https://docs.jax.dev/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html
   """
-
+  @partial(core.jax_boundary, api_name="jax.custom_vjp", api_constructor=True)
   def __init__(self,
                fun: Callable[..., ReturnValue],
                nondiff_argnums: Sequence[int] = (),
@@ -587,6 +591,7 @@ def __init__(self,
 
   __getattr__ = custom_api_util.forward_attr
 
+  @partial(core.jax_boundary, api_name="jax_defvjp")
   def defvjp(self,
              fwd: Callable[..., tuple[ReturnValue, Any]],
              bwd: Callable[..., tuple[Any, ...]],
@@ -686,6 +691,7 @@ def defvjp(self,
           "remat optimization for custom_vjp does not support symbolic zeros")
 
   @traceback_util.api_boundary
+  @partial(core.jax_boundary, api_name="jax.custom_vjp_call")
   def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:  # pytype: disable=invalid-annotation
     debug_fun = debug_info("custom_vjp fun", self.fun, args, kwargs,
                            static_argnums=self.nondiff_argnums)

jax/_src/interpreters/partial_eval.py

@@ -2382,7 +2382,6 @@ def trace_to_jaxpr_dynamic(
     in_tracers = _input_type_to_tracers(
         partial(trace.new_arg, source_info=source_info), in_avals)
     in_tracers = [t for t, keep in zip(in_tracers, keep_inputs) if keep]
-
     with core.set_current_trace(trace):
       ans = fun.call_wrapped(*in_tracers)
     _check_returned_jaxtypes(fun.debug_info, ans)

jax/_src/lax/control_flow/conditionals.py

@@ -400,6 +400,7 @@ def _capitalize(s):
   return s[0].capitalize() + s[1:]
 
 @api_boundary
+@partial(core.jax_boundary, api_name="jax.lax.cond")
 @functools.wraps(_cond)
 def cond(*args, **kwargs):
   # detect an attempt to call the former, deprecated cond

jax/_src/lax/control_flow/loops.py

@@ -110,6 +110,7 @@ def _promote_weak_typed_inputs(in_vals, in_avals, out_avals):
 Y = TypeVar('Y')
 
 @api_boundary
+@partial(core.jax_boundary, api_name="jax.lax.scan")
 def scan(f: Callable[[Carry, X], tuple[Carry, Y]],
          init: Carry,
          xs: X | None = None,
@@ -1603,6 +1604,7 @@ def _move_right(lst, to_move):
 ### while_loop
 
 @api_boundary
+@partial(core.jax_boundary, api_name="jax.lax.while_loop")
 def while_loop(cond_fun: Callable[[T], BooleanNumeric],
                body_fun: Callable[[T], T],
                init_val: T) -> T: