jax-ml/jax · error · ValueError

Barrier can only be indexed with integers or slices, got {id

Error message

Barrier can only be indexed with integers or slices, got {idx}

What it means

Barrier address computation only accepts integer-like indices (Python int, MLIR ir.Value, FragmentedArray, TypedNdArray) or slices. Any other index type (e.g. a float, a tuple, or an unrecognized object) is rejected with this ValueError.

Source

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

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

          idx = arith_dialect.muli(idx, lowering._as_index(stride))  # pylint: disable=protected-access
          if base_index is None:
            base_index = idx
          else:
            base_index = arith_dialect.addi(base_index, idx)
      case _:
        raise ValueError("Barrier does not support arbitrary transforms")
  return base_index


barrier_arrive_p = jax_core.Primitive("barrier_arrive")
barrier_arrive_p.multiple_results = True

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert the index to an int (e.g. int(i) or i.item()) before indexing the barrier.
2. Check the index expression for accidental tuple/float values; barriers are 1D and take a single integer or slice.

Example fix

# before
barriers[np.float32(i)].arrive()
# after
barriers[int(i)].arrive()

def is_valid_barrier_index(idx):
    return idx is None or isinstance(idx, (int, slice)) or hasattr(idx, 'type')  # ir.Value / FragmentedArray

When it happens

Trigger: Indexing a barrier ref with an unsupported type, e.g. barrier[1.0], barrier[(0, 1)], or a custom indexer object.

Common situations: Passing a numpy scalar or non-integer loop variable as a barrier index; typos producing nested tuples in multi-dimensional indexing.

Related errors

• Unsupported index: {v} of type {type(v)}
• Barrier does not support slice with `stride != 1`
• ref must be a reference
• Unsupported type: {x}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s

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