jax-ml/jax · error · NotImplementedError

Batching over dynamic grid values is not supported yet.

Error message

Batching over dynamic grid values is not supported yet.

What it means

Raised by the Mosaic Pallas pipeline when vmap-style batching is applied to a kernel whose grid depends on dynamic (traced) values. The batching transform can replicate block mappings for a known batch size, but it cannot reconstruct index maps when grid extents themselves are computed at trace time, so it refuses.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:2696

  dimension_semantics = (PARALLEL,) + dimension_semantics
  all_args: EmitPipelinePrimitiveArgs = args_tree.unflatten(args_flat)

  _, dynamic_dims, _, _, flat_ref_dims, alloc_dims = jax_util.split_list(dims, [
      len(all_args.all_index_map_consts),
      len(all_args.dynamic_grid_spec),
      int(all_args.has_core_id),
      len(all_args.body_consts),
      len(all_args.refs_flat)])

  if any(d is not None for d in alloc_dims):
    raise NotImplementedError(
        "Batching over custom allocations is not supported yet."
    )

  batch_size = axis_data.size

  if any(d is not None for d in dynamic_dims):
    raise NotImplementedError(
      "Batching over dynamic grid values is not supported yet.")

  batched_block_mappings = map(
      functools.partial(_batch_block_mapping, grid_mapping, batch_size),
      map(_ref_to_value_aval, all_args.refs_flat),
      flat_ref_dims, grid_mapping.block_mappings)

  index_map_tree_args, index_map_tree_kwargs = (
      grid_mapping.index_map_tree.unflatten(grid_mapping.index_map_avals))
  assert not index_map_tree_kwargs
  batched_index_map_args = (
      pallas_core.index_map_grid_aval, *index_map_tree_args)
  batched_index_map_avals, batched_index_map_tree = tree_util.tree_flatten(
      (batched_index_map_args, {}))

  axis_size_is_dynamic = not isinstance(batch_size, int)
  new_grid_dim = (pallas_core.dynamic_grid_dim
                  if axis_size_is_dynamic else batch_size)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the grid static: compute grid extents from Python ints or from shapes that do not depend on the batched axis (e.g. hardcode the per-example grid and let vmap add a leading grid dimension)
2. Restructure so the batch appears as an extra grid dimension that the kernel indexes via program_id instead of a dynamic grid value
3. Move batching inside the kernel body (manual loop over the batch within one launch) instead of using vmap on the outside
4. Check for a newer JAX version — batching support for Pallas is actively expanding

Example fix

# before
grid = (x.shape[0] // 128,)  # x.shape[0] depends on vmapped axis
kern = pallas_call(fn, out_shape=..., grid=grid)
out = jax.vmap(lambda x: kern(x))(xs)

# after
def fn(x_ref, o_ref):
  ...
# static per-example grid; batch becomes grid dim
kern = pallas_call(fn, out_shape=..., grid=(xs.shape[0], xs.shape[1] // 128))
out = kern(xs)

def check_grid_static(kernel_fn, *args, **kwargs):
  import jax
  try:
    jax.make_jaxpr(lambda *a: kernel_fn(*a, **kwargs))(*args)
  except Exception as e:
    raise RuntimeError(f"kernel trace failed (dynamic grid?)") from e
# ensure grid extents are Python ints:
assert all(isinstance(g, int) for g in grid), f"dynamic grid: {grid}"

When it happens

Trigger: Calling jax.vmap (or another batching transform) over a function that invokes a pallas_call/mosaic kernel whose grid tuple contains dynamic values (e.g. grid=(x.shape[0],) where the shape depends on the batched axis, or a grid built from a traced computation).

Common situations: Batching a TPU Pallas kernel that computes its grid from an input whose leading dim is being vmapped; stacking per-example kernels where grid=(num_examples,) and num_examples is derived from the batch.

Related errors

• Batching over custom allocations is not supported yet.
• dma_start not implemented in LoJAX yet.
• vmapping pallas_call with no arguments.
• ragged_dot vmap over any dim but 0 - NYI
• reduce_window batching is not implemented for initial values

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/9d8e4f09adb596f8. Report an issue: GitHub.