jax-ml/jax · error · NotImplementedError

Only single indexer is supported.

Error message

Only single indexer is supported.

What it means

Raised by the Pallas Triton lowering rule for atomic operations when the reference being atomically updated is indexed by more than one indexer (i.e., a chained/multi-dimensional indexing that produces multiple NDIndexer elements). The Triton backend can only compute a single flat pointer per atomic op, so compound indexing cannot be lowered.

Source

Thrown at jax/_src/pallas/triton/primitives.py:461

@lowering.register_lowering(atomic_rmw_p)
def _atomic_lowering_rule(
    ctx: lowering.LoweringRuleContext,
    *args_flat,
    args_tree,
    atomic_type: AtomicOpType,
):
  block_info, *_ = ctx.block_infos
  assert block_info is not None
  ptr, indexers, val, mask = args_tree.unflatten(args_flat)
  *_, value_aval, mask_aval = args_tree.unflatten(ctx.avals_in)
  indexers = list(indexers)
  if not indexers or not isinstance(indexers[-1], indexing.NDIndexer):
    ref_aval = state.transform_type(indexers, ctx.avals_in[0])
    assert isinstance(ref_aval, state.AbstractRef)
    indexers.append(indexing.NDIndexer.make_trivial_indexer(ref_aval.shape))
  if len(indexers) != 1:
    raise NotImplementedError("Only single indexer is supported.")
  idx = indexers[0]
  ptr = lowering._compute_pointers_from_indices(ptr, block_info, idx)
  val = lowering._ensure_ir_value(val, value_aval)
  if mask is not None:
    mask = lowering._ensure_ir_value(mask, mask_aval)
  if atomic_type == AtomicOpType.XCHG:
    op = tt_dialect.RMWOp.XCHG
  elif atomic_type == AtomicOpType.ADD:
    if isinstance(val.type, ir.IntegerType):
      op = tt_dialect.RMWOp.ADD
    else:
      op = tt_dialect.RMWOp.FADD
  elif atomic_type == AtomicOpType.MIN:
    if isinstance(val.type, ir.IntegerType):
      op = (
        tt_dialect.RMWOp.MIN
        if jnp.issubdtype(value_aval.dtype, jnp.signedinteger)
        else tt_dialect.RMWOp.UMIN

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Flatten the array (or use a single NDIndexer) so the atomic access uses one indexer, e.g. ref[i * cols + j] instead of ref[i, j] via chained views
2. If indexing a block, index into a temporary non-Ref array first, then do a scalar atomic store
3. File a feature request / check newer JAX versions if multi-indexer atomics were added

Example fix

# before
p.atomic_add(ref[i], j, val)  # chained -> multiple indexers
# after
p.atomic_add(ref, i * ref.shape[1] + j, val)  # single flat indexer

idxers = idx if isinstance(idx, tuple) else (idx,)
assert len(idxers) == 1, 'use a single (flat) indexer for atomics'

When it happens

Trigger: Calling pallas.triton primitives like atomic_add/atomic_or/atomic_xchg with an idx argument that is a tuple of multiple indexers, or applying an atomic op to a Ref produced via chained indexing (e.g. ref[i][j] or a view chain) inside a pallas triton kernel.

Common situations: Porting a GPU kernel that does atomic updates on a 2D location using two separate bracket operations; using jnp-style fancy indexing inside pallas kernels which silently creates multiple indexers.

Related errors

• unsupported atomic operation: {atomic_type}
• Out-of-bounds swap of ({device_id} {local_core_id} {memory_s
• Out-of-bounds masked swap of ({device_id} {local_core_id} {m
• Unsupported range type: {type(r)}.
• Cannot do int indexing on TPU

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