jax-ml/jax · error · ValueError

Scoped allocation must have the thread axis in its collectiv

Error message

Scoped allocation must have the thread axis in its collective axes.

What it means

Mosaic GPU interpret mode requires scoped allocations to be collective along the thread axis. Since the interpreter executes one thread's view of the block, a scoped allocation that is not collective on the thread axis would have undefined semantics, so it is rejected.

Source

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

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


# TODO(nrink): Try unifying this function with `_extract_barrier_slice_base`
# from `jax._src.pallas.mosaic_gpu.primitives`.
def _get_index_for_barrier_allocation_key(
    transforms_treedef, transforms_leaves,
) -> indexing.DimIndexer | None:
  # TODO(nrink): The working out of `transforms` and the returned index below
  # may need tidying up. Specifically, GPU interpret mode should correctly
  # support legal ways to index into barriers. (Here, 'legal' is to be read as
  # 'allowed by the Pallas GPU semantics'.)
  if transforms_treedef is None:
    return None
  transforms = jax.tree.unflatten(transforms_treedef, transforms_leaves)

  if not transforms:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Add the thread axis name (from mesh.axis_names[-1]) to collective_axes
2. Pass collective_axes=None/empty only if the API permits non-collective scope, otherwise include thread axis
3. Check WarpMesh construction to confirm the intended thread axis name

Example fix

// before
run_scoped(buf, collective_axes=())
// after
run_scoped(buf, collective_axes=('thread',))

thread_axis = mesh.axis_names[-1] if mesh is not None else None
if thread_axis is not None:
    assert thread_axis in (collective_axes or ()), 'thread axis must be collective'

When it happens

Trigger: Calling run_scoped with collective_axes that omit the thread axis name of the current WarpMesh (for non-REGS memory spaces, which route through _raise_if_unsupported_collective_axes).

Common situations: Passing an empty collective_axes tuple or only cluster axis names; refactoring kernels and accidentally dropping the thread axis from collective_axes.

Related errors

• collective_axes is not supported in pallas_call. Use plgpu.k
• Only thread-collective allocations are supported in run_scop
• WGMMA accumulators can only be allocated non-collectively. H
• Only thread-collective allocations are supported in multithr
• Only collective_axes that include all JAX device mesh axes

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