jax-ml/jax · error · ValueError

Invalid dtype for `swap`. Ref dtype: {expected_out_ty.dtype}

Error message

Invalid dtype for `swap`. Ref dtype: {expected_out_ty.dtype}. Value dtype: {val_aval.dtype}. 

What it means

Raised during abstract evaluation of the `swap` primitive when the value being swapped into a Ref has a dtype different from the Ref's element dtype (after applying index transforms). JAX's state primitives require exact dtype match between the written value and the reference; no implicit casting is performed. The error names both dtypes so the mismatch is immediately visible.

Source

Thrown at jax/_src/state/primitives.py:438

        ref_aval, val_aval, *args, tree=tree)
  out_aval: core.AbstractValue
  if not isinstance(ref_aval, AbstractRef):
    raise ValueError(f"`swap` must be called on `Ref` types: {ref_aval}.")
  if isinstance(val_aval, AbstractRef):
    raise ValueError("Cannot store a Ref into another Ref. "
                     "Did you forget to load from it using `[...]`?")
  if isinstance(ref_aval.inner_aval, core.ShapedArray):
    assert isinstance(val_aval, core.ShapedArray)
    expected_out_ty = transform_type(transforms, ref_aval.inner_aval)
    assert isinstance(expected_out_ty, core.ShapedArray)
    if expected_out_ty.shape != val_aval.shape:
      raise ValueError("Invalid shape for `swap`. "
                       f"Ref shape: {ref_aval.shape}. "
                       f"Expected shape: {expected_out_ty.shape}. "
                       f"Value shape: {val_aval.shape}. "
                       f"Transforms: {transforms}. ")
    if expected_out_ty.dtype != val_aval.dtype:
      raise ValueError(
          "Invalid dtype for `swap`. "
          f"Ref dtype: {expected_out_ty.dtype}. "
          f"Value dtype: {val_aval.dtype}. "
      )
    out_aval = expected_out_ty
  else:
    if transforms:
      raise ValueError("Cannot index non-shaped array with nontrivial indices.")
    out_aval = ref_aval.inner_aval
  return (out_aval, {WriteEffect(0)})
swap_p.def_effectful_abstract_eval(_swap_abstract_eval)


def _addupdate_abstract_eval(ref_aval: AbstractRef,
                             val_aval: core.AbstractValue,
                             *args: Any, tree):
  transforms = tree_util.tree_unflatten(tree, args)
  if not isinstance(ref_aval, AbstractRef):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the value to the Ref's dtype before swapping: `x.astype(ref.dtype)` (or `ref_aval.inner_aval.dtype`).
2. Initialize the Ref with the same dtype as the values you will write (e.g. `jnp.zeros(shape, dtype=x.dtype)`).
3. Check for accidental weak-typed scalars: wrap Python scalars in `jnp.asarray(x, dtype=...)`.
4. Under `jax.experimental.enable_x64`, verify both sides weren't created under different x64 settings.

Example fix

// before
ref = state.Ref(jnp.zeros((n,), jnp.int32))
old = ref.swap(0, jnp.float32(1.5))
// after
ref = state.Ref(jnp.zeros((n,), jnp.int32))
old = ref.swap(0, jnp.asarray(1, dtype=jnp.int32))

val = jnp.asarray(val)
assert val.dtype == ref.aval.inner_aval.dtype, (val.dtype, ref.aval.inner_aval.dtype)
old = ref.swap(idx, val)

When it happens

Trigger: Calling `ref.swap(...)` (or `swap_p.bind`, or lax.fori_loop bodies / scan carry updates that compile to swap) with a value whose dtype differs from the ref, e.g. writing a float32 into an int32 Ref, or a weakly-typed Python scalar array promoted differently.

Common situations: Creating a Ref via `jax.experimental.state.ref` or getters like `Ref(np.zeros(..., dtype=np.int32))` then swapping a float loss/gradient; mixing f32/f64 under enabled x64; converting numpy arrays whose default dtype differs from the initialized buffer.

Related errors

• Invalid dtype for `addupdate`. Ref dtype: {ref_aval.dtype}.
• Incorrect output dtype for return value #{i}: Expected: {out
• {} function carry input and carry output must have equal typ
• lower and upper arguments to fori_loop must have equal types
• the first argument to pow must have an inexact dtype (float

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