Add an impl rule to empty2. Fixes #32404

PiperOrigin-RevId: 816219967

Author: yashk2810

jax/_src/lax/lax.py

@@ -9090,6 +9090,7 @@ def _empty_batcher(axis_data, vals_in, dims_in, *, shape, dtype, out_sharding):
 def empty2(dtype, *, memory_space):
   return empty2_p.bind(dtype=dtype, memory_space=memory_space)
 empty2_p = core.Primitive('empty2')
+dispatch.simple_impl(empty2_p)
 
 def _empty2_abstract_eval(*, dtype, memory_space):
   return core.ShapedArray((), dtype, memory_space=memory_space)

tests/python_callback_test.py

@@ -604,6 +604,17 @@ def f(x):
     x = np.arange(8, dtype=dtype)
     np.testing.assert_array_equal(jax.jit(f)(x), np.arange(8, dtype=dtype))
 
+  def test_pure_callback_sequential_vmap_method_eval_jaxpr(self):
+    def f(x):
+      return jax.pure_callback(
+          lambda x: x, jax.ShapeDtypeStruct(shape=(), dtype=jnp.float32),
+          x, vmap_method="sequential")
+
+    jaxpr = jax.make_jaxpr(lambda: jax.vmap(f)(
+        jnp.zeros(100, dtype=jnp.float32)))()
+    with jax.ensure_compile_time_eval():
+      jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)  # doesn't crash
+
   @parameterized.parameters("int2", "int4", "uint2", "uint4", "float4_e2m1fn")
   def test_subbyte_results(self, dtype: str):
     if "2" in dtype and jtu.test_device_matches(["tpu"]):