jax-ml/jax · error · NotImplementedError

Invalid grid {grid_mapping.grid} in grid_mapping: expected g

Error message

Invalid grid {grid_mapping.grid} in grid_mapping: expected grid to have the same size as {reconstructed_grid}

What it means

Consistency check: the product of the grid mapping's dimensions must equal the product of the reconstructed (grid × cluster × num_threads) dims from the mesh. A mismatch means the mesh and the declared grid describe different amounts of parallelism.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py:76


def _get_grid_and_cluster_dims_and_num_threads(
    grid_mapping: pallas_core.GridMapping, mesh: mosaic_gpu_core.Mesh | None
) -> tuple[tuple[int, ...], tuple[int, ...], int]:
  if not mesh:
    num_threads = 1
    cluster_dims = ()
    grid_dims = _get_grid_bounds(grid_mapping)
  elif isinstance(mesh, mosaic_gpu_core.Mesh):
    num_threads = int(mesh.num_threads or 1)
    cluster_dims = tuple(mesh.cluster) if mesh.cluster is not None else ()
    grid_dims = tuple(mesh.grid)
  else:
    raise ValueError(f"Unsupported mesh type: {type(mesh)}")

  reconstructed_grid = grid_dims + cluster_dims + (num_threads,)
  if math.prod(_get_grid_bounds(grid_mapping)) != math.prod(reconstructed_grid):
    raise NotImplementedError(
        f"Invalid grid {grid_mapping.grid} in grid_mapping: expected grid to"
        f" have the same size as {reconstructed_grid}"
    )

  return grid_dims, cluster_dims, num_threads


def _allocate_buffers_for_inputs(
    token: jax.Array,
    device: memory.Device,
    invars: Sequence[Any],
    inputs: Sequence[jax.Array],
) -> tuple[jax.Array, list[jax.Array]]:
  """Allocates `GMEM` buffers for the `inputs` of a `pallas_call`."""
  # TODO(nrink): This code is a simplified version to the corresponding TPU
  # interpreter code. Eventually, we should merge the two.
  input_buffer_keys = []
  for var, value in safe_zip(invars, inputs):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make the products match: ensure prod(grid_mapping.grid) == prod(mesh.grid) * prod(cluster) * num_threads
Usually set the pallas_call grid from the mesh: grid=tuple(mesh.grid) (times cluster dims)
Recompute launch dims after changing num_threads
Check the printed expected tuple in the message and align to it

Example fix

# before
kernel = pallas_call(fn, out, grid=(8,))
kernel(..., grid=Mesh(grid=(4,), num_threads=128, ...))  # mismatch
# after
mesh = Mesh(grid=(4,), num_threads=128, thread_name='tid')
kernel = pallas_call(fn, out, grid=tuple(mesh.grid))

Defensive patterns

Strategy: validation

Validate before calling

import math
expected = tuple(mesh.grid) + (tuple(mesh.cluster) if mesh.cluster else ()) + (int(mesh.num_threads or 1),)
assert math.prod(grid_mapping.grid) == math.prod(expected)

Prevention

Derive pallas_call grid from the Mesh to keep products consistent
Re-validate launch dims after changing num_threads

When it happens

Trigger: Declaring a pallas_call grid whose total size differs from grid*cluster*threads of the Mosaic Mesh — e.g. grid=(8,) but Mesh(grid=(4,), num_threads=128) (or similar product mismatch).

Common situations: Editing one of grid/mesh/num_threads without updating the others; copy-paste kernels with mismatched launch configs; migrating kernels between thread counts.

Related errors

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