jax-ml/jax · error · TypeError

The inputs to the closure produced by closure_convert must h

Error message

The inputs to the closure produced by closure_convert must have the same Pytree structure as the example arguments passed when closure_convert was called. Expected {in_tree}, but got {in_tree2}

What it means

jax.closure_convert traces a function against example arguments and returns a converted closure that can only be called with arguments having the identical Pytree structure. This TypeError fires when the closure is later invoked with a differently-structured argument tree.

Source

Thrown at jax/_src/custom_derivatives.py:1447

  closed_jaxpr, out_avals = pe.trace_to_jaxpr(
      fun, ft.treedef_args_to_ft(in_tree, in_avals), debug_info)
  jaxpr, consts = pe.separate_consts(closed_jaxpr)
  out_tree = out_avals.tree

  (closure_consts, const_args), merge = partition_list(_maybe_perturbed, consts)
  num_consts = len(const_args)

  def converted_fun(*args_hconsts):
    num_args = len(args_hconsts) - num_consts
    args, const_args = split_list(args_hconsts, [num_args])
    consts = merge(closure_consts, const_args)
    all_args, in_tree2 = tree_flatten((tuple(args), {}))
    if in_tree != in_tree2:
      msg = ("The inputs to the closure produced by closure_convert must have "
             "the same Pytree structure as the example arguments passed when "
             f"closure_convert was called. Expected {in_tree}, but got "
             f"{in_tree2}")
      raise TypeError(msg)
    out_flat = core.eval_jaxpr(jaxpr, consts, *all_args)
    return tree_unflatten(out_tree, out_flat)

  return converted_fun, const_args

def partition_list(choice, lst):
  out = [], []
  which = [out[choice(elt)].append(elt) or choice(elt) for elt in lst]
  def merge(l1, l2):
    i1, i2 = iter(l1), iter(l2)
    return [next(i2 if snd else i1) for snd in which]
  return out, merge


### Custom transposition

def linear_call(fun: Callable,
                fun_transpose: Callable, residual_args,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call converted_fun with exactly the same pytree structure as the example arguments used in closure_convert
2. Re-run closure_convert with new example args if the structure legitimately changes
3. Flatten/uniformize arguments (e.g. always tuples of arrays) before conversion and at every call site

Example fix

# before
conv, consts = closure_convert(f, (1.0, 2.0))
conv([1.0, 2.0])  # list vs tuple -> structure mismatch

# after
conv, consts = closure_convert(f, (1.0, 2.0))
conv((3.0, 4.0))

from jax.tree_util import tree_structure
expected = tree_structure((example_args, {}))
assert tree_structure((new_args, {})) == expected

def same_tree(expected, args):
    return tree_structure((expected, {})) == tree_structure((args, {}))

When it happens

Trigger: Calling the function returned by closure_convert(fun, *example_args) with args whose pytree structure (shapes of containers, not values) differs from example_args, e.g. passing a list where a tuple was given, or extra/missing leaves.

Common situations: Using closure_convert to capture non-JAX-type constants (custom objects, C++ pointers) in a function later called from jit/scan with different container nesting; refactoring call sites after conversion.

Related errors

• scan body output must be a pair, got {}.
• 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

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