jax-ml/jax · error · ValueError

Requesting collective allocations, but no explicit thread ax

Error message

Requesting collective allocations, but no explicit thread axis specified.

What it means

The interpreter validates collective (cross-thread) allocations: if any axis is marked collective but the mesh has no named thread axis (mesh.thread_name is None/empty), there is no defined group of threads to make the allocation collective over, so it raises ValueError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/jaxpr_interpret.py:93

    space: mosaic_gpu_core.MemorySpace | None,
):
  # TODO(nrink): Support more memory spaces.
  if space is not None and space not in [
      mosaic_gpu_core.MemorySpace.GMEM,
      mosaic_gpu_core.MemorySpace.SMEM,
      mosaic_gpu_core.MemorySpace.TMEM,
      mosaic_gpu_core.MemorySpace.REGS,
  ]:
    raise NotImplementedError(f"Unsupported memory space: {space}")


def _raise_if_unsupported_collective_axes(
    mesh: mosaic_gpu_core.Mesh | None,
    is_collective_by_thread_cluster_axis: tuple[bool, ...],
):
  if not mesh or not mesh.thread_name:
    if any(is_collective_by_thread_cluster_axis):
      raise ValueError(
          "Requesting collective allocations, but no explicit thread axis"
          " specified."
      )
  else:
    # Note that the leading entries in `is_collective_by_thread__cluster_axis`
    # correspond to the cluster axes, while the last entry corresponds to the
    # thread axis within a block.
    *is_collective_by_cluster_axis, is_thread_axis_collective = (
        is_collective_by_thread_cluster_axis
    )
    if any(is_collective_by_cluster_axis):
      raise ValueError(
          "Collective allocations along cluster axes are not supported."
      )
    if not is_thread_axis_collective:
      raise ValueError(
          "Scoped allocation must have the thread axis in its collective axes."
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add thread_name (and num_threads) to the Mesh: Mesh(grid=..., thread_name='tid', num_threads=N)
2. If the allocation shouldn't be shared, drop the collective axis from the allocation
3. Align collective axis flags with the actual mesh axis layout (cluster axes first, thread axis last)

Example fix

# before
mesh = Mesh(grid=(1,), num_threads=128)  # no thread_name
# after
mesh = Mesh(grid=(1,), thread_name='tid', num_threads=128)

if any(is_collective_by_thread_cluster_axis):
    assert mesh is not None and mesh.thread_name, 'collective alloc needs a named thread axis'

def mesh_has_thread_axis(mesh) -> bool:
    return mesh is not None and bool(mesh.thread_name)

When it happens

Trigger: Creating a Block/scratch with collective axes (is_collective_by_thread_cluster_axis contains True) while the kernel's Mesh lacks thread_name, in interpret mode.

Common situations: Sharing scratch across warps without declaring the thread axis in the Mesh; meshes built only with grid/cluster dims; refactor removing thread_name while keeping collective buffers.

Related errors

• Unsupported mesh type: {type(mesh)}
• Invalid grid {grid_mapping.grid} in grid_mapping: expected g
• No axis names are available. Make sure you are using `pl.cor
• Unknown cluster axis {axis_name}, available axes: {[*axis_na
• No axis names are available. Make sure you are using `pl.cor

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