Repository: google/jax

Issue:
C++ jit dispatch failure due to stale cache entries
The following program
```
import jax

@jax.jit
def a(x):
    return ()

@jax.jit
def b(x):
    return a(x)


@jax.jit
def g(x):
    return x, x

@jax.jit
def h(x):
    return g(x)


jaxpr = jax.make_jaxpr(h)(7)
jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 7)

b(8)
```

fails with:
```
Traceback (most recent call last):
  File "/Users/phawkins/p/jax/t.py", line 27, in <module>
    b(8)
  File "/Users/phawkins/p/jax/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/_src/api.py", line 603, in cache_miss
    out_flat = xla.xla_call(
  File "/Users/phawkins/p/jax/jax/core.py", line 1953, in bind
    return call_bind(self, fun, *args, **params)
  File "/Users/phawkins/p/jax/jax/core.py", line 1969, in call_bind
    outs = top_trace.process_call(primitive, fun_, tracers, params)
  File "/Users/phawkins/p/jax/jax/core.py", line 695, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/Users/phawkins/p/jax/jax/_src/dispatch.py", line 233, in _xla_call_impl
    compiled_fun = xla_callable(fun, device, backend, name, donated_invars,
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 295, in memoized_fun
    ans = call(fun, *args)
  File "/Users/phawkins/p/jax/jax/_src/dispatch.py", line 367, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars, False,
  File "/Users/phawkins/p/jax/jax/_src/profiler.py", line 313, in wrapper
    return func(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/_src/dispatch.py", line 453, in lower_xla_callable
    jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
  File "/Users/phawkins/p/jax/jax/_src/profiler.py", line 313, in wrapper
    return func(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 2080, in trace_to_jaxpr_final2
    jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(fun, main, debug_info)
  File "/Users/phawkins/p/jax/jax/interpreters/partial_eval.py", line 2030, in trace_to_subjaxpr_dynamic2
    ans = fun.call_wrapped(*in_tracers_)
  File "/Users/phawkins/p/jax/jax/linear_util.py", line 168, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/Users/phawkins/p/jax/t.py", line 12, in b
    return a(x)
  File "/Users/phawkins/p/jax/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/phawkins/p/jax/jax/_src/api.py", line 624, in cache_miss
    fastpath_data = _device_array_use_fast_path(execute, out_pytree_def, args_flat, out_flat)
  File "/Users/phawkins/p/jax/jax/_src/api.py", line 557, in _device_array_use_fast_path
    assert len(avals) == len(out_flat)
jax._src.traceback_util.UnfilteredStackTrace: AssertionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/phawkins/p/jax/t.py", line 27, in <module>
    b(8)
  File "/Users/phawkins/p/jax/t.py", line 12, in b
    return a(x)
AssertionError
```

`eval_jaxpr` leaves `most_recent_entry` populated on the `xla_callable` cache for `g`, since it uses `xla_call`'s `impl` rule but never steals the cache entry.

Later, the `jit` cache miss logic gets confused because it finds the stale entry for `g` when it calls `a`. `a` doesn't get a `most_recent_entry` because it is being traced into a jaxpr and its `impl` rule never runs.


Discussion:

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -537,6 +537,8 @@ def _device_array_use_fast_path(execute, out_pytree_def, args_flat, out_flat):
       not execute.args[5] and not execute.args[6] and
       # Has no host callbacks
       not execute.args[8] and
+      # impl rule must have been called, i.e. top trace is an EvalTrace
+      isinstance(core.find_top_trace(args_flat), core.EvalTrace) and
       # Not supported: ShardedDeviceArray
       all(device_array.type_is_device_array(x) for x in out_flat) and
       # Not supported: dynamic shapes


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -3872,6 +3872,53 @@ def test_jit_negative_static_argnums(self):
     g = jax.jit(lambda x, y: x * y, static_argnums=-1)
     g(1, 2)  # doesn't crash
 
+  def test_fastpath_cache_confusion(self):
+    # https://github.com/google/jax/issues/12542
+    @jax.jit
+    def a(x):
+      return ()
+
+    @jax.jit
+    def b(x):
+      return a(x)
+
+
+    @jax.jit
+    def g(x):
+      return x, x
+
+    @jax.jit
+    def h(x):
+      return g(x)
+
+    jaxpr = jax.make_jaxpr(h)(7)
+    jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 7)
+
+    b(8)  # don't crash
+
+  def test_fastpath_cache_confusion2(self):
+    @jax.jit
+    def a():  # note nullary function, still staged out though
+      return ()
+
+    @jax.jit
+    def b(x):
+      return a()
+
+
+    @jax.jit
+    def g(x):
+      return x, x
+
+    @jax.jit
+    def h(x):
+      return g(x)
+
+    jaxpr = jax.make_jaxpr(h)(7)
+    jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, 7)
+
+    b(8)  # don't crash
+
 
 @jtu.with_config(jax_experimental_subjaxpr_lowering_cache=True)
 class SubcallTraceCacheTest(jtu.JaxTestCase):