Add dynamic grid support to emit_pipeline

PiperOrigin-RevId: 623393190

Author: sharadmv

jax/_src/pallas/core.py

@@ -37,8 +37,9 @@
 # mypy: ignore-errors
 
 partial = functools.partial
-Grid = tuple[Union[int, None], ...]  # None indicates that the bound is dynamic.
-StaticGrid = tuple[int, ...]  # None indicates that the bound is dynamic.
+Grid = tuple[Union[int, jax_core.Array, None], ...]  # None indicates that the bound is dynamic.
+DynamicGrid = tuple[Union[int, jax_core.Array], ...]
+StaticGrid = tuple[int, ...]
 split_list = util.split_list
 
 map, unsafe_map = util.safe_map, map

jax/_src/pallas/mosaic/pipeline.py

@@ -120,7 +120,9 @@ def _tree_map_with_kwargs(f, *args, **kwargs):
   )
 
 
-def _get_next_indices(grid: core.StaticGrid, indices: GridIndices) -> GridIndices:
+def _get_next_indices(
+    grid: core.DynamicGrid, indices: GridIndices
+) -> GridIndices:
   """Takes a grid and current indices and returns the next indices.
 
   grid: (3, 4, 5)
@@ -135,7 +137,7 @@ def _get_next_indices(grid: core.StaticGrid, indices: GridIndices) -> GridIndice
     Incremented indices.
   """
   next_indices = []
-  carry = True
+  carry: Union[bool, jax.Array] = True
   for dim_size, index in reversed(list(zip(grid, indices))):
     i = jnp.where(carry, index + 1, index)
     carry = dim_size == i
@@ -494,7 +496,7 @@ def __call__(
 def emit_pipeline_with_allocations(
     body: PipelineBody,
     *,
-    grid: core.StaticGrid,
+    grid: core.DynamicGrid,
     in_specs: PipelineBlockSpecs,
     out_specs: PipelineBlockSpecs,
     should_accumulate_out: Union[Sequence[bool], Any] = False,
@@ -1236,7 +1238,7 @@ def set_buffer_ref(buffer_ref, buffer):
 def emit_pipeline(
     body: PipelineBody,
     *,
-    grid: core.StaticGrid,
+    grid: core.DynamicGrid,
     in_specs: PipelineBlockSpecs,
     out_specs: PipelineBlockSpecs,
     should_accumulate_out: Union[Sequence[bool], Any] = False,