jax-ml/jax · error · ValueError

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

Error message

Invalid shape for `addupdate`. Ref shape: {ref_aval.shape}. Expected shape: {expected_out_ty.shape}. Value shape: {val_aval.shape}. Transforms: {transforms}. 

What it means

During abstract evaluation of `addupdate` (the primitive behind `ref[idx] += value`), the value's shape must equal the Ref's element shape after applying the index transforms. This error fires when they differ, printing ref shape, expected (post-transform) shape, value shape, and the transforms for diagnosis.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/broadcast the value to the indexed target shape: `x = jnp.broadcast_to(x, expected_shape)` or `x.reshape(...)`.
2. Print shapes of the ref slice and the value before the update to find the mismatch.
3. If a singleton dim snuck in, squeeze the value: `x.squeeze(-1)`.

Example fix

// before
ref = state.Ref(jnp.zeros((10, 2)))
ref[i] += jnp.ones(2)  # wrong ndim
// after
ref[i] += jnp.ones(2) * 0  # ensure shape (2,)
ref[i] += jnp.ones((2,))
// (key: value shape must equal the (2,) slice shape)

expected = ref.aval.inner_aval.shape  # after transform, e.g. slice shape
upd = jnp.broadcast_to(upd, expected)
ref[i] += upd

When it happens

Trigger: `ref[idx] += x` where x's shape doesn't match the indexed slice, e.g. `ref[i] += jnp.ones((3,))` when `ref[i]` selects a scalar or a (2,) slice; or accumulating with a broadcast-incompatible array inside lax.fori_loop/scan.

Common situations: In-place accumulation in loops where the carry shape drifted; mixing (n,) buffers with (n,1) updates; using fancy indexing that returns a different-length result than assumed; refactors of buffer shapes without updating updates.

Related errors

• Invalid dtype for `addupdate`. Ref dtype: {ref_aval.dtype}.
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out

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