{"record":{"id":"79fc6a3737c74eee","repo":"jax-ml/jax","slug":"the-gradient-for-vjp-arg-name-jaxpr-in-tree-id","errorCode":null,"errorMessage":"the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which is Ref-typed, can't be returned as a value. In the arguments to the VJP function's `with_refs` method, pass a `Ref` for it, to accumulate the gradient into the ref in-place, or pass `jax.ad.DontWant()` to skip computing this argument's gradient.","messagePattern":"the gradient for (.+?), which is Ref-typed, can't be returned as a value\\. In the arguments to the VJP function's `with_refs` method, pass a `Ref` for it, to accumulate the gradient into the ref in-place, or pass `jax\\.ad\\.DontWant\\(\\)` to skip computing this argument's gradient\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1771,"sourceCode":"  if isinstance(x, ad.GradAccum):\n    return check_accum(v.aval.to_ct_aval(), x)\n  elif _is_ref(x):\n    expected_aval = _ref_aval(v.aval).to_ct_aval()\n    given_aval = _ref_aval(typeof(x))\n    if (not core.typecompat(expected_aval, given_aval) and\n        not _temporary_dtype_exception(given_aval, expected_aval)):\n      raise ValueError(\n          \"unexpected JAX type (e.g. shape/dtype) for gradient ref passed to \"\n          f\"the VJP function's `with_refs` method for \"\n          f\"{_vjp_arg_name(jaxpr, in_tree, idx)}: the given ref has type \"\n          f\"{typeof(x).str_short()}, but accumulating this argument's \"\n          f\"gradient requires a ref of type Ref{{{expected_aval.str_short()}}}\")\n    return ad.RefAccum(expected_aval, x)\n  elif isinstance(x, DontWant):\n    return ad.NullAccum(v.aval.to_ct_aval())\n  elif _is_ref_aval(v.aval):\n    if explicit_refs:\n      raise ValueError(\n          f\"the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which is \"\n          \"Ref-typed, can't be returned as a value. In the arguments to the \"\n          \"VJP function's `with_refs` method, pass a `Ref` for it, to \"\n          \"accumulate the gradient into the ref in-place, or pass \"\n          \"`jax.ad.DontWant()` to skip computing this argument's gradient.\")\n    else:\n      raise ValueError(\n          f\"{_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its \"\n          \"gradient must be accumulated into a ref, but no gradient ref was \"\n          \"provided. Bind one using the VJP function's `with_refs` method \"\n          \"before applying it, as in `f_vjp.with_refs(grad_ref)(ct)`; the \"\n          \"gradient will be accumulated into `grad_ref` in-place via \"\n          \"addition. Or, to skip computing this argument's gradient, pass \"\n          \"`jax.ad.DontWant()` in place of a gradient ref.\")\n  else:\n    return ad.ValAccum(v.aval.to_ct_aval())\n\ndef _vjp_arg_name(jaxpr, in_tree, idx):","sourceCodeStart":1753,"sourceCodeEnd":1789,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1753-L1789","documentation":"If an argument to the original function is Ref-typed, its VJP gradient cannot be returned by value; with_refs demands either a Ref to accumulate into or jax.ad.DontWant(). When explicit refs were provided but one is missing (explicitRefs true) this error is raised.","triggerScenarios":"jax.vjp(f)(x) where f takes a Ref argument and the returned vjp is applied without binding a gradient ref via with_refs, while other args did get refs.","commonSituations":"Functions that mutate state via Refs (e.g. in-place accumulation, new Ref API) combined with vjp; partially converting a training step to with_refs.","solutions":["Bind a gradient ref for the Ref-typed argument via f_vjp.with_refs(...)(ct)","Pass jax.ad.DontWant() for that argument to skip its gradient","Redesign so the differentiated function takes plain arrays and mutates state outside"],"exampleFix":"# before\nout, f_vjp = jax.vjp(step_with_ref, x)\ng = f_vjp(ct)\n# after\nout, f_vjp = jax.vjp(step_with_ref, x, ref)\nf_vjp.with_refs(grad_ref, jax.ad.DontWant())(ct)  # or supply a ref for it","handlingStrategy":"validation","validationCode":"from jax.tree import tree_leaves as tl\nassert all(not _is_ref_aval(jax.core.typeof(v)) or bound for v, bound in zip(primal_leaves, provided)), 'Ref-typed args need a bound gradient ref or DontWant'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Supply a gradient ref (or jax.ad.DontWant()) for every Ref-typed argument","Keep differentiated function signatures Ref-free where possible","Document with_refs requirements next to vjp usage"],"tags":["jax","vjp","refs","autodiff"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}