{"record":{"id":"297081fbba15518c","repo":"jax-ml/jax","slug":"the-inputs-to-the-closure-produced-by-closure-conv","errorCode":null,"errorMessage":"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}","messagePattern":"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 (.+?), but got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":1447,"sourceCode":"  closed_jaxpr, out_avals = pe.trace_to_jaxpr(\n      fun, ft.treedef_args_to_ft(in_tree, in_avals), debug_info)\n  jaxpr, consts = pe.separate_consts(closed_jaxpr)\n  out_tree = out_avals.tree\n\n  (closure_consts, const_args), merge = partition_list(_maybe_perturbed, consts)\n  num_consts = len(const_args)\n\n  def converted_fun(*args_hconsts):\n    num_args = len(args_hconsts) - num_consts\n    args, const_args = split_list(args_hconsts, [num_args])\n    consts = merge(closure_consts, const_args)\n    all_args, in_tree2 = tree_flatten((tuple(args), {}))\n    if in_tree != in_tree2:\n      msg = (\"The inputs to the closure produced by closure_convert must have \"\n             \"the same Pytree structure as the example arguments passed when \"\n             f\"closure_convert was called. Expected {in_tree}, but got \"\n             f\"{in_tree2}\")\n      raise TypeError(msg)\n    out_flat = core.eval_jaxpr(jaxpr, consts, *all_args)\n    return tree_unflatten(out_tree, out_flat)\n\n  return converted_fun, const_args\n\ndef partition_list(choice, lst):\n  out = [], []\n  which = [out[choice(elt)].append(elt) or choice(elt) for elt in lst]\n  def merge(l1, l2):\n    i1, i2 = iter(l1), iter(l2)\n    return [next(i2 if snd else i1) for snd in which]\n  return out, merge\n\n\n### Custom transposition\n\ndef linear_call(fun: Callable,\n                fun_transpose: Callable, residual_args,","sourceCodeStart":1429,"sourceCodeEnd":1465,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L1429-L1465","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Call converted_fun with exactly the same pytree structure as the example arguments used in closure_convert","Re-run closure_convert with new example args if the structure legitimately changes","Flatten/uniformize arguments (e.g. always tuples of arrays) before conversion and at every call site"],"exampleFix":"# before\nconv, consts = closure_convert(f, (1.0, 2.0))\nconv([1.0, 2.0])  # list vs tuple -> structure mismatch\n\n# after\nconv, consts = closure_convert(f, (1.0, 2.0))\nconv((3.0, 4.0))","handlingStrategy":"type-guard","validationCode":"from jax.tree_util import tree_structure\nexpected = tree_structure((example_args, {}))\nassert tree_structure((new_args, {})) == expected","typeGuard":"def same_tree(expected, args):\n    return tree_structure((expected, {})) == tree_structure((args, {}))","tryCatchPattern":null,"preventionTips":["Fix argument container types (always tuples) at all call sites of the converted closure","Re-run closure_convert whenever the argument structure changes"],"tags":["jax","closure-convert","pytree","typeerror"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}