{"record":{"id":"da61de635092e726","repo":"jax-ml/jax","slug":"unsupported-atomic-operation-atomic-type","errorCode":null,"errorMessage":"unsupported atomic operation: {atomic_type}","messagePattern":"unsupported atomic operation: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/primitives.py","lineNumber":499,"sourceCode":"    else:\n      return _expand_atomic_fp_min_max(atomic_type, ptr, val, mask=mask)\n  elif atomic_type == AtomicOpType.MAX:\n    if isinstance(val.type, ir.IntegerType):\n      op = (\n        tt_dialect.RMWOp.MAX\n        if jnp.issubdtype(value_aval.dtype, jnp.signedinteger)\n        else tt_dialect.RMWOp.UMAX\n      )\n    else:\n      return _expand_atomic_fp_min_max(atomic_type, ptr, val, mask=mask)\n  elif atomic_type == AtomicOpType.AND:\n    op = tt_dialect.RMWOp.AND\n  elif atomic_type == AtomicOpType.OR:\n    op = tt_dialect.RMWOp.OR\n  elif atomic_type == AtomicOpType.XOR:\n    op = tt_dialect.RMWOp.XOR\n  else:\n    raise NotImplementedError(f\"unsupported atomic operation: {atomic_type}\")\n  return lowering._atomic_rmw(op, ptr, val, mask=mask)\n\n\ndef atomic_xchg(x_ref_or_view, idx, val, *, mask: Any | None = None):\n  \"\"\"Atomically exchanges the given value with the value at the given index.\n\n  Args:\n    x_ref_or_view: The ref to operate on.\n    idx: The indexer to use.\n    mask: TO BE DOCUMENTED.\n\n  Returns:\n    The value at the given index prior to the aupdate.\n  \"\"\"\n  return _atomic_rmw(\n      x_ref_or_view, idx, val, mask=mask, atomic_type=AtomicOpType.XCHG\n  )\n","sourceCodeStart":481,"sourceCodeEnd":517,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/primitives.py#L481-L517","documentation":"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.","triggerScenarios":"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.","commonSituations":"Using internal pallas APIs directly; version mismatch where a custom/new AtomicOpType exists on one side but the Triton lowering was not updated.","solutions":["Use only the public helpers (atomic_add, atomic_min, atomic_max, atomic_and, atomic_or, atomic_xor, atomic_xchg, atomic_cas)","Upgrade (or align) JAX to a version where the Triton backend supports the op you need","Replace the unsupported op with a supported composition (e.g. compare-and-swap loop using atomic_cas)"],"exampleFix":"# before\ncustom_atomic(AtomicOpType.SOMETHING_NEW, ref, idx, val)\n# after\np.atomic_add(ref, idx, val)  # supported op","handlingStrategy":"validation","validationCode":"from jax._src.pallas.triton.primitives import AtomicOpType\nSUPPORTED = {AtomicOpType.ADD, AtomicOpType.MIN, AtomicOpType.MAX, AtomicOpType.AND, AtomicOpType.OR, AtomicOpType.XOR, AtomicOpType.XCHG}\nassert op in SUPPORTED","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Stick to public atomic helper functions","Pin JAX version when using internal pallas APIs"],"tags":["jax","pallas","triton","atomics"],"backgroundTag":"unsupported-operation-lowering","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}