jax-ml/jax · error · NotImplementedError

Barrier does not support slice with `stride != 1`

Error message

Barrier does not support slice with `stride != 1`

What it means

When computing a barrier's base address, the indexing slice must have stride 1 because the lowering computes a single linear base offset. Slices with a step (e.g. barrier_ref[0::2]) cannot be represented.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:1489

      collective_axes=collective_axes,
      leader_tracked=leader_tracked,
  )
  return None


def _get_barrier_base_index(aval, transforms) -> ir.Value | None:
  if not transforms:
    return None
  strides = list(pallas_utils.strides_from_shape(aval.shape))
  base_index: ir.Value | None = None
  while transforms:
    match transforms:
      case [indexing.NDIndexer() as indexer, *transforms]:
        num_int_idxs = 0
        for i, (idx, stride) in enumerate(zip(indexer.indices, strides[:])):
          if isinstance(idx, indexing.Slice):
            if idx.stride != 1:
              raise NotImplementedError(
                  "Barrier does not support slice with `stride != 1`"
              )
            idx = idx.start
          else:
            # This dimension is absent for any corresponding `NDIndexer`s, so
            # we remove the corresponding stride.
            strides.pop(i - num_int_idxs)
            num_int_idxs += 1

          if isinstance(
              idx, (int, ir.Value, mgpu.FragmentedArray, literals.TypedNdArray)
          ):
            idx = lowering._as_index(idx)  # pylint: disable=protected-access
          else:
            raise ValueError(
                "Barrier can only be indexed with integers or slices, got"
                f" {idx}"
            )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Index barriers individually or with contiguous slices (stride 1), e.g. barrier[i] or barrier[0:4].
2. Restructure so you don't need strided barrier selection (allocate separate barriers per subgroup).

Example fix

# before
barriers[0:8:2].arrive()
# after
for i in range(0, 8, 2):
    barriers[i].arrive()

def check_slice(idx):
    if isinstance(idx, slice):
        assert idx.step in (None, 1), f'stride {idx.step} unsupported on barriers'
for i in indices: check_slice(i)

When it happens

Trigger: Indexing a barrier/mma accumulator ref with a strided slice, e.g. barrier[0:4:2], in any op that lowers through _get_barrier_base_index (barrier.arrive/wait/test, async_store_smem, tcgen05_mma, copies).

Common situations: Using Python slice syntax with a step on barrier refs while trying to select every other barrier in an array of barriers.

Related errors

• Barrier can only be indexed with integers or slices, got {id
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Unsupported barrier type: {type(barrier)}
• tcgen05_mma only allows arriving on a Barrier
• tcgen05_mma only allows arriving on a Barrier that orders te

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