Reverts a01fbc2

PiperOrigin-RevId: 814008079

Author: yashk2810

jax/_src/lax/lax.py

@@ -73,7 +73,7 @@
 from jax._src.typing import Array, ArrayLike, DimSize, DuckTypedArray, DType, DTypeLike, Shape
 from jax._src.util import (cache, canonicalize_axis,
                            safe_map, safe_zip, split_list, weakref_lru_cache,
-                           foreach)
+                           foreach, tuple_insert)
 
 _max = builtins.max
 _min = builtins.min
@@ -9071,6 +9071,15 @@ def _empty_lower(ctx, *, shape, dtype, out_sharding):
   return [mlir.lower_with_sharding_in_types(ctx, out, phys_aval)]
 mlir.register_lowering(empty_p, _empty_lower)
 
+def _empty_batcher(axis_data, vals_in, dims_in, *, shape, dtype, out_sharding):
+  batched_shape = tuple_insert(shape, 0, axis_data.size)
+  batched_out_sharding = (
+      None if out_sharding is None else
+      batching.get_sharding_for_vmap(axis_data, out_sharding, 0))
+  y = empty_p.bind(shape=batched_shape, dtype=dtype,
+                   out_sharding=batched_out_sharding)
+  return y, 0
+batching.fancy_primitive_batchers[empty_p] = _empty_batcher
 
 # TODO(yashkatariya): Delete `empty2` and replace scan's usage with `empty` once
 # AllocateBuffer issues are fixed

tests/api_test.py

@@ -4467,6 +4467,18 @@ def test_lax_real_empty(self):
     self.assertEqual(out.shape, (2, 2))
     self.assertEqual(out.dtype, jnp.float32)
 
+  @jtu.run_on_devices('gpu', 'tpu')
+  def test_lax_empty_vmap(self):
+    inp = np.arange(8, dtype=jnp.int32).reshape(4, 2)
+
+    def f(x):
+      return jax.lax.empty(x.shape, x.dtype)
+
+    f = jax.jit(jax.vmap(f))
+    f(inp)  # doesn't crash
+    lowered_text = f.lower(inp).as_text()
+    self.assertIn('@AllocateBuffer() : () -> tensor<4x2xi32>', lowered_text)
+
   def test_leaked_tracer_issue_7613(self):
     # from https://github.com/jax-ml/jax/issues/7613
     import numpy.random as npr