{"record":{"id":"f8622bdf85b7cba6","repo":"jax-ml/jax","slug":"cmp-must-be-scalar","errorCode":null,"errorMessage":"cmp must be scalar.","messagePattern":"cmp must be scalar\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/primitives.py","lineNumber":625,"sourceCode":"  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.\n    cmp: The expected value to compare against.\n    val: The value to swap in.\n\n  Returns:","sourceCodeStart":607,"sourceCodeEnd":643,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/primitives.py#L607-L643","documentation":"The comparison operand of pallas.triton atomic_cas must be a scalar (empty shape). A non-scalar cmp (e.g. a (1,)-shaped array) fails validation in the abstract eval.","triggerScenarios":"Passing cmp as an array, e.g. p.atomic_cas(ref[0], jnp.zeros(1), 0.0), or a value produced by an operation that left a trailing dimension.","commonSituations":"Values coming out of vectorized computations retain shape (1,) instead of (); mixing block-shaped intermediates into scalar atomics.","solutions":["Squeeze the cmp value: cmp = jnp.squeeze(cmp) or cmp.reshape(())","Construct scalars directly with dtype-qualified constants, e.g. jnp.float32(1.0)"],"exampleFix":"# before\np.atomic_cas(ref[0], cmp_array, val)  # cmp_array.shape == (1,)\n# after\np.atomic_cas(ref[0], jnp.squeeze(cmp_array), val)","handlingStrategy":"type-guard","validationCode":"cmp = jnp.squeeze(cmp)","typeGuard":"def is_scalar(x): return getattr(x, 'shape', ()) == ()","tryCatchPattern":null,"preventionTips":["Build scalars with jnp.float32(0.0)-style constructors"],"tags":["jax","pallas","triton","atomic-cas","shape"],"backgroundTag":"invalid-argument-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}