{"record":{"id":"8933ba951d1335ee","repo":"jax-ml/jax","slug":"cotangent-tree-does-not-match-function-output-exp","errorCode":null,"errorMessage":"cotangent tree does not match function output, expected {out_tree()} but got {out_tree2}","messagePattern":"cotangent tree does not match function output, expected (.+?) but got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":2121,"sourceCode":"\n  in_pvals = map(pe.PartialVal.unknown, in_avals)\n  jaxpr, out_pvals, const = pe.trace_to_jaxpr_nounits(flat_fun, in_pvals,\n                                                      instantiate=True)\n  jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True)\n  out_avals, _ = unzip2(out_pvals)\n  out_dtypes = [a.dtype for a in out_avals if not a.is_high]\n  if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes)\n          or all(dtypes.issubdtype(d, np.integer)\n                 for d in in_dtypes + out_dtypes)):\n    raise TypeError(\"linear_transpose only supports [float or complex] -> \"\n                    \"[float or complex], and integer -> integer functions, \"\n                    f\"but got {in_dtypes} -> {out_dtypes}.\")\n\n  @api_boundary\n  def transposed_fun(const, out_cotangent):\n    out_cts, out_tree2 = tree_flatten(out_cotangent)\n    if out_tree() != out_tree2:\n      raise TypeError(\"cotangent tree does not match function output, \"\n                      f\"expected {out_tree()} but got {out_tree2}\")\n    if not all(map(core.typecheck, out_avals, out_cts)):\n      raise TypeError(\"cotangent type does not match function output, \"\n                      f\"expected {out_avals} but got {out_cts}\")\n    dummies = [ad.UndefinedPrimal(a.to_ct_aval()) for a in in_avals]\n    in_cts = ad.backward_pass(jaxpr, True, const, dummies, out_cts)\n    in_cts = map(ad.instantiate_zeros, in_cts)\n    return tree_unflatten(in_tree, in_cts)\n\n  # Ensure that transposed_fun is a PyTree\n  return Partial(transposed_fun, const)\n\n\n@overload\ndef make_jaxpr(\n    fun: Callable,\n    static_argnums: int | Sequence[int] = (),\n    axis_env: Sequence[tuple[AxisName, int]] | None = None,","sourceCodeStart":2103,"sourceCodeEnd":2139,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L2103-L2139","documentation":"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.","triggerScenarios":"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.","commonSituations":"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).","solutions":["Match the cotangent structure to the original output: inspect it once, e.g. out_struct = jax.tree_util.tree_structure(f(*primals))","Build cotangents via tree_map over the primal output: jax.tree_util.tree_map(jnp.zeros_like, f(*primals))","Re-read the linear_transpose contract: the transposed function maps output-space cotangents to input-space cotangents"],"exampleFix":"# before\nct = jnp.ones_like(x)      # shaped like the input\nin_ct = f_t(ct)\n# after\nout = f(x)\nct = jax.tree_util.tree_map(jnp.ones_like, out)\nin_ct = f_t(ct)","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\nexpected = jtu.tree_structure(f(*primals))\nassert jtu.tree_structure(cotangent) == expected, 'cotangent tree mismatch'","typeGuard":null,"tryCatchPattern":"try:\n    in_ct = f_t(ct)\nexcept TypeError as e:\n    if 'cotangent tree' in str(e):\n        ct = jtu.tree_map(jnp.zeros_like, f(*primals))\n        in_ct = f_t(ct)\n    else:\n        raise","preventionTips":["Cache tree_structure(f(*primals)) and validate before calling the transpose","Derive cotangents from actual outputs with tree_map"],"tags":["jax","linear-transpose","pytree"],"backgroundTag":"cotangent-tree-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}