{"record":{"id":"cddec3bb188f398e","repo":"jax-ml/jax","slug":"invalid-shape-for-addupdate-ref-shape-ref-ava","errorCode":null,"errorMessage":"Invalid shape for `addupdate`. Ref shape: {ref_aval.shape}. Expected shape: {expected_out_ty.shape}. Value shape: {val_aval.shape}. Transforms: {transforms}. ","messagePattern":"Invalid shape for `addupdate`\\. Ref shape: (.+?)\\. Expected shape: (.+?)\\. Value shape: (.+?)\\. Transforms: (.+?)\\. ","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/primitives.py","lineNumber":463,"sourceCode":"    if transforms:\n      raise ValueError(\"Cannot index non-shaped array with nontrivial indices.\")\n    out_aval = ref_aval.inner_aval\n  return (out_aval, {WriteEffect(0)})\nswap_p.def_effectful_abstract_eval(_swap_abstract_eval)\n\n\ndef _addupdate_abstract_eval(ref_aval: AbstractRef,\n                             val_aval: core.AbstractValue,\n                             *args: Any, tree):\n  transforms = tree_util.tree_unflatten(tree, args)\n  if not isinstance(ref_aval, AbstractRef):\n    raise ValueError(f\"`addupdate` must be called on `Ref` types: {ref_aval}.\")\n  if isinstance(ref_aval.inner_aval, core.ShapedArray):\n    expected_out_ty = transform_type(transforms, ref_aval.inner_aval)\n    assert isinstance(val_aval, core.ShapedArray)\n    assert isinstance(expected_out_ty, core.ShapedArray)\n    if expected_out_ty.shape != val_aval.shape:\n      raise ValueError(\n          \"Invalid shape for `addupdate`. \"\n          f\"Ref shape: {ref_aval.shape}. \"\n          f\"Expected shape: {expected_out_ty.shape}. \"\n          f\"Value shape: {val_aval.shape}. \"\n          f\"Transforms: {transforms}. \"\n      )\n    if expected_out_ty.dtype != val_aval.dtype:\n      raise ValueError(\"Invalid dtype for `addupdate`. \"\n                       f\"Ref dtype: {ref_aval.dtype}. \"\n                       f\"Value shape: {val_aval.dtype}. \")\n    out_sharding = expected_out_ty.sharding\n    if ((out_sharding.mesh._any_axis_explicit or\n         val_aval.sharding.mesh._any_axis_explicit) and\n        out_sharding != val_aval.sharding):\n      raise ValueError(\"Invalid sharding for `addupdate`. \"\n                       f\"Ref sharding: {ref_aval.sharding}. \"\n                       f\"Value sharding: {val_aval.sharding}. \")\n  else:","sourceCodeStart":445,"sourceCodeEnd":481,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/primitives.py#L445-L481","documentation":"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.","triggerScenarios":"`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.","commonSituations":"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.","solutions":["Reshape/broadcast the value to the indexed target shape: `x = jnp.broadcast_to(x, expected_shape)` or `x.reshape(...)`.","Print shapes of the ref slice and the value before the update to find the mismatch.","If a singleton dim snuck in, squeeze the value: `x.squeeze(-1)`."],"exampleFix":"// before\nref = state.Ref(jnp.zeros((10, 2)))\nref[i] += jnp.ones(2)  # wrong ndim\n// after\nref[i] += jnp.ones(2) * 0  # ensure shape (2,)\nref[i] += jnp.ones((2,))\n// (key: value shape must equal the (2,) slice shape)","handlingStrategy":"validation","validationCode":"expected = ref.aval.inner_aval.shape  # after transform, e.g. slice shape\nupd = jnp.broadcast_to(upd, expected)\nref[i] += upd","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Broadcast updates to the target slice shape before adding.","Watch for (n,) vs (n,1) drift when refactoring buffers.","Unit-test accumulator shapes under the loop body."],"tags":["jax","shape-mismatch","state-primitives","in-place-update"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}