[shape_poly] Fix handling of shape polymorphism for pallas_call_batching

The previous code was not handling the case of a symbolic batch
dimension (one of the most common uses of shape polymorphism)

Author: gnecula

jax/_src/pallas/core.py

@@ -1163,7 +1163,8 @@ def get_grid_mapping(
     dim_check : Any = jax_core.is_constant_dim  # type: ignore[no-redef]
   assert all(i is None or dim_check(i) for i in grid_spec.grid)
   grid_mapping_grid = tuple(
-      dynamic_grid_dim if d is None else d for d in grid_spec.grid
+      dynamic_grid_dim if (d is None or not jax_core.is_constant_dim(d)) else d
+      for d in grid_spec.grid
   )
   # The inputs for the index maps
   index_map_avals = (

jax/_src/pallas/pallas_call.py

@@ -717,8 +717,11 @@ def temp_f(*args):
   batched_index_map_avals, batched_index_map_tree = tree_util.tree_flatten(
       (batched_index_map_args, {}))
 
+  axis_size_is_dynamic = not jax_core.is_constant_dim(axis_size)
+  new_grid_dim = pallas_core.dynamic_grid_dim if axis_size_is_dynamic else axis_size
+
   batched_grid_mapping = grid_mapping.replace(
-      grid=(axis_size, *grid_mapping.grid),
+      grid=(new_grid_dim, *grid_mapping.grid),
       block_mappings=tuple(batched_block_mappings),
       index_map_avals=tuple(batched_index_map_avals),
       index_map_tree=batched_index_map_tree,
@@ -729,7 +732,7 @@ def temp_f(*args):
   # Avoid scaling the cost estimate by the batch size if the batch size is a
   # dynamic shape (DimExpr).
   # https://docs.jax.dev/en/latest/export/shape_poly.html#computing-with-dimension-variables
-  if cost_estimate is not None and isinstance(axis_size, int):
+  if cost_estimate is not None and not axis_size_is_dynamic:
     batched_cost_estimate = CostEstimate(
         flops=cost_estimate.flops * axis_size,
         bytes_accessed=cost_estimate.bytes_accessed * axis_size,
@@ -762,6 +765,8 @@ def temp_f(*args):
     with jax_core.remove_explicit_mesh_axis_names(ema):
       bind = shard_map(bind, out_specs=P(ema), axis_names=set(ema))
 
+  if axis_size_is_dynamic:
+    dynamic_grid_args = [axis_size, *dynamic_grid_args]
   out = bind(*dynamic_grid_args, *args)
   return out, (0,) * len(out)
 

tests/pallas/export_pallas_test.py

@@ -97,9 +97,6 @@ def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
 
 class ExportTestWithMosaicTPU(jtu.JaxTestCase):
   def test_dynamic_shapes_export(self):
-    if jtu.device_under_test() != "tpu":
-      self.skipTest("Mosaic TPU test only runs on TPU")
-
     def add_vectors_kernel(x_ref, y_ref, o_ref):
       block_b = x_ref.shape[0]
 
@@ -142,11 +139,45 @@ def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
     with pallas_export_experimental(dynamic_shapes=True):
       f_k = f_e(x_shape, y_shape)
 
-    print(f_k.mlir_module())
     self.assertRegex(
         f_k.mlir_module(),
         r"stablehlo.custom_call @tpu_custom_call.+kernel_name\s*=\s*\"my_custom_kernel_name\"")
 
+  def test_export_vmap(self):
+    def add_vectors_kernel(x_ref, y_ref, o_ref):
+      o_ref[...] = x_ref[...] + y_ref[...]
+
+    def add_vectors(x, y):
+      block_size = 128
+      # Grid depends on input shape, which will be symbolic
+      grid = (x.shape[0] // block_size, x.shape[1] // block_size)
+      block_spec = pl.BlockSpec((block_size, block_size), lambda i, j: (i, j))
+      return pl.pallas_call(
+          add_vectors_kernel,
+          out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
+          grid=grid,
+          in_specs=[block_spec, block_spec],
+          out_specs=block_spec
+      )(x, y)
+
+    b, m, n = jax.export.symbolic_shape("b,m,n")
+    x_info = jax.ShapeDtypeStruct((b, m, n), jnp.float32)
+
+    exporter = jax.export.export(jax.jit(jax.vmap(add_vectors)),
+                                 platforms=["tpu"])
+
+    with pallas_export_experimental(dynamic_shapes=True):
+      exp = exporter(x_info, x_info)  # No crash
+
+    if jtu.device_under_test() == "tpu":
+      x = y = jnp.ones((4, 128, 128))
+      res = exp.call(x, y)
+      self.assertAllClose(res, x + y)
+
+      x = y = jnp.ones((4, 192, 192))  # Not multiple of 128
+      res = exp.call(x, y)
+      self.assertAllClose(res, x + y)
+
 
 class ExportTestWithMosaicGpu(ExportTestWithTriton):