jax-ml/jax · error · ValueError

`addupdate` must be called on `Ref` types: {ref_aval}.

Error message

`addupdate` must be called on `Ref` types: {ref_aval}.

What it means

The `addupdate` primitive (ref[idx] += x lowered form) requires its first argument to be an `AbstractRef`. The abstract evaluator raises this when the first tracer passed is not a Ref aval, i.e. addupdate was bound to a plain array.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a `jax.experimental.state.AbstractRef`/Ref instance as the first argument.
2. If you hold an array, wrap or convert it to a ref with the state API (`state.ref(x)`) before addupdate.
3. Upgrade JAX — internal signatures of state primitives have changed across versions; align with the version you depend on.

Example fix

// before
addupdate_p.bind(arr, val)
// after
ref = state.ref(arr)
addupdate_p.bind(ref, val)

from jax.experimental.state import AbstractRef
assert isinstance(ref_aval, AbstractRef), ref_aval

def is_ref_aval(a) -> bool:
    from jax.experimental.state import AbstractRef
    return isinstance(a, AbstractRef)

When it happens

Trigger: Calling `addupdate_p.bind(...)` or internal `ref.addupdate` paths on a non-Ref value; typically from custom lowering code or misuse of internal state APIs rather than the public `Ref` class.

Common situations: Writing custom Jaxpr interpreters / primitives that reuse addupdate; passing an array where a state ref is expected after a refactor; version changes in jax.experimental.state internals changing expected argument order.

Related errors

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.
• Cannot interpret value of type {typ} as an abstract array; i

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