Repository: jax-ml/jax
Pull request #27873 (2025-04-09 13:03:24)

Make `trace` and `lower` class attributes for `jax.jit`.
Previously, jax.jit returned a function with extra attributes, e.g., `trace`, and `lower`, such that we can use:

```
jax.jit(f).trace(...)
```

The new attributes create problems when `jax.jit` is used along `functools.wraps`. Essentially, `functools.wraps(jax.jit(f))(wrapper)` is supposed to result in a function that when invoked will invoke `wrapper` and then presumably `jax.jit(f)`. This works as expected if you just call the result, but if you try to use it with `lower` and `trace`, the `wrapper` is bypassed. This is because `wraps` copies the attributes `trace` and `lower` from `jax.jit(f)` onto the resulting function, so when `trace` is invoked the `wrapper` is bypassed entirely.

See #27829 and #27825.

The solution proposed here is to make the `trace` and `lower` be class attributes, so that they are not copied by `functools.wraps` (it only copies the object attributes). Since they are class attributes calling `jax.jit(f).lower()` will
still work because it will look up the `lower` attribute in the class if the object does not have it.
Bug, if you try to use `lower` or `trace` on the result of `functools.wraps(jax.jit(f))()` you will get an error that
`lower` or `trace` does not exist.
That is better than silently ignoring the wrapper. The workaround is to apply `jax.jit` last among your wrappers.

Fixes: #27829
jax-fixit

gnecula: > This looks great! Can you run the benchmarks in api_benchmark.py once please on CPU?
> 
> Instructions in go/jax-benchmarks

I ran the benchmarks and there is no regression. There is no reason to expect a regression, because nothing changes except the path for invoking `trace` and `lower`, which now involve another attribute lookup on the class after the lookup on the object fails.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 00c96b228..d118ef10f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -24,6 +24,12 @@ When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.
     which was added temporarily in v0.4.36 to allow users to opt out of the
     new "stackless" tracing machinery.
   * Removed the `config.jax_eager_pmap` config option.
+  * Disallow the calling of `lower` and `trace` AOT APIs on the result
+    of `jax.jit` if there have been subsequent wrappers applied.
+    Previously this worked, but silently ignored the wrappers.
+    The workaround is to apply `jax.jit` last among the wrappers,
+    and similarly for `jax.pmap`.
+    See {jax-issue}`#27873`.
 
 * Changes
   * The minimum CuDNN version is v9.8.
diff --git a/jax/_src/api.py b/jax/_src/api.py
index 2149510a2..d170d2632 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1569,11 +1569,13 @@ def _cpp_pmap(
       out_axes)
   del static_broadcasted_argnums, donate_argnums
 
+  prepare_pmap_fn = partial(_prepare_pmap,
+        fun, in_axes, out_axes, static_broadcasted_tuple, donate_tuple,
+        devices, backend, axis_size)
+
   @api_boundary
   def cache_miss(*args, **kwargs):
-    p = _prepare_pmap(fun, in_axes, out_axes, static_broadcasted_tuple,
-                      donate_tuple, devices, backend,
-                      axis_size, args, kwargs)
+    p = prepare_pmap_fn(args, kwargs)
     for arg in p.flat_args:
       dispatch.check_arg(arg)
 
