{"record":{"id":"4df9bc0c64160faf","repo":"jax-ml/jax","slug":"addupdate-must-be-called-on-ref-types-ref-av","errorCode":null,"errorMessage":"`addupdate` must be called on `Ref` types: {ref_aval}.","messagePattern":"`addupdate` must be called on `Ref` types: (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/state/primitives.py","lineNumber":457,"sourceCode":"          \"Invalid dtype for `swap`. \"\n          f\"Ref dtype: {expected_out_ty.dtype}. \"\n          f\"Value dtype: {val_aval.dtype}. \"\n      )\n    out_aval = expected_out_ty\n  else:\n    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","sourceCodeStart":439,"sourceCodeEnd":475,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/state/primitives.py#L439-L475","documentation":"The `addupdate` primitive (ref[idx] += x lowered form) requires its first argument to be an `AbstractRef`. The abstract evaluator raises this when the first tracer passed is not a Ref aval, i.e. addupdate was bound to a plain array.","triggerScenarios":"Calling `addupdate_p.bind(...)` or internal `ref.addupdate` paths on a non-Ref value; typically from custom lowering code or misuse of internal state APIs rather than the public `Ref` class.","commonSituations":"Writing custom Jaxpr interpreters / primitives that reuse addupdate; passing an array where a state ref is expected after a refactor; version changes in jax.experimental.state internals changing expected argument order.","solutions":["Pass a `jax.experimental.state.AbstractRef`/Ref instance as the first argument.","If you hold an array, wrap or convert it to a ref with the state API (`state.ref(x)`) before addupdate.","Upgrade JAX — internal signatures of state primitives have changed across versions; align with the version you depend on."],"exampleFix":"// before\naddupdate_p.bind(arr, val)\n// after\nref = state.ref(arr)\naddupdate_p.bind(ref, val)","handlingStrategy":"type-guard","validationCode":"from jax.experimental.state import AbstractRef\nassert isinstance(ref_aval, AbstractRef), ref_aval","typeGuard":"def is_ref_aval(a) -> bool:\n    from jax.experimental.state import AbstractRef\n    return isinstance(a, AbstractRef)","tryCatchPattern":null,"preventionTips":["Use the public Ref API rather than binding primitives directly.","Pin your JAX version when relying on internal state primitives."],"tags":["jax","state-primitives","type-error","addupdate"],"backgroundTag":"wrong-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}