{"record":{"id":"c156fa51ae25eb91","repo":"jax-ml/jax","slug":"invalid-sharding-for-addupdate-ref-sharding-r","errorCode":null,"errorMessage":"Invalid sharding for `addupdate`. Ref sharding: {ref_aval.sharding}. Value sharding: {val_aval.sharding}. ","messagePattern":"Invalid sharding for `addupdate`\\. Ref sharding: (.+?)\\. Value sharding: (.+?)\\. ","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/primitives.py","lineNumber":478,"sourceCode":"    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:\n    # Check that the transforms are valid\n    if transforms:\n      raise ValueError(\"Cannot index non-shaped array with nontrivial indices.\")\n  return [], {AccumEffect(0)}\naddupdate_p.def_effectful_abstract_eval(_addupdate_abstract_eval)\n\n## Pretty printing for `get` and `swap` in jaxprs\n\npp_ref_var = partial(pp.color, intensity=pp.Intensity.NORMAL,\n                 foreground=pp.Color.GREEN)\n\n\ndef _pp_transforms(\n    context: core.JaxprPpContext,\n    transforms: tuple[Transform, ...],","sourceCodeStart":460,"sourceCodeEnd":496,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/primitives.py#L460-L496","documentation":"`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.","triggerScenarios":"Using `ref[idx] += x` inside pjit/sharded_jit with NamedShardings where ref and value have different sharding specs and the mesh has explicit axes.","commonSituations":"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.","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."],"exampleFix":"// before\nref[idx] += x  # ref replicated, x sharded across 'data'\n// after\nx = jax.lax.with_sharding_constraint(x, ref_sharding)\nref[idx] += x","handlingStrategy":"validation","validationCode":"x = jax.lax.with_sharding_constraint(x, ref_aval.sharding)\nref[i] += x","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Create SPMD buffers and their updates with the same NamedSharding.","Assert sharding equality in single-device tests to catch drift before scaling."],"tags":["jax","sharding","spmd","state-primitives"],"backgroundTag":"sharding-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}