jax-ml/jax · error · ValueError

unexpected JAX type (e.g. shape/dtype) for gradient ref pass

Error message

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()}}}

What it means

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.

Source

Thrown at jax/_src/api.py:1760

                      isinstance(args_res_[i.idx], NotSaveable)]:
    _vjp_not_saveable_error(jaxpr, in_tree, not_restored)
  residuals = [args_res_[i.idx] if i.primal else opaque_res[i.idx] for i in spec]
  arg_invars = jaxpr.invars[len(spec):]  # skip the residual invars
  maybe_accums = [_vjp_accum(jaxpr, in_tree, explicit_refs, idx, v, x)
                  for idx, (v, x) in enumerate(unsafe_zip(arg_invars, maybe_ct_refs_flat))]
  return Partial(partial(_vjp3_bwd, in_tree, out_tree, out_zeros, jaxpr,
                         out_primal_avals, want_logs), residuals, structured_res,
                 maybe_accums)

def _vjp_accum(jaxpr, in_tree, explicit_refs, idx, v, x):
  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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate each gradient ref with jax.empty_like(param) (matching shape and dtype) inside jax.tree.map over the primal args
2. Fix dtype policy so accumulation refs match parameter dtypes (cast explicitly after if needed)
3. Verify with core.typecompat or by comparing .shape/.dtype of ref vs param before calling

Example fix

# before
grad_ref = jax.make_ref(jnp.zeros(n, dtype=jnp.float16))
f_vjp.with_refs(grad_ref)(ct)
# after
grad_ref = jax.tree.map(lambda p: jax.make_ref(jnp.zeros_like(p)), primals)
f_vjp.with_refs(grad_ref)(ct)

refs = jax.tree.map(lambda p: jax.make_ref(jnp.zeros_like(p)), primal_args)  # guarantees matching aval

def ref_matches(ref, param): return core.typecompat(core.typeof(param).to_ct_aval(), _ref_aval(jax.core.typeof(ref)))

When it happens

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

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

Related errors

• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• But the tree structures differ:
• Accumulator aval mismatch: expected {aval}, got {acc.aval}
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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