jax-ml/jax · error · NotImplementedError

Indexing barrier with {transforms} not supported in GPU inte

Error message

Indexing barrier with {transforms} not supported in GPU interpret mode

What it means

When a barrier is indexed in interpret mode, only a single indexing transform (an NDIndexer) is supported. Multiple transforms or non-index transforms cannot be mapped to a barrier allocation key by the interpreter.

Source

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

# 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:
    return None
  if not hasattr(transforms, "__len__") or len(transforms) != 1:
    raise NotImplementedError(
        f"Indexing barrier with {transforms} not supported in GPU interpret"
        " mode"
    )
  if not isinstance(transforms[0], indexing.NDIndexer):
    raise ValueError(f"Expected an `NDIndexer`, but got {transforms[0]}")
  if len(transforms[0].indices) == 1:
    return transforms[0].indices[0]
  return tuple(transforms[0].indices)


def _get_barrier_allocation_key_from_inval(
    inval, transforms_treedef, transforms_leaves
) -> jax.Array:
  # `inval` is expected to correspond to a barrier. Since we are interpreting,
  # `inval` will in fact contain the allocation key (which is a Jax array) for
  # the barrier.
  allocation_key_as_array = inval

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Index the barrier once with a single NDIndexer (slice/index) and pass it directly
2. Allocate separate scalar barriers instead of an array of barriers when per-element access is needed
3. Simplify barrier transforms before interpret mode runs

Example fix

// before
op(barrier[0:2][1])  # chained transforms
// after
op(barrier[1])      # single NDIndexer

# ensure barrier refs carry at most one indexing transform
assert len(ref.transforms or ()) <= 1, 'multiple transforms on barrier unsupported in interpret mode'

def is_simple_indexed_barrier(ref) -> bool:
    ts = getattr(ref, 'transforms', ()) or ()
    return len(ts) <= 1 and all(
        type(t).__name__ == 'NDIndexer' for t in ts)

When it happens

Trigger: Passing a barrier reference with two or more transforms (e.g. slicing then unswizzling), or a transforms tuple without __len__, into copy_gmem_to_smem, barrier_arrive, barrier_wait, or tcgen05 ops during interpretation.

Common situations: Composing multiple reference transformations on a barrier allocated as an array of barriers; chaining slicing with other transforms generated by compiler plumbing.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unsupported transform: {type(transform)}
• Expected an `NDIndexer`, but got {transforms[0]}

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