{"record":{"id":"b086ee6689245c90","repo":"jax-ml/jax","slug":"transpose-output-pytree-structure-must-match-that","errorCode":null,"errorMessage":"transpose output pytree structure must match that of linear inputs, got output structure {t_out_avals.tree} and input structure {lin_tree}.","messagePattern":"transpose output pytree structure must match that of linear inputs, got output structure (.+?) and input structure (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":1575,"sourceCode":"      fun,\n      ft.pack(((ft.FTPyTree(res_avals, res_tree),\n                ft.FTPyTree(lin_avals, lin_tree)), {})),\n      debug_info(\"linear_call fun\", fun, (residual_args, linear_args), {}))\n  f_jaxpr_closed, f_consts = pe.separate_consts(f_jaxpr)\n  out_avals = f_jaxpr_closed.out_avals\n  out_tree = f_out_avals.tree\n\n  @pe._memoize\n  def transpose_thunk():\n    t_jaxpr, t_out_avals = pe.trace_to_jaxpr(\n        fun_transpose,\n        ft.pack(((ft.FTPyTree(res_avals, res_tree),\n                  ft.FTPyTree(list(out_avals), out_tree)), {})),\n        # TODO(necula): the fun_transpose takes residual and output of fun!\n        debug_info(\"linear_call fun_transpose\", fun_transpose,\n                   (residual_args, linear_args), {}).with_unknown_names())\n    if t_out_avals.tree != lin_tree:\n      raise TypeError(\n          'transpose output pytree structure must match that of linear inputs, '\n          f'got output structure {t_out_avals.tree} '\n          f'and input structure {lin_tree}.')\n    return pe.separate_consts(t_jaxpr)\n\n  out = linear_call_p.bind(*f_consts, *operands_res, *operands_lin,\n                           callee=f_jaxpr_closed,\n                           transpose_thunk=transpose_thunk,\n                           num_callee_consts=len(f_consts),\n                           num_res=len(operands_res))\n\n  return tree_unflatten(out_tree, out)\n\ndef _linear_call_impl(*args, callee, transpose_thunk, num_callee_consts,\n                      num_res):\n  del transpose_thunk, num_callee_consts, num_res\n  return core.eval_jaxpr(callee, (), *args)\n","sourceCodeStart":1557,"sourceCodeEnd":1593,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L1557-L1593","documentation":"jax.linear_call invokes the user-supplied transpose function and requires its outputs to have the same Pytree structure as the linear inputs. This TypeError is raised when the transposed function returns a differently-structured pytree.","triggerScenarios":"Using jax.linear_call(fun, fun_transpose, ...) where fun_transpose returns e.g. a single array instead of a tuple matching the linear args structure, or nests containers differently.","commonSituations":"Writing custom linear operations (e.g. custom matmul/conv) with linear_call and forgetting the transpose must mirror the linear-input structure; returning a scalar where a tuple is expected.","solutions":["Make fun_transpose return a pytree structurally identical to the linear_args argument it receives","If there is one linear input, still return a matching structure (e.g. tuple/list as passed)","Write a quick unit test calling the transpose directly and compare tree_structure(output) to tree_structure(linear_args)"],"exampleFix":"# before\ndef trans(res, lin_out):\n  return jnp.dot(A, lin_out)  # single array, but lin args were (x, y)\njax.linear_call(f, trans, ...)\n\n# after\ndef trans(res, lin_out):\n  return (jnp.dot(A, lin_out[0]), jnp.dot(B, lin_out[1]))\njax.linear_call(f, trans, ...)","handlingStrategy":"validation","validationCode":"from jax.tree_util import tree_structure\nout = fun_transpose(res_example, lin_example)\nassert tree_structure(out) == tree_structure(lin_example)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Unit-test the transpose function standalone before wiring it into linear_call","Mirror the linear-args pytree exactly in transpose outputs"],"tags":["jax","linear-call","transpose","pytree"],"backgroundTag":"transpose-pytree-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}