jax-ml/jax · error · ValueError

Invalid dtype for `addupdate`. Ref dtype: {ref_aval.dtype}.

Error message

Invalid dtype for `addupdate`. Ref dtype: {ref_aval.dtype}. Value shape: {val_aval.dtype}.

What it means

The `addupdate` primitive requires the accumulated value's dtype to exactly match the Ref's element dtype. This check runs after the shape check and reports both dtypes (note the message's 'Value shape' label is a copy-paste artifact — it prints the value dtype).

Source

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

                             val_aval: core.AbstractValue,
                             *args: Any, tree):
  transforms = tree_util.tree_unflatten(tree, args)
  if not isinstance(ref_aval, AbstractRef):
    raise ValueError(f"`addupdate` must be called on `Ref` types: {ref_aval}.")
  if isinstance(ref_aval.inner_aval, core.ShapedArray):
    expected_out_ty = transform_type(transforms, ref_aval.inner_aval)
    assert isinstance(val_aval, core.ShapedArray)
    assert isinstance(expected_out_ty, core.ShapedArray)
    if expected_out_ty.shape != val_aval.shape:
      raise ValueError(
          "Invalid shape for `addupdate`. "
          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 `addupdate`. "
                       f"Ref dtype: {ref_aval.dtype}. "
                       f"Value shape: {val_aval.dtype}. ")
    out_sharding = expected_out_ty.sharding
    if ((out_sharding.mesh._any_axis_explicit or
         val_aval.sharding.mesh._any_axis_explicit) and
        out_sharding != val_aval.sharding):
      raise ValueError("Invalid sharding for `addupdate`. "
                       f"Ref sharding: {ref_aval.sharding}. "
                       f"Value sharding: {val_aval.sharding}. ")
  else:
    # Check that the transforms are valid
    if transforms:
      raise ValueError("Cannot index non-shaped array with nontrivial indices.")
  return [], {AccumEffect(0)}
addupdate_p.def_effectful_abstract_eval(_addupdate_abstract_eval)

## Pretty printing for `get` and `swap` in jaxprs

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Cast the update: `ref[idx] += x.astype(ref.dtype)`.
Create the accumulator ref with the update dtype: `state.Ref(jnp.zeros(shape, dtype=upd.dtype))`.
For mixed precision, use explicit loss scaling + cast rather than relying on implicit promotion.

Example fix

// before
ref = state.Ref(jnp.zeros((n,), jnp.float16))
ref[i] += grad  # grad is float32
// after
ref[i] += grad.astype(jnp.float16)

Defensive patterns

Strategy: validation

Validate before calling

upd = jnp.asarray(upd).astype(ref.aval.inner_aval.dtype)
ref[i] += upd

Prevention

Define one dtype for buffers and updates in a config dataclass.
Cast at every write boundary in mixed-precision code.

When it happens

Trigger: `ref[idx] += x` where x has a different dtype than the ref, e.g. adding float32 gradients into a float16/int ref, or adding a Python int to a float ref under changed x64 settings.

Common situations: Mixed-precision training (f32 grads, f16/bf16 buffers); accumulators created with `jnp.zeros(..., dtype=int)` then updated with floats; enabling/disabling x64 changing scalar literal dtypes.

Related errors

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