@@ -1650,48 +1652,56 @@ def _cpp_pmap(
   _pmap_cache_clears.add(cpp_mapped_f)
 
   pmap_f = wraps(fun)(cpp_mapped_f)
+  # Store some data for the `lower` and `trace` methods
   pmap_f._fun = fun
+  pmap_f._prepare_pmap = prepare_pmap_fn
+  pmap_f._backend = backend
+  pmap_f._axis_name = axis_name
+  pmap_f._donate_tuple = donate_tuple
+
+  # TODO(necula): move these to top-level; we don't need to do this for
+  # every pmap
+  cpp_mapped_f_class = type(pmap_f)
+  cpp_mapped_f_class.lower = _cpp_mapped_lower
+  cpp_mapped_f_class.trace = _cpp_mapped_trace
+  # We return directly the function produced by pmap_lib.pmap, because we do not
+  # want to have Python in the dispatch path.
+  return pmap_f
 
-  @api_boundary
-  def lower(*args, **kwargs):
-    return trace(*args, **kwargs).lower()
+@api_boundary
+def _cpp_mapped_trace(pmap_f, *args, **kwargs):
+  p = pmap_f._prepare_pmap(args, kwargs)
+  abstract_args = list(map(shaped_abstractify, p.flat_args))
+  closed_jaxpr, xc_backend, replicas, shards, pci = pxla.get_pmap_jaxpr(
+      p.flat_fun, pmap_f._backend, pmap_f._axis_name,
+      axis_size=p.local_axis_size, global_axis_size=p.global_axis_size,
+      devices=p.devices,
+      name=p.flat_fun.__name__,
+      in_axes=p.in_axes_flat,
+      out_axes_thunk=p.out_axes_thunk,
+      avals=abstract_args)
+  lower_callable = partial(
+      pxla.lower_parallel_callable, p.flat_fun, pmap_f._axis_name,
+      axis_size=p.local_axis_size, global_axis_size=p.global_axis_size,
+      devices=p.devices,
+      name=p.flat_fun.__name__,
+      in_axes=p.in_axes_flat,
+      donated_invars=p.donated_invars,
+      is_explicit_global_axis_size=p.is_explicit_global_axis_size,
+      avals=abstract_args,
+      closed_jaxpr=closed_jaxpr,
+      backend=xc_backend,
+      replicas=replicas,
+      shards=shards,
+      pci=pci)
+  args_info = stages.make_args_info(p.in_tree, abstract_args, pmap_f._donate_tuple)
+  return stages.Traced(closed_jaxpr, args_info, p.flat_fun.__name__,
+                       p.out_tree(), lower_callable)
 
-  @api_boundary
-  def trace(*args, **kwargs):
-    p = _prepare_pmap(
-        fun, in_axes, out_axes, static_broadcasted_tuple, donate_tuple,
-        devices, backend, axis_size, args, kwargs)
-    abstract_args = list(map(shaped_abstractify, p.flat_args))
-    closed_jaxpr, xc_backend, replicas, shards, pci = pxla.get_pmap_jaxpr(
-        p.flat_fun, backend, axis_name,
-        axis_size=p.local_axis_size, global_axis_size=p.global_axis_size,
-        devices=p.devices,
-        name=p.flat_fun.__name__,
-        in_axes=p.in_axes_flat,
-        out_axes_thunk=p.out_axes_thunk,
-        avals=abstract_args)
-    lower_callable = partial(
-        pxla.lower_parallel_callable, p.flat_fun, axis_name,
-        axis_size=p.local_axis_size, global_axis_size=p.global_axis_size,
-        devices=p.devices,
-        name=p.flat_fun.__name__,
-        in_axes=p.in_axes_flat,
-        donated_invars=p.donated_invars,
-        is_explicit_global_axis_size=p.is_explicit_global_axis_size,
-        avals=abstract_args,
-        closed_jaxpr=closed_jaxpr,
-        backend=xc_backend,
-        replicas=replicas,
-        shards=shards,
-        pci=pci)
-    args_info = stages.make_args_info(p.in_tree, abstract_args, donate_tuple)
-    return stages.Traced(closed_jaxpr, args_info, p.flat_fun.__name__,
-                         p.out_tree(), lower_callable)
-
-  pmap_f.lower = lower
-  pmap_f.trace = trace
+@api_boundary
+def _cpp_mapped_lower(pmap_f, *args, **kwargs):
+  return _cpp_mapped_trace(pmap_f, *args, **kwargs).lower()
 
-  return pmap_f
 
 _pmap_cache_clears = weakref.WeakSet()  # type: ignore
 
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index 83ddb8709..871fb0a6c 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -308,12 +308,6 @@ def _get_fastpath_data(
   return fastpath_data
 
 
-def _cpp_pjit_evict_fn(self):
-  self._clear_cache()
-  _create_pjit_jaxpr.evict_function(self._fun)  # pytype: disable=attribute-error
-  _infer_params_cached.cache_clear()
-
-
 # The entries are doubled here from the default 4096 because _pjit_call_impl
 # also has a cpp dispatch path and that would double the number of entries in
 # the global shared cache.
@@ -374,9 +368,50 @@ def _cpp_pjit(fun: Callable, jit_info: PjitInfo):
 
   cpp_pjitted_f = wraps(fun)(cpp_pjit_f)
   cpp_pjitted_f._fun = fun
-  type(cpp_pjitted_f).clear_cache = _cpp_pjit_evict_fn
+  cpp_pjitted_f._jit_info = jit_info
+  # TODO(necula): move these to top-level; we don't need to do this for
+  # every jit
+  cpp_jitted_f_class = type(cpp_pjitted_f)
+  # TODO(necula): make clear_cache private, no need to have it part of the API
+  cpp_jitted_f_class.clear_cache = jit_evict_fn
+  cpp_jitted_f_class.lower = jit_lower
+  cpp_jitted_f_class.trace = jit_trace
+  cpp_jitted_f_class.eval_shape = jit_eval_shape
+  # We return directly the function produced by _xla.pjit, because we do not
+  # want to have Python in the dispatch path.
   return cpp_pjitted_f
 
+@api_boundary
+def jit_trace(jit_func, *args, **kwargs) -> stages.Traced:
+  p, args_flat = _infer_params(jit_func._fun, jit_func._jit_info, args, kwargs)
+  donate_argnums = tuple(i for i, d in enumerate(p.donated_invars) if d)
+  args_info = stages.make_args_info(p.in_tree, p.in_avals, donate_argnums)
+  lower_callable = partial(_resolve_and_lower, args_flat, **p.params,
+                           pgle_profiler=None)
+  return stages.Traced(
+      p.params['jaxpr'], args_info, p.params["name"], p.out_tree,
+      lower_callable, args_flat, p.arg_names, p.num_consts)
+
+
+@api_boundary
+def jit_lower(jit_func, *args, **kwargs):
+  return jit_trace(jit_func, *args, **kwargs).lower()
+
+@api_boundary
+def jit_eval_shape(jit_func, *args, **kwargs):
+  p, _ = _infer_params(jit_func._fun, jit_func._jit_info, args, kwargs)
+  out_s = [None if isinstance(s, UnspecifiedValue) else s for s in p.params['out_shardings']]
+  # TODO(yashkatariya): Add `Layout` to SDS.
+  out = [api.ShapeDtypeStruct(x.shape, x.dtype, sharding=s,
+                              weak_type=x.weak_type)
+         for x, s in zip(p.params['jaxpr'].out_avals, out_s)]
+  return tree_unflatten(p.out_tree, out)
+
+def jit_evict_fn(self):
+  self._clear_cache()
+  _create_pjit_jaxpr.evict_function(self._fun)  # pytype: disable=attribute-error
+  _infer_params_cached.cache_clear()
+
 
 def _split_layout_and_sharding(entries):
   entries_flat, treedef = tree_flatten(entries, is_leaf=lambda x: x is None)
@@ -484,41 +519,6 @@ def _parse_jit_arguments(fun: Callable, in_shardings: Any, out_shardings: Any,
         use_resource_env=use_resource_env,
         compiler_options_kvs=compiler_options_kvs)
 
-
-def _make_jit_wrapper(fun: Callable, jit_info: PjitInfo):
-
-  @api_boundary
-  def lower(*args, **kwargs):
-    return trace(*args, **kwargs).lower()
-
-  @api_boundary
-  def eval_shape(*args, **kwargs):
-    p, _ = _infer_params(fun, jit_info, args, kwargs)
-    out_s = [None if isinstance(s, UnspecifiedValue) else s for s in p.params['out_shardings']]
-    # TODO(yashkatariya): Add `Layout` to SDS.
-    out = [api.ShapeDtypeStruct(x.shape, x.dtype, sharding=s,
-                                weak_type=x.weak_type)
-           for x, s in zip(p.params['jaxpr'].out_avals, out_s)]
-    return tree_unflatten(p.out_tree, out)
-
-  @api_boundary
-  def trace(*args, **kwargs) -> stages.Traced:
-    p, args_flat = _infer_params(fun, jit_info, args, kwargs)
-    donate_argnums = tuple(i for i, d in enumerate(p.donated_invars) if d)
-    args_info = stages.make_args_info(p.in_tree, p.in_avals, donate_argnums)
-    lower_callable = partial(_resolve_and_lower, args_flat, **p.params,
-                             pgle_profiler=None)
-    return stages.Traced(
-        p.params['jaxpr'], args_info, p.params["name"], p.out_tree,
-        lower_callable, args_flat, p.arg_names, p.num_consts)
-
-  wrapped = _cpp_pjit(fun, jit_info)
-  wrapped.lower = lower
-  wrapped.eval_shape = eval_shape
-  wrapped.trace = trace
-  return wrapped
-
-
 def make_jit(fun: Callable, in_shardings: Any, out_shardings: Any,
              donate_argnums: int | Sequence[int] | None,
              donate_argnames: str | Iterable[str] | None,
@@ -533,7 +533,7 @@ def make_jit(fun: Callable, in_shardings: Any, out_shardings: Any,
         fun, in_shardings, out_shardings, donate_argnums, donate_argnames,
         static_argnums, static_argnames, device, backend, abstracted_axes,
         keep_unused, inline, compiler_options, use_resource_env)
-  return _make_jit_wrapper(fun, jit_info)
+  return _cpp_pjit(fun, jit_info)
 
 
 class PjitParams(NamedTuple):


Test patch:
diff --git a/tests/memories_test.py b/tests/memories_test.py
index fca365702..203bc4abb 100644
--- a/tests/memories_test.py
+++ b/tests/memories_test.py
@@ -1501,8 +1501,8 @@ class ComputeOffload(jtu.BufferDonationTestCase):
     s = NamedSharding(mesh, P(), memory_kind='pinned_host')
     s_dev = s.with_memory_kind('device')
 
-    @compute_on('device_host')
     @functools.partial(jax.jit, out_shardings=(s, s_dev), donate_argnums=(0, 1))
+    @compute_on('device_host')
     def f(inp1, inp2):
       return inp1 * 2, inp2 * 2
 
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 2db75be18..f90b02cbd 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -14,7 +14,7 @@
 
 from collections import OrderedDict, namedtuple
 import re
-from functools import partial
+from functools import partial, wraps
 import logging
 import json
 import math
@@ -940,6 +940,18 @@ class PJitTest(jtu.BufferDonationTestCase):
     # Make sure this doesn't crash
     pjit(lambda x: x, in_shardings=None, out_shardings=None)(key)
 
+  def test_lower_with_wrapper_error(self):
+    @jax.jit
+    def f(x):
+      return x
+
+    self.assertAllClose(1., f(1.))
+    self.assertAllClose(1., f.lower(1.).compile()(1.))
+    wrapped_f = wraps(f)(lambda x: f(x + 1))
+
+    with self.assertRaisesRegex(AttributeError, "has no attribute 'lower'"):
+      wrapped_f.lower(1.)
+
   @jtu.with_mesh([('x', 2), ('y', 2)])
   def testLowerCompile(self):
     @partial(pjit,
diff --git a/tests/xla_metadata_test.py b/tests/xla_metadata_test.py
index 33fd7a08b..ba2120fcb 100644
--- a/tests/xla_metadata_test.py
+++ b/tests/xla_metadata_test.py
@@ -288,10 +288,10 @@ class XlaMetadataTest(jtu.JaxTestCase):
       with set_xla_metadata(a="b"):
         return (x + y, y * 2.0)
 
-    f_vmap_jaxpr = jax.make_jaxpr(jax.vmap(f2, in_axes=(0, None)))
+    f2_vmap = jax.vmap(f2, in_axes=(0, None))
     self.assertIn(
         'mhlo.frontend_attributes = {a = "b"}',
-        f_vmap_jaxpr.lower(jnp.arange(5.0), 1.0).as_text(),
+        jax.jit(f2_vmap).lower(jnp.arange(5.0), 1.0).as_text(),
     )
 
   def test_multiple_instructions(self):