[c++ jit] only set use_fastpath in cache_miss if all args are DeviceA…

…rrays

fixes #12542

Co-authored-by: Peter Hawkins <phawkins@google.com>
Co-authored-by: Kuangyuan Chen <chky@google.com>

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

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):