When the size of the remainder array is 0, don't append it to the remainder_leaves list. This fixes usage of lax.map in sharding-in-types mode.

jax/_src/lax/control_flow/loops.py

@@ -2509,19 +2509,23 @@ def fori_loop(lower, upper, body_fun, init_val):
 
 def _batch_and_remainder(x, batch_size: int):
   leaves, treedef = tree_flatten(x)
-
-  scan_leaves = []
-  remainder_leaves = []
-
-  for leaf in leaves:
-    num_batches, _ = divmod(leaf.shape[0], batch_size)
-    total_batch_elems = num_batches * batch_size
-    scan_leaves.append(leaf[:total_batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:]))
-    remainder_leaves.append(leaf[total_batch_elems:])
-
-  scan_tree = treedef.unflatten(scan_leaves)
-  remainder_tree = treedef.unflatten(remainder_leaves)
-  return scan_tree, remainder_tree
+  if not leaves:
+    return x, None
+  num_batches, remainder = divmod(leaves[0].shape[0], batch_size)
+  total_batch_elems = num_batches * batch_size
+  if remainder:
+    scan_leaves, remainder_leaves = [], []
+    for leaf in leaves:
+      scan_leaves.append(leaf[:total_batch_elems].reshape(
+          num_batches, batch_size, *leaf.shape[1:]))
+      remainder_leaves.append(leaf[total_batch_elems:])
+    return treedef.unflatten(scan_leaves), treedef.unflatten(remainder_leaves)
+  else:
+    scan_leaves = [
+        leaf[:total_batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:])
+        for leaf in leaves
+    ]
+    return treedef.unflatten(scan_leaves), None
 
 @api_boundary
 def map(f, xs, *, batch_size: int | None = None):
@@ -2576,11 +2580,14 @@ def map(f, xs):
     scan_xs, remainder_xs = _batch_and_remainder(xs, batch_size)
     g = lambda _, x: ((), api.vmap(f)(x))
     _, scan_ys = scan(g, (), scan_xs)
-    remainder_ys = api.vmap(f)(remainder_xs)
     flatten = lambda x: x.reshape(-1, *x.shape[2:])
-    ys = tree_map(
-      lambda x, y: lax.concatenate([flatten(x), y], dimension=0), scan_ys, remainder_ys,
-    )
+    if remainder_xs is not None:
+      remainder_ys = api.vmap(f)(remainder_xs)
+      ys = tree_map(
+        lambda x, y: lax.concatenate([flatten(x), y], dimension=0), scan_ys,
+        remainder_ys)
+    else:
+      ys = tree_map(flatten, scan_ys)
   else:
     g = lambda _, x: ((), f(x))
     _, ys = scan(g, (), xs)

tests/pjit_test.py

@@ -7893,6 +7893,20 @@ def test_nn_constant(self, mesh):
     self.assertArraysEqual(out, jnp.full((8, 2), -7, dtype=jnp.float32))
     self.assertEqual(out.sharding, NamedSharding(mesh, P('x', None)))
 
+  @config.numpy_rank_promotion('allow')
+  @jtu.with_explicit_mesh((2, 2), ('x', 'y'))
+  def test_lax_map(self, mesh):
+    def simple_func(w, x):
+      return jnp.sum(w * x, axis=-1)
+
+    w = jax.device_put(np.arange(4, dtype=np.float32), P('x'))
+    x = jax.device_put(np.ones((4, 2, 4), dtype=np.float32),
+                       P(None, 'y', None))
+
+    jax.lax.map(lambda _x: simple_func(w, _x), x)  # doesn't crash
+
+    jax.lax.map(lambda _x: simple_func(w, _x), x, batch_size=2)  # doesn't crash
+
 
 @jtu.pytest_mark_if_available('multiaccelerator')
 class PJitErrorTest(jtu.JaxTestCase):