jax-ml/jax · error · NotImplementedError

unsupported atomic operation: {atomic_type}

Error message

unsupported atomic operation: {atomic_type}

What it means

The Pallas Triton atomic lowering only maps a fixed set of AtomicOpType values (add, min, max, and/or/xor, xchg) to Triton RMW ops. Any other AtomicOpType reaching the lowering rule raises this NotImplementedError.

Source

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

    else:
      return _expand_atomic_fp_min_max(atomic_type, ptr, val, mask=mask)
  elif atomic_type == AtomicOpType.MAX:
    if isinstance(val.type, ir.IntegerType):
      op = (
        tt_dialect.RMWOp.MAX
        if jnp.issubdtype(value_aval.dtype, jnp.signedinteger)
        else tt_dialect.RMWOp.UMAX
      )
    else:
      return _expand_atomic_fp_min_max(atomic_type, ptr, val, mask=mask)
  elif atomic_type == AtomicOpType.AND:
    op = tt_dialect.RMWOp.AND
  elif atomic_type == AtomicOpType.OR:
    op = tt_dialect.RMWOp.OR
  elif atomic_type == AtomicOpType.XOR:
    op = tt_dialect.RMWOp.XOR
  else:
    raise NotImplementedError(f"unsupported atomic operation: {atomic_type}")
  return lowering._atomic_rmw(op, ptr, val, mask=mask)


def atomic_xchg(x_ref_or_view, idx, val, *, mask: Any | None = None):
  """Atomically exchanges the given value with the value at the given index.

  Args:
    x_ref_or_view: The ref to operate on.
    idx: The indexer to use.
    mask: TO BE DOCUMENTED.

  Returns:
    The value at the given index prior to the aupdate.
  """
  return _atomic_rmw(
      x_ref_or_view, idx, val, mask=mask, atomic_type=AtomicOpType.XCHG
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use only the public helpers (atomic_add, atomic_min, atomic_max, atomic_and, atomic_or, atomic_xor, atomic_xchg, atomic_cas)
2. Upgrade (or align) JAX to a version where the Triton backend supports the op you need
3. Replace the unsupported op with a supported composition (e.g. compare-and-swap loop using atomic_cas)

Example fix

# before
custom_atomic(AtomicOpType.SOMETHING_NEW, ref, idx, val)
# after
p.atomic_add(ref, idx, val)  # supported op

from jax._src.pallas.triton.primitives import AtomicOpType
SUPPORTED = {AtomicOpType.ADD, AtomicOpType.MIN, AtomicOpType.MAX, AtomicOpType.AND, AtomicOpType.OR, AtomicOpType.XOR, AtomicOpType.XCHG}
assert op in SUPPORTED

When it happens

Trigger: Constructing or invoking a pallas atomic primitive with an AtomicOpType not covered by the if/elif chain (e.g. a newly added or internal op type) via pallas.triton.primitives internals rather than the public helpers.

Common situations: Using internal pallas APIs directly; version mismatch where a custom/new AtomicOpType exists on one side but the Triton lowering was not updated.

Related errors

• Only single indexer is supported.
• Out-of-bounds swap of ({device_id} {local_core_id} {memory_s
• Out-of-bounds masked swap of ({device_id} {local_core_id} {m
• Only positional arguments are supported by debug_print on Pa
• unsupported dtypes: {x_aval.dtype} and {y_aval.dtype}

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