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.

Solutions

1. Cast both to the ref's dtype: cmp = cmp.astype(ref.dtype), val = val.astype(ref.dtype)
2. 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))

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

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

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

• a and b must have the same element type, but got
• cannot cast from
• cannot cast to
• cmp must be scalar.
• erf_inv_lowering_helper not implemented for

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

Appendix: 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)