code: def make_transpose_from_thunk(thunk, lin_tree):
      transpose_jaxpr, transpose_consts = thunk()
      transpose_jaxpr = core.ClosedJaxpr(
          pe.convert_constvars_jaxpr(transpose_jaxpr), ()
      )
  
      def transpose(res_arg, ct_out):
          args_flat = tree_leaves((res_arg, ct_out))
          ct_ins = core.jaxpr_as_fun(transpose_jaxpr)(
              *transpose_consts, *args_flat
          )
          return tree_unflatten(lin_tree, ct_ins)
  
      return transpose
  
imports: from __future__ import annotations
  import functools
  from typing import Any, Callable
  from jax._src import ad_util
  from jax._src import api_util
  from jax._src import core
  from jax._src import custom_api_util
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import traceback_util
  from jax._src import util
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.interpreters import xla
  from jax._src.tree_util import tree_flatten, tree_leaves, tree_map, tree_structure, treedef_tuple, tree_unflatten
label_desc: Use image node_id 8 for calling a global function with example usage: make_transpose_from_thunk(thunk, lin_tree) and returns: transpose, tree_unflatten