jax-ml/jax · error · ValueError

the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which

Error message

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.

What it means

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.

Source

Thrown at jax/_src/api.py:1771

  if isinstance(x, ad.GradAccum):
    return check_accum(v.aval.to_ct_aval(), x)
  elif _is_ref(x):
    expected_aval = _ref_aval(v.aval).to_ct_aval()
    given_aval = _ref_aval(typeof(x))
    if (not core.typecompat(expected_aval, given_aval) and
        not _temporary_dtype_exception(given_aval, expected_aval)):
      raise ValueError(
          "unexpected JAX type (e.g. shape/dtype) for gradient ref passed to "
          f"the VJP function's `with_refs` method for "
          f"{_vjp_arg_name(jaxpr, in_tree, idx)}: the given ref has type "
          f"{typeof(x).str_short()}, but accumulating this argument's "
          f"gradient requires a ref of type Ref{{{expected_aval.str_short()}}}")
    return ad.RefAccum(expected_aval, x)
  elif isinstance(x, DontWant):
    return ad.NullAccum(v.aval.to_ct_aval())
  elif _is_ref_aval(v.aval):
    if explicit_refs:
      raise ValueError(
          f"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.")
    else:
      raise ValueError(
          f"{_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.")
  else:
    return ad.ValAccum(v.aval.to_ct_aval())

def _vjp_arg_name(jaxpr, in_tree, idx):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Bind a gradient ref for the Ref-typed argument via f_vjp.with_refs(...)(ct)
2. Pass jax.ad.DontWant() for that argument to skip its gradient
3. Redesign so the differentiated function takes plain arrays and mutates state outside

Example fix

# before
out, f_vjp = jax.vjp(step_with_ref, x)
g = f_vjp(ct)
# after
out, f_vjp = jax.vjp(step_with_ref, x, ref)
f_vjp.with_refs(grad_ref, jax.ad.DontWant())(ct)  # or supply a ref for it

from jax.tree import tree_leaves as tl
assert 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'

When it happens

Trigger: 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.

Common situations: 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.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• But the tree structures differ:
• structure of the differentiated function {jaxpr.debug_info.f
• Pure callbacks do not support transpose. Please use `jax.cus

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/79fc6a3737c74eee. Report an issue: GitHub.