jax-ml/jax · error · ValueError

cmp must be scalar.

Error message

cmp must be scalar.

What it means

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.

Source

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

  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.
    cmp: The expected value to compare against.
    val: The value to swap in.

  Returns:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Squeeze the cmp value: cmp = jnp.squeeze(cmp) or cmp.reshape(())
2. Construct scalars directly with dtype-qualified constants, e.g. jnp.float32(1.0)

Example fix

# before
p.atomic_cas(ref[0], cmp_array, val)  # cmp_array.shape == (1,)
# after
p.atomic_cas(ref[0], jnp.squeeze(cmp_array), val)

cmp = jnp.squeeze(cmp)

def is_scalar(x): return getattr(x, 'shape', ()) == ()

When it happens

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

Common situations: Values coming out of vectorized computations retain shape (1,) instead of (); mixing block-shaped intermediates into scalar atomics.

Related errors

• ref must be scalar.
• val must be scalar.
• Only arguments with shape [..., 1] are supported.
• Only equal-sized splits are supported.
• other cannot be a block if pointer is not a block

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