Merge code between make_jaxpr and jit(f).trace.

The semantics of `make_jaxpr` are preserved here i.e. `make_jaxpr` still closes over tracers but `jit(f).trace` doesn't.

Since we can keep the existing behavior and still merge the implementation is a good cleanup!

Fixes #21116

PiperOrigin-RevId: 641347140

Author: yashk2810

jax/_src/api.py

@@ -71,7 +71,8 @@
 from jax._src.layout import Layout, AutoLayout
 from jax._src.traceback_util import api_boundary
 from jax._src import tree_util
-from jax._src.util import unzip2, safe_map, safe_zip, wrap_name, wraps
+from jax._src.util import (unzip2, safe_map, safe_zip, wrap_name, wraps,
+                           split_list)
 from jax._src import util
 
 from jax._src.interpreters import ad
@@ -2358,43 +2359,34 @@ def make_jaxpr(fun: Callable,
       g:f32[] = mul f c
     in (g,) }
   """
-  check_callable(fun)
-  static_argnums = _ensure_index_tuple(static_argnums)
-
-  def abstractify(args, kwargs):
-    flat_args, in_tree = tree_flatten((args, kwargs))
-    if abstracted_axes is None:
-      return map(shaped_abstractify, flat_args), in_tree, [True] * len(flat_args)
-    else:
-      axes_specs = _flat_axes_specs(abstracted_axes, *args, **kwargs)
-      in_type = pe.infer_lambda_input_type(axes_specs, flat_args)
-      in_avals, keep_inputs = unzip2(in_type)
-      return in_avals, in_tree, keep_inputs
+  try:
+    hash(fun)
+    weakref.ref(fun)
+  except TypeError:
+    fun = partial(fun)
 
   @wraps(fun)
   @api_boundary
   def make_jaxpr_f(*args, **kwargs):
-    f = lu.wrap_init(fun)
-    if static_argnums:
-      dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
-      f, args = argnums_partial(f, dyn_argnums, args)
-    in_avals, in_tree, keep_inputs = abstractify(args, kwargs)
-    in_type = tuple(zip(in_avals, keep_inputs))
-    f, out_tree = flatten_fun(f, in_tree)
-    f = lu.annotate(f, in_type)
-    debug_info = pe.debug_info(fun, in_tree, out_tree, True, 'make_jaxpr')
     with ExitStack() as stack:
       for axis_name, size in axis_env or []:
         stack.enter_context(core.extend_axis_env(axis_name, size, None))
-      jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(
-          f, debug_info=debug_info)
-    closed_jaxpr = core.ClosedJaxpr(jaxpr, consts)
+      traced = jit(fun, static_argnums=static_argnums,
+                   abstracted_axes=abstracted_axes).trace(*args, **kwargs)
+    # `jit` converts tracers in consts to args but that breaks the semantics of
+    # `make_jaxpr`. Hence convert the tracers in args back to consts in jaxpr.
+    if traced._num_consts:
+      consts, _ = split_list(traced._args_flat, [traced._num_consts])
+      jaxpr_ = pe.convert_invars_to_constvars(traced.jaxpr.jaxpr,
+                                              traced._num_consts)
+      jaxpr = core.ClosedJaxpr(jaxpr_, consts)
+    else:
+      jaxpr = traced.jaxpr
     if return_shape:
-      out_avals, _ = unzip2(out_type)
-      out_shapes_flat = [
-          ShapeDtypeStruct(a.shape, a.dtype, a.named_shape) for a in out_avals]
-      return closed_jaxpr, tree_unflatten(out_tree(), out_shapes_flat)
-    return closed_jaxpr
+      out = [ShapeDtypeStruct(o.shape, o.dtype, getattr(o, 'named_shape', None))
+             for o in jaxpr.out_avals]
+      return jaxpr, tree_unflatten(tree_structure(traced.out_info), out)
+    return jaxpr
 
   make_jaxpr_f.__module__ = "jax"
   if hasattr(fun, "__qualname__"):

jax/_src/pjit.py

@@ -171,7 +171,7 @@ class PjitInfo(NamedTuple):
 
 
 def _python_pjit_helper(jit_info, *args, **kwargs):
-  (args_flat, params, _, out_tree, _, arg_names,
+  (args_flat, params, in_avals, _, out_tree, _, arg_names, _,
    attrs_tracked) = _infer_params(jit_info, args, kwargs)
 
   for arg in args_flat:
@@ -197,7 +197,7 @@ def _python_pjit_helper(jit_info, *args, **kwargs):
     if params['jaxpr'].consts:
       raise TypeError(e.args[0]) from e
     else:
-      for arg, name, aval in zip(args_flat, arg_names, params['jaxpr'].in_avals):
+      for arg, name, aval in zip(args_flat, arg_names, in_avals):
         try:
           xla.canonicalize_dtype(arg)
         except xla.InvalidInputException as _:
@@ -491,7 +491,7 @@ def lower(*args, **kwargs):
 
   @api_boundary
   def eval_shape(*args, **kwargs):
-    _, params, _, out_tree, _, _, _ = _infer_params(jit_info, args, kwargs)
+    _, params, _, _, out_tree, _, _, _, _ = _infer_params(jit_info, args, kwargs)
     out_s = [None if is_unspecified(s) else s for s in params['out_shardings']]
     # TODO(yashkatariya): Add `Layout` to SDS.
     out = [api.ShapeDtypeStruct(x.shape, x.dtype, x.named_shape, sharding=s)
@@ -503,16 +503,15 @@ def trace(*args, **kwargs) -> stages.Traced:
     lowering_parameters = kwargs.pop(
         '_experimental_lowering_parameters', mlir.LoweringParameters())
 
-    (args_flat, params, in_tree, out_tree, donated_invars,
-     arg_names, _) = _infer_params(jit_info, args, kwargs)
+    (args_flat, params, in_avals, in_tree, out_tree, donated_invars,
+     arg_names, num_consts, _) = _infer_params(jit_info, args, kwargs)
 
     donate_argnums = tuple(i for i, d in enumerate(donated_invars) if d)
-    jaxpr = params['jaxpr']
-    args_info = stages.make_args_info(in_tree, jaxpr.in_avals, donate_argnums)
+    args_info = stages.make_args_info(in_tree, in_avals, donate_argnums)
     lower_callable = partial(_resolve_and_lower, args_flat, **params,
                              lowering_parameters=lowering_parameters)
-    return stages.Traced(jaxpr, args_info, params["name"], out_tree,
-                         lower_callable, args_flat, arg_names)
+    return stages.Traced(params['jaxpr'], args_info, params["name"], out_tree,
+                         lower_callable, args_flat, arg_names, num_consts)
 
   wrapped = _cpp_pjit(jit_info)
   wrapped.lower = lower
@@ -662,8 +661,9 @@ def _infer_params(jit_info, args, kwargs):
       keep_unused=keep_unused,
       inline=inline,
   )
-  return (consts + args_flat, params, in_tree, out_tree(),
-          donated_invars, dbg.arg_names if dbg else None, attrs_tracked)
+  return (consts + args_flat, params, in_avals, in_tree, out_tree(),
+          donated_invars, dbg.arg_names if dbg else None, len(consts),
+          attrs_tracked)
 
 def _extract_implicit_args(
   in_type: Sequence[tuple[core.AbstractValue, bool]],

jax/_src/stages.py

@@ -427,17 +427,19 @@ class CompiledCallParams(NamedTuple):
 
 class Traced(Stage):
   __slots__ = ["jaxpr", "args_info", "fun_name", "_out_tree", "_lower_callable",
-               "_args_flat", "_arg_names"]
+               "_args_flat", "_arg_names", "_num_consts"]
 
   def __init__(self, jaxpr: core.ClosedJaxpr, args_info, fun_name, out_tree,
-               lower_callable, args_flat=None, arg_names=None):
+               lower_callable, args_flat=None, arg_names=None,
+               num_consts: int = 0):
     self.jaxpr = jaxpr
     self.args_info = args_info
     self.fun_name = fun_name
     self._out_tree = out_tree
     self._lower_callable = lower_callable
     self._args_flat = args_flat
     self._arg_names = arg_names
+    self._num_consts = num_consts
 
   @property
   def out_info(self):

tests/api_test.py

@@ -4058,6 +4058,16 @@ def f():
       return jnp.exp(dtype(0))
     f()  # doesn't error
 
+  def test_vmap_make_jaxpr_close_over_tracer(self):
+    def run(inp):
+      def f(x, y):
+        return x + y
+      g = lambda x: f(x, inp)
+      jaxpr = jax.make_jaxpr(g)(1)
+      return jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 1)
+
+    jax.vmap(run)(jnp.arange(2))  # doesn't crash
+
   def test_large_python_ints(self):
     with self.assertRaises(OverflowError):
       jnp.multiply(2 ** 100, 3.)

tests/lax_control_flow_test.py

@@ -2518,7 +2518,7 @@ def f(c, a):
     scan_fun = lambda c, xs: lax.scan(f, c, xs)
 
     def new_jaxpr():
-      jaxpr = jax.make_jaxpr(scan_fun)(c, xs).jaxpr
+      jaxpr = jax.make_jaxpr(partial(scan_fun))(c, xs).jaxpr
       scan = next(eqn for eqn in jaxpr.eqns if eqn.primitive.name == 'scan')
       return jaxpr, scan