{"record":{"id":"7ef5b8f6ee293afc","repo":"jax-ml/jax","slug":"cmp-and-val-must-have-identical-dtypes-and-shapes","errorCode":null,"errorMessage":"cmp and val must have identical dtypes and shapes","messagePattern":"cmp and val must have identical dtypes and shapes","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/primitives.py","lineNumber":621,"sourceCode":"    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 atomic operation.\n  \"\"\"\n  return _atomic_rmw(\n      x_ref_or_view, idx, val, mask=mask, atomic_type=AtomicOpType.XOR\n  )\n\n\natomic_cas_p = jax_core.Primitive(\"atomic_cas\")\n\n\n@atomic_cas_p.def_effectful_abstract_eval\ndef _atomic_cas_abstract_eval(ref_aval, cmp_aval, val_aval):\n  if cmp_aval.dtype != val_aval.dtype or cmp_aval.shape != val_aval.shape:\n    raise ValueError(\"cmp and val must have identical dtypes and shapes\")\n  if ref_aval.shape:\n    raise ValueError(\"ref must be scalar.\")\n  if cmp_aval.shape:\n    raise ValueError(\"cmp must be scalar.\")\n  if val_aval.shape:\n    raise ValueError(\"val must be scalar.\")\n  return jax_core.ShapedArray(val_aval.shape, val_aval.dtype), {\n      state.WriteEffect(0)\n  }\n\n\ndef atomic_cas(ref, cmp, val):\n  \"\"\"Performs an atomic compare-and-swap of the value in the ref with the\n\n  given value.\n\n  Args:\n    ref: The ref to operate on.","sourceCodeStart":603,"sourceCodeEnd":639,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/primitives.py#L603-L639","documentation":"atomic_cas in pallas.triton requires the comparison value and the replacement value to have both the same dtype and the same shape; they are validated in the abstract eval before lowering. Mismatched dtypes (e.g. cmp as int32, val as float32) or differing shapes trigger this ValueError.","triggerScenarios":"Calling p.torch.atomic_cas(ref, cmp, val) (or atomic_cas) where cmp and val differ in dtype or shape, e.g. cmp=jnp.float32(1.0) and val=jnp.int32(2).","commonSituations":"Implicit-dtype assumptions from NumPy/Python scalars; refactoring kernels where cmp was computed in a different precision than val.","solutions":["Cast both to the ref's dtype: cmp = cmp.astype(ref.dtype), val = val.astype(ref.dtype)","Ensure both are scalar (shape ()) of the same dtype before calling atomic_cas"],"exampleFix":"# before\np.atomic_cas(ref, jnp.int32(0), jnp.float32(1.0))\n# after\np.atomic_cas(ref, jnp.float32(0.0), jnp.float32(1.0))","handlingStrategy":"type-guard","validationCode":"cmp = cmp.astype(val.dtype).reshape(())\nval = val.reshape(())","typeGuard":"def cas_ready(cmp, val):\n    return (cmp.dtype == val.dtype and cmp.shape == val.shape == ())","tryCatchPattern":null,"preventionTips":["Always cast cmp and val to the ref's dtype before atomic_cas"],"tags":["jax","pallas","triton","dtype","atomic-cas"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}