jax-ml/jax · error · ValueError

But the tree structures differ:

Error message

But the tree structures differ:

What it means

When applying a VJP function bound via with_refs, the pytree structure of the refs must match the structure of the original primal arguments. This error enumerates each path where a node was, e.g., a dict in the primal arguments but a tuple in the with_refs arguments, and raises with the itemized diff.

Source

Thrown at jax/_src/api.py:1819

`with_refs` method.

The arguments to `with_refs` must match the pytree structure of the primal
arguments of the differentiated function {jaxpr.debug_info.func_src_info},
with one entry for each array argument. Each entry must be a `Ref` (to
accumulate this argument's gradient into the ref in-place), a
`jax.ad.GradValue()` (to have this argument's gradient returned as a value,
the default behavior), or a `jax.ad.DontWant()` (to skip computing this
argument's gradient). Note that `None` is an empty pytree, so it can't be
used as a placeholder entry.

But the tree structures differ:
"""
  msg += '\n'.join(f"  * args{keystr(path)} was a {thing1} in the primal "
                   f"arguments, but a {thing2} in the `with_refs` arguments, "
                   f"so {explanation}."
                   for path, thing1, thing2, explanation
                   in equality_errors_pytreedef(in_tree, refs_tree))
  raise ValueError(msg)

def _vjp_not_saveable_error(jaxpr, in_tree, idxs):
  msg = """the VJP function was applied before restoring its not-saveable residuals.

Because `saveable_args` was passed to `jax.vjp`, some argument values that
would have been saved for the backward pass were instead replaced with
`NotSaveable()` sentinels. Before the VJP function can be applied, these
values must be restored, e.g. by assigning to the VJP function's `args_res`
attribute. The values not yet restored correspond to:
"""
  msg += '\n'.join(f"  * {_vjp_arg_name(jaxpr, in_tree, idx)};" for idx in idxs)
  raise ValueError(msg)

def check_accum(aval, acc):
  if not core.typecompat(acc.aval, aval):
    raise ValueError(f"Accumulator aval mismatch: expected {aval}, got {acc.aval}")
  return acc

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Build the refs from the primal tree directly: jax.tree.map(lambda p: jax.make_ref(jnp.zeros_like(p)), primals)
2. Fix container types/keys to mirror the primal arguments exactly
3. Log jax.tree.structure(primals) vs jax.tree.structure(refs) to diff structures

Example fix

# before
primals = {'w': w, 'b': b}
f_vjp.with_refs((g_w, g_b))(ct)
# after
primals = {'w': w, 'b': b}
refs = jax.tree.map(lambda p: jax.make_ref(jnp.zeros_like(p)), primals)
f_vjp.with_refs(refs)(ct)

assert jax.tree.structure(refs_tree) == jax.tree.structure(primal_tree), 'with_refs tree must mirror primal argument tree'

def refs_tree_ok(primals, refs): return jax.tree.structure(primals) == jax.tree.structure(refs)

When it happens

Trigger: f_vjp.with_refs(grads_tuple)(ct) when primals were passed as a dict; mismatched leaf counts, container types, or dict keys anywhere in the tree.

Common situations: Gradient buffers built by a different utility than the parameters (e.g. optax state tree vs param tree); refactoring parameter dataclass fields; mixing namedtuple vs tuple.

Related errors

• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• structure of the differentiated function {jaxpr.debug_info.f
• numpy masked arrays are not supported as direct inputs to JA

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