{"record":{"id":"fc68e9128d55558c","repo":"jax-ml/jax","slug":"vjp-arg-name-jaxpr-in-tree-idx-is-ref-typed","errorCode":null,"errorMessage":"{_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gradient must be accumulated into a ref, but no gradient ref was provided. Bind one using the VJP function's `with_refs` method before applying it, as in `f_vjp.with_refs(grad_ref)(ct)`; the gradient will be accumulated into `grad_ref` in-place via addition. Or, to skip computing this argument's gradient, pass `jax.ad.DontWant()` in place of a gradient ref.","messagePattern":"(.+?) is Ref-typed, so its gradient must be accumulated into a ref, but no gradient ref was provided\\. Bind one using the VJP function's `with_refs` method before applying it, as in `f_vjp\\.with_refs\\(grad_ref\\)\\(ct\\)`; the gradient will be accumulated into `grad_ref` in-place via addition\\. Or, to skip computing this argument's gradient, pass `jax\\.ad\\.DontWant\\(\\)` in place of a gradient ref\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1778,"sourceCode":"      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):\n  try:\n    dummy_args = tree_unflatten(in_tree, list(range(in_tree.num_leaves)))\n    path, _ = list(generate_key_paths(dummy_args))[idx]\n    position = f\"args{keystr(path)}\"\n  except Exception:  # unflattening custom pytree nodes can reject dummy leaves\n    position = f\"flat argument index {idx}\"\n  return (f\"the argument at position {position} of the \"","sourceCodeStart":1760,"sourceCodeEnd":1796,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1760-L1796","documentation":"Same constraint as above but for the case where no refs were explicitly provided at all: a Ref-typed argument's gradient must be accumulated into a Ref bound via with_refs, or skipped with jax.ad.DontWant(); by-value return is impossible.","triggerScenarios":"jax.vjp on a function accepting a Ref parameter, then calling the vjp function directly (no with_refs), expecting a gradient tuple back.","commonSituations":"Applying grad/vjp to stateful (Ref-mutating) code; upgrading older functional code to the Ref API without updating the backward pass plumbing.","solutions":["Call f_vjp.with_refs(grad_ref)(ct) binding a ref for the Ref-typed arg","Pass jax.ad.DontWant() in place of that gradient ref","Restructure the function so Refs are not part of the differentiated signature"],"exampleFix":"# before\n_, f_vjp = jax.vjp(loss, x, state_ref)\ng = f_vjp(ct)\n# after\n_, f_vjp = jax.vjp(loss, x, state_ref)\nf_vjp.with_refs(jnp.zeros_like_grad, state_grad_ref)(ct)","handlingStrategy":"validation","validationCode":"ref_typed = [i for i, v in enumerate(primal_avals) if _is_ref_aval(v)]\nassert all(i in bound_ref_indices for i in ref_typed), 'bind gradient refs via with_refs first'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use f_vjp.with_refs(grad_ref)(ct) whenever the function takes Refs","Pass jax.ad.DontWant() to skip unneeded gradients","Centralize with_refs plumbing in one helper"],"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"}