jax-ml/jax · error · TypeError

cotangent tree does not match function output, expected {out

Error message

cotangent tree does not match function output, expected {out_tree()} but got {out_tree2}

What it means

The transposed function returned by jax.linear_transpose must be called with a cotangent whose pytree structure equals the structure of the original function's output. This fires when the passed cotangent tree (tuples/dicts/leaf counts) differs.

Source

Thrown at jax/_src/api.py:2121

  in_pvals = map(pe.PartialVal.unknown, in_avals)
  jaxpr, out_pvals, const = pe.trace_to_jaxpr_nounits(flat_fun, in_pvals,
                                                      instantiate=True)
  jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True)
  out_avals, _ = unzip2(out_pvals)
  out_dtypes = [a.dtype for a in out_avals if not a.is_high]
  if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes)
          or all(dtypes.issubdtype(d, np.integer)
                 for d in in_dtypes + out_dtypes)):
    raise TypeError("linear_transpose only supports [float or complex] -> "
                    "[float or complex], and integer -> integer functions, "
                    f"but got {in_dtypes} -> {out_dtypes}.")

  @api_boundary
  def transposed_fun(const, out_cotangent):
    out_cts, out_tree2 = tree_flatten(out_cotangent)
    if out_tree() != out_tree2:
      raise TypeError("cotangent tree does not match function output, "
                      f"expected {out_tree()} but got {out_tree2}")
    if not all(map(core.typecheck, out_avals, out_cts)):
      raise TypeError("cotangent type does not match function output, "
                      f"expected {out_avals} but got {out_cts}")
    dummies = [ad.UndefinedPrimal(a.to_ct_aval()) for a in in_avals]
    in_cts = ad.backward_pass(jaxpr, True, const, dummies, out_cts)
    in_cts = map(ad.instantiate_zeros, in_cts)
    return tree_unflatten(in_tree, in_cts)

  # Ensure that transposed_fun is a PyTree
  return Partial(transposed_fun, const)


@overload
def make_jaxpr(
    fun: Callable,
    static_argnums: int | Sequence[int] = (),
    axis_env: Sequence[tuple[AxisName, int]] | None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match the cotangent structure to the original output: inspect it once, e.g. out_struct = jax.tree_util.tree_structure(f(*primals))
2. Build cotangents via tree_map over the primal output: jax.tree_util.tree_map(jnp.zeros_like, f(*primals))
3. Re-read the linear_transpose contract: the transposed function maps output-space cotangents to input-space cotangents

Example fix

# before
ct = jnp.ones_like(x)      # shaped like the input
in_ct = f_t(ct)
# after
out = f(x)
ct = jax.tree_util.tree_map(jnp.ones_like, out)
in_ct = f_t(ct)

import jax.tree_util as jtu
expected = jtu.tree_structure(f(*primals))
assert jtu.tree_structure(cotangent) == expected, 'cotangent tree mismatch'

try:
    in_ct = f_t(ct)
except TypeError as e:
    if 'cotangent tree' in str(e):
        ct = jtu.tree_map(jnp.zeros_like, f(*primals))
        in_ct = f_t(ct)
    else:
        raise

When it happens

Trigger: f_t((a, b)) when f returned a dict; f_t(1.0) when f returned a pair; passing a flat array where the output was a tuple of arrays.

Common situations: Assuming the transpose takes inputs shaped like the function's inputs rather than outputs; the wrapped function restructures its output (e.g. returns (out, aux) with has_aux-like patterns).

Related errors

• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n
• The return value of the policies should be a boolean. Got: {

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