jax-ml/jax · error · ValueError

structure of the differentiated function {jaxpr.debug_info.f

Error message

structure of the differentiated function {jaxpr.debug_info.func_src_info}.

But the tree structures differ:

What it means

In jax.vjp, the cotangent argument you pass to the returned backward function must have the same pytree structure as the primal function's output. JAX detected a structural mismatch (e.g. output was a dict but cotangent is a tuple) and reports exactly which paths differ.

Source

Thrown at jax/_src/api.py:1966

If we instead call `f_vjp(2.0, 2.0)`, with the values 'splatted out' as
arguments rather than in a tuple, this error can arise.
""".format


def _vjp_ct_tree_error(jaxpr, out_tree, ct_tree):
  msg = f"""unexpected tree structure.

The argument to a VJP function returned by `jax.vjp` must match the pytree
structure of the differentiated function {jaxpr.debug_info.func_src_info}.

But the tree structures differ:
"""
  msg += '\n'.join(f"  * out{keystr(path)} was a {thing1} in the original "
                   f"output, but a {thing2} here, so {explanation}."
                   for path, thing1, thing2, explanation
                   in equality_errors_pytreedef(out_tree, ct_tree))
  raise ValueError(msg)


def _vjp_check_ct_avals(cts, primal_avals):
  # TODO(mattjj): improve this error  by flattening with keys in the first place
  for ct, aval in zip(cts, primal_avals):
    if isinstance(ct, ad.Zero): continue
    ct_aval = typeof(ct)
    ct_aval_expected = aval.to_ct_aval()
    if (not core.typecompat(ct_aval, ct_aval_expected) and
        not _temporary_dtype_exception(ct_aval, ct_aval_expected)):
      raise ValueError(
          "unexpected JAX type (e.g. shape/dtype) for argument to VJP function: "
          f"got {ct_aval.str_short()}, but expected {ct_aval_expected.str_short()} "
          "because the corresponding output of the differentiated function had JAX type "
          f"{aval.str_short()}")


@register_dataclass

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass cotangents matching tree_structure(out) exactly
2. Recompute out, _ = jax.vjp(f, *args) and mirror the structure of out when building cts
3. Use jax.tree_util.tree_map(lambda x: jnp.zeros_like(x), out) to build a correctly-structured cotangent

Example fix

// before
out, vjp_fn = jax.vjp(f, x)
in_cts = vjp_fn(1.0)            # f returns (a, b)
// after
out, vjp_fn = jax.vjp(f, x)
in_cts = vjp_fn((jnp.ones_like(a), jnp.ones_like(b)))

import jax.tree_util as jtu
out, vjp_fn = jax.vjp(f, *args)
cts = jtu.tree_map(jnp.zeros_like, out)   # guaranteed structural match
in_cts = vjp_fn(cts)

try:
    in_cts = vjp_fn(cts)
except ValueError as e:
    if 'tree structures differ' in str(e):
        cts = jax.tree_util.tree_map(jnp.zeros_like, out)
        in_cts = vjp_fn(cts)
    else:
        raise

When it happens

Trigger: Calling vjp_fn(cts) where cts is a tuple but the differentiated function returned a dict, or a scalar where a pair was returned; changing how outputs are packaged between the forward call and the backward call.

Common situations: Function returns multiple values that get auto-packed into a tuple but user passes a single cotangent; using has_aux or auxiliary outputs so the cotangent tree differs from what the user expects; passing jnp.ones(shape) where the output was a pytree of arrays.

Related errors

• primal and tangent arguments to jax.jvp must have the same t
• 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:
• 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/a4682e9a75b7eaf8. Report an issue: GitHub.