jax-ml/jax · error · ValueError

Unsupported mesh type: {type(mesh)}

Error message

Unsupported mesh type: {type(mesh)}

What it means

When interpreting a Mosaic GPU kernel, the interpreter derives grid/cluster/thread dimensions from the kernel's mesh. It only understands mosaic_gpu_core.Mesh (and the no-mesh case); any other mesh type is rejected.

Source

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

    # entries in the grid should be ints.
    assert isinstance(x, int)
    result.append(x)
  return tuple(result)


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`."""

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the mesh as mosaic_gpu_core.Mesh (or None with a plain int grid)
2. Ensure jax and any mosaic plugin versions match
3. Check the printed type in the message to find which object leaked through

Example fix

# before
grid = some_tpu_mesh  # wrong mesh type
# after
from jax._src.pallas.mosaic_gpu.core import Mesh
grid = Mesh(grid=(g,), thread_name='tid', num_threads=128)

def is_supported_mesh(mesh) -> bool:
    from jax._src.pallas.mosaic_gpu import core
    return mesh is None or isinstance(mesh, core.Mesh)

When it happens

Trigger: Passing a grid mapping whose mesh is of an unexpected type (not None, not a GridMapping-derived int grid, not mosaic_gpu_core.Mesh) — usually from API misuse or version mismatch.

Common situations: Mixing TPU pallas Mesh types with Mosaic GPU kernels; stale jax versions where mesh classes moved; custom grid wrapper objects.

Related errors

• {axis} mixes JAX mesh and Pallas mesh grid axes
• {non_mesh_axes}
• `thread_id` must be zero when allocating a buffer for all th
• `block_id` must be zero when allocating a buffer for all thr
• Out-of-bounds read of {allocation_key}: reading [{read_range

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