jax-ml/jax · error · ValueError

Attempting to operate on barrier without indexing, but…

Error message

Attempting to operate on barrier without indexing, but `num_barriers = {num_barriers}`

What it means

Consistency guard in the GPU interpreter's barrier key handling: a barrier operation without an index transform is only valid when there is exactly one barrier (num_barriers == 1); with multiple barriers the index is required, so this raises NotImplementedError. The input at fault is the unindexed barrier ref in a multi-barrier kernel.

Solutions

1. Index the barrier: barrier_arrive(barrier[i], ...) with the stage index
2. Allocate with num_barriers=1 if only one barrier is needed
3. Check loop bodies that use barrier arrays to ensure the loop induction variable indexes the barrier

Example fix

// before
barrier_arrive(barrier, sem)
// after
barrier_arrive(barrier[i], sem)

if num_barriers > 1:
    assert barrier_transforms_present, f'index barrier: num_barriers={num_barriers}'

When it happens

Trigger: Calling barrier_arrive, barrier_wait, copy_gmem_to_smem, tcgen05_mma, or tcgen05_commit_arrive with an un-indexed barrier ref where the barrier was allocated with num_barriers > 1.

Common situations: Allocating a barrier array for pipelined copies (multiple buffers) but forgetting to index it per pipeline stage; refactoring from a single barrier to multiple barriers.

Related errors

• arrive_expect_tx is only supported on Hopper+ hardware
• Barrier can only be indexed with integers or slices, got
• Barrier does not support slice with `stride != 1`
• Barrier offset is out of bounds
• Barriers are required for TMA GMEM -> SMEM copies

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

Appendix: source

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

  allocation_key_as_array = inval

  # Assert to check internal consistency: `allocation_key_as_array` should be
  # at least a 2D array, and the size of the last dimension is 5 (which matches the
  # fields count of HostAllocationKey).
  assert len(allocation_key_as_array.shape) >= 2
  assert (
      allocation_key_as_array.shape[-1:]
      == gpu_callbacks.HostAllocationKey.shape_and_dtype().shape
  )
  num_barriers = math.prod(allocation_key_as_array.shape[:-1])

  index = _get_index_for_barrier_allocation_key(
      transforms_treedef, transforms_leaves
  )

  if index is None:
    if num_barriers != 1:
      raise ValueError(
          "Attempting to operate on barrier without indexing, but"
          f" `num_barriers = {num_barriers}`"
      )
    idx = (0,) * (len(allocation_key_as_array.shape) - 1)
    return allocation_key_as_array[idx]
  else:
    return allocation_key_as_array[index]


def _get_num_threads_sharing_collective_allocation(
    axes_dims: tuple[int, ...],
    is_last_thread_axis_collective: bool,
) -> int:
  """Returns the number of threads that share a collective allocation."""
  if is_last_thread_axis_collective:
    return axes_dims[-1]
  else:
    return 1

View on GitHub (pinned to 1e1c6a8fc0)