{"record":{"id":"7bd44484fdf1a82d","repo":"jax-ml/jax","slug":"unexpected-jax-type-e-g-shape-dtype-for-gradien","errorCode":null,"errorMessage":"unexpected JAX type (e.g. shape/dtype) for gradient ref passed to the VJP function's `with_refs` method for {_vjp_arg_name(jaxpr, in_tree, idx)}: the given ref has type {typeof(x).str_short()}, but accumulating this argument's gradient requires a ref of type Ref{{{expected_aval.str_short()}}}","messagePattern":"unexpected JAX type \\(e\\.g\\. shape/dtype\\) for gradient ref passed to the VJP function's `with_refs` method for (.+?): the given ref has type (.+?), but accumulating this argument's gradient requires a ref of type Ref(.+?)\\}\\}","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1760,"sourceCode":"                      isinstance(args_res_[i.idx], NotSaveable)]:\n    _vjp_not_saveable_error(jaxpr, in_tree, not_restored)\n  residuals = [args_res_[i.idx] if i.primal else opaque_res[i.idx] for i in spec]\n  arg_invars = jaxpr.invars[len(spec):]  # skip the residual invars\n  maybe_accums = [_vjp_accum(jaxpr, in_tree, explicit_refs, idx, v, x)\n                  for idx, (v, x) in enumerate(unsafe_zip(arg_invars, maybe_ct_refs_flat))]\n  return Partial(partial(_vjp3_bwd, in_tree, out_tree, out_zeros, jaxpr,\n                         out_primal_avals, want_logs), residuals, structured_res,\n                 maybe_accums)\n\ndef _vjp_accum(jaxpr, in_tree, explicit_refs, idx, v, x):\n  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(","sourceCodeStart":1742,"sourceCodeEnd":1778,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1742-L1778","documentation":"When calling a VJP function's with_refs method, gradient refs supplied for arguments must have the JAX type (shape/dtype) required for accumulating that argument's gradient (its ct_aval, wrapped in a Ref). This error fires when the passed ref's aval is incompatible.","triggerScenarios":"f_vjp.with_refs(grad_ref)(ct) where grad_ref was created with a different shape or dtype than the corresponding input to jax.vjp, e.g. f32 input but f16 ref, or wrong-shaped buffer.","commonSituations":"Pre-allocating gradient accumulation buffers with a global dtype policy (e.g. f16 for memory) that differs from parameter dtype; reusing one buffer shape for many parameters; refactors changing parameter shapes.","solutions":["Allocate each gradient ref with jax.empty_like(param) (matching shape and dtype) inside jax.tree.map over the primal args","Fix dtype policy so accumulation refs match parameter dtypes (cast explicitly after if needed)","Verify with core.typecompat or by comparing .shape/.dtype of ref vs param before calling"],"exampleFix":"# before\ngrad_ref = jax.make_ref(jnp.zeros(n, dtype=jnp.float16))\nf_vjp.with_refs(grad_ref)(ct)\n# after\ngrad_ref = jax.tree.map(lambda p: jax.make_ref(jnp.zeros_like(p)), primals)\nf_vjp.with_refs(grad_ref)(ct)","handlingStrategy":"type-guard","validationCode":"refs = jax.tree.map(lambda p: jax.make_ref(jnp.zeros_like(p)), primal_args)  # guarantees matching aval","typeGuard":"def ref_matches(ref, param): return core.typecompat(core.typeof(param).to_ct_aval(), _ref_aval(jax.core.typeof(ref)))","tryCatchPattern":null,"preventionTips":["Always allocate gradient refs from the primal tree with zeros_like","Keep accumulation buffer dtype equal to parameter dtype","Validate ref aval vs param before with_refs calls"],"tags":["jax","vjp","refs","gradient-accumulation","dtype"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}