jax-ml/jax · error · ValueError

Expected an `NDIndexer`, but got {transforms[0]}

Error message

Expected an `NDIndexer`, but got {transforms[0]}

What it means

Barrier indexing in interpret mode only accepts transforms that are indexing.NDIndexer instances. Any other transform type on a barrier reference is rejected with this ValueError.

Source

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

    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

  # 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 (

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove non-indexer transforms from barrier references
2. Apply unswizzle/untile transforms only to regular SMEM buffers, never to ClusterBarrierType refs
3. Verify the ref passed to barrier ops is the raw allocation

Example fix

// before
op(unswizzle(barrier_ref))
// after
op(barrier_ref)

from jax._src.pallas import indexing
def barrier_transforms_are_valid(transforms) -> bool:
    return all(isinstance(t, indexing.NDIndexer) for t in transforms)

When it happens

Trigger: Applying a non-NDIndexer transform (e.g. UnswizzleRef, UntilingTransform, or any custom Transform) to a barrier before barrier_arrive/wait, copy_gmem_to_smem, or tcgen05 MMA ops in interpret mode.

Common situations: Copying transform pipelines written for SMEM buffers onto barrier refs; generic code that applies unswizzle transforms to all refs in a kernel.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Indexing barrier with {transforms} not supported in GPU inte
• Attempting to operate on barrier without indexing, but `num_
• Expected a single barrier, got a barrier reference with shap
• copy_gmem_to_smem with a barrier is only supported Hopper an
• inline_mgpu_p only supports plgpu.ShapeDtypeStruct return ty

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