Expand test case to repro a crash.

jax-ml · Dec 23, 2024 · 122aec2 · 122aec2
1 parent 6c85e54
commit 122aec2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
@@ -4325,6 +4325,8 @@ def _f(x, y):
         out_specs=P(None, 'i')))
     f(jnp.zeros((2, 16)), jnp.ones(2))
 
+    jax.jit(_f)(jnp.zeros((2, 16)), jnp.ones(2))
+
   def test_jit_trace_lower_and_compile(self):
     def f(x):
       return x * 2