jax-ml/jax · error · ValueError

cmp and val must have identical dtypes and shapes

Error message

cmp and val must have identical dtypes and shapes

What it means

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.

Source

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

    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 atomic operation.
  """
  return _atomic_rmw(
      x_ref_or_view, idx, val, mask=mask, atomic_type=AtomicOpType.XOR
  )


atomic_cas_p = jax_core.Primitive("atomic_cas")


@atomic_cas_p.def_effectful_abstract_eval
def _atomic_cas_abstract_eval(ref_aval, cmp_aval, val_aval):
  if cmp_aval.dtype != val_aval.dtype or cmp_aval.shape != val_aval.shape:
    raise ValueError("cmp and val must have identical dtypes and shapes")
  if ref_aval.shape:
    raise ValueError("ref must be scalar.")
  if cmp_aval.shape:
    raise ValueError("cmp must be scalar.")
  if val_aval.shape:
    raise ValueError("val must be scalar.")
  return jax_core.ShapedArray(val_aval.shape, val_aval.dtype), {
      state.WriteEffect(0)
  }


def atomic_cas(ref, cmp, val):
  """Performs an atomic compare-and-swap of the value in the ref with the

  given value.

  Args:
    ref: The ref to operate on.

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
p.atomic_cas(ref, jnp.int32(0), jnp.float32(1.0))
# after
p.atomic_cas(ref, jnp.float32(0.0), jnp.float32(1.0))

Defensive patterns

Strategy: type-guard

Validate before calling

cmp = cmp.astype(val.dtype).reshape(())
val = val.reshape(())

Type guard

def cas_ready(cmp, val):
    return (cmp.dtype == val.dtype and cmp.shape == val.shape == ())

Prevention

Always cast cmp and val to the ref's dtype before atomic_cas

When it happens

Trigger: 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).

Common situations: Implicit-dtype assumptions from NumPy/Python scalars; refactoring kernels where cmp was computed in a different precision than val.

Related errors

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