jax-ml/jax · error · ValueError

{_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr

Error message

{_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.

What it means

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.

Source

Thrown at jax/_src/api.py:1778

      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):
  try:
    dummy_args = tree_unflatten(in_tree, list(range(in_tree.num_leaves)))
    path, _ = list(generate_key_paths(dummy_args))[idx]
    position = f"args{keystr(path)}"
  except Exception:  # unflattening custom pytree nodes can reject dummy leaves
    position = f"flat argument index {idx}"
  return (f"the argument at position {position} of the "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call f_vjp.with_refs(grad_ref)(ct) binding a ref for the Ref-typed arg
2. Pass jax.ad.DontWant() in place of that gradient ref
3. Restructure the function so Refs are not part of the differentiated signature

Example fix

# before
_, f_vjp = jax.vjp(loss, x, state_ref)
g = f_vjp(ct)
# after
_, f_vjp = jax.vjp(loss, x, state_ref)
f_vjp.with_refs(jnp.zeros_like_grad, state_grad_ref)(ct)

ref_typed = [i for i, v in enumerate(primal_avals) if _is_ref_aval(v)]
assert all(i in bound_ref_indices for i in ref_typed), 'bind gradient refs via with_refs first'

When it happens

Trigger: jax.vjp on a function accepting a Ref parameter, then calling the vjp function directly (no with_refs), expecting a gradient tuple back.

Common situations: Applying grad/vjp to stateful (Ref-mutating) code; upgrading older functional code to the Ref API without updating the backward pass plumbing.

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

• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• 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/fc68e9128d55558c. Report an issue: GitHub.