jax-ml/jax · error · ValueError

Invalid sharding for `addupdate`. Ref sharding: {ref_aval.sh

Error message

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

What it means

`addupdate` checks that when the ref's output sharding mesh has any explicit axis, the value's sharding must equal the ref's sharding. This fires under sharded/mesh execution (jax.sharding) when the accumulated value is sharded differently from the reference, since in-place updates can't reconcile distinct explicit-mesh shardings.

Source

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

    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

pp_ref_var = partial(pp.color, intensity=pp.Intensity.NORMAL,
                 foreground=pp.Color.GREEN)


def _pp_transforms(
    context: core.JaxprPpContext,
    transforms: tuple[Transform, ...],

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make the value's sharding match the ref: `x = jax.lax.with_sharding_constraint(x, ref_sharding)` before the update.
Create the ref with the same sharding as the incoming updates.
Replicate one side explicitly (e.g. `jax.lax.all_gather`) if semantics allow.

Example fix

// before
ref[idx] += x  # ref replicated, x sharded across 'data'
// after
x = jax.lax.with_sharding_constraint(x, ref_sharding)
ref[idx] += x

Defensive patterns

Strategy: validation

Validate before calling

x = jax.lax.with_sharding_constraint(x, ref_aval.sharding)
ref[i] += x

Prevention

Create SPMD buffers and their updates with the same NamedSharding.
Assert sharding equality in single-device tests to catch drift before scaling.

When it happens

Trigger: Using `ref[idx] += x` inside pjit/sharded_jit with NamedShardings where ref and value have different sharding specs and the mesh has explicit axes.

Common situations: SPMD code where the buffer was created with one sharding and the update with another (e.g. ref replicated, value sharded); multi-host training loop accumulators; changes in jax sharding propagation between versions.

Related errors

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