{"record":{"id":"5de2f15da081bcd1","repo":"jax-ml/jax","slug":"invalid-dtype-for-addupdate-ref-dtype-ref-ava","errorCode":null,"errorMessage":"Invalid dtype for `addupdate`. Ref dtype: {ref_aval.dtype}. Value shape: {val_aval.dtype}. ","messagePattern":"Invalid dtype for `addupdate`\\. Ref dtype: (.+?)\\. Value shape: (.+?)\\. ","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/primitives.py","lineNumber":471,"sourceCode":"                             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:\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","sourceCodeStart":453,"sourceCodeEnd":489,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/primitives.py#L453-L489","documentation":"The `addupdate` primitive requires the accumulated value's dtype to exactly match the Ref's element dtype. This check runs after the shape check and reports both dtypes (note the message's 'Value shape' label is a copy-paste artifact — it prints the value dtype).","triggerScenarios":"`ref[idx] += x` where x has a different dtype than the ref, e.g. adding float32 gradients into a float16/int ref, or adding a Python int to a float ref under changed x64 settings.","commonSituations":"Mixed-precision training (f32 grads, f16/bf16 buffers); accumulators created with `jnp.zeros(..., dtype=int)` then updated with floats; enabling/disabling x64 changing scalar literal dtypes.","solutions":["Cast the update: `ref[idx] += x.astype(ref.dtype)`.","Create the accumulator ref with the update dtype: `state.Ref(jnp.zeros(shape, dtype=upd.dtype))`.","For mixed precision, use explicit loss scaling + cast rather than relying on implicit promotion."],"exampleFix":"// before\nref = state.Ref(jnp.zeros((n,), jnp.float16))\nref[i] += grad  # grad is float32\n// after\nref[i] += grad.astype(jnp.float16)","handlingStrategy":"validation","validationCode":"upd = jnp.asarray(upd).astype(ref.aval.inner_aval.dtype)\nref[i] += upd","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Define one dtype for buffers and updates in a config dataclass.","Cast at every write boundary in mixed-precision code."],"tags":["jax","dtype-mismatch","state-primitives","in-place-update"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}