jax-ml/jax · error · TypeError

transpose output pytree structure must match that of linear

Error message

transpose output pytree structure must match that of linear inputs, got output structure {t_out_avals.tree} and input structure {lin_tree}.

What it means

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.

Source

Thrown at jax/_src/custom_derivatives.py:1575

      fun,
      ft.pack(((ft.FTPyTree(res_avals, res_tree),
                ft.FTPyTree(lin_avals, lin_tree)), {})),
      debug_info("linear_call fun", fun, (residual_args, linear_args), {}))
  f_jaxpr_closed, f_consts = pe.separate_consts(f_jaxpr)
  out_avals = f_jaxpr_closed.out_avals
  out_tree = f_out_avals.tree

  @pe._memoize
  def transpose_thunk():
    t_jaxpr, t_out_avals = pe.trace_to_jaxpr(
        fun_transpose,
        ft.pack(((ft.FTPyTree(res_avals, res_tree),
                  ft.FTPyTree(list(out_avals), out_tree)), {})),
        # TODO(necula): the fun_transpose takes residual and output of fun!
        debug_info("linear_call fun_transpose", fun_transpose,
                   (residual_args, linear_args), {}).with_unknown_names())
    if t_out_avals.tree != lin_tree:
      raise TypeError(
          'transpose output pytree structure must match that of linear inputs, '
          f'got output structure {t_out_avals.tree} '
          f'and input structure {lin_tree}.')
    return pe.separate_consts(t_jaxpr)

  out = linear_call_p.bind(*f_consts, *operands_res, *operands_lin,
                           callee=f_jaxpr_closed,
                           transpose_thunk=transpose_thunk,
                           num_callee_consts=len(f_consts),
                           num_res=len(operands_res))

  return tree_unflatten(out_tree, out)

def _linear_call_impl(*args, callee, transpose_thunk, num_callee_consts,
                      num_res):
  del transpose_thunk, num_callee_consts, num_res
  return core.eval_jaxpr(callee, (), *args)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make fun_transpose return a pytree structurally identical to the linear_args argument it receives
2. If there is one linear input, still return a matching structure (e.g. tuple/list as passed)
3. Write a quick unit test calling the transpose directly and compare tree_structure(output) to tree_structure(linear_args)

Example fix

# before
def trans(res, lin_out):
  return jnp.dot(A, lin_out)  # single array, but lin args were (x, y)
jax.linear_call(f, trans, ...)

# after
def trans(res, lin_out):
  return (jnp.dot(A, lin_out[0]), jnp.dot(B, lin_out[1]))
jax.linear_call(f, trans, ...)

from jax.tree_util import tree_structure
out = fun_transpose(res_example, lin_example)
assert tree_structure(out) == tree_structure(lin_example)

When it happens

Trigger: 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.

Common situations: 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.

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/b086ee6689245c90. Report an issue: GitHub.