code: def transpose_jaxpr(
      jaxpr: core.Jaxpr, which_linear: list[bool]
  ) -> core.Jaxpr:
      def trans(i, *args):
          # First we want to run the computation to read all the residual refs. We can
          # do that by using partial evaluation with all linear inputs unknown.
          res_jaxpr, tangent_jaxpr_, *_ = _partial_eval_jaxpr_custom(
              jaxpr, [False, *which_linear], _save_everything
          )
          res_args = [
              x for x, lin in zip(args, which_linear) if not lin
          ]
          res = core.eval_jaxpr(res_jaxpr, (), i, *res_args)
  
          # Now that we have residual values, we run the tangent jaxpr. It takes as
          # input the residuals, the loop index, and all the refs (at least, the ones
          # that are used in the body). Luckily, `tangent_jaxpr_` has all known and
          # unknown inputs!
          tangent_jaxpr, used = pe.dce_jaxpr(tangent_jaxpr_, [])
          used_res, (used_i,), used_ct = split_list(used, [len(res), 1])
          primals_args = [*(r for u, r in zip(used_res, res) if u)]
          if used_i:
              primals_args = [*primals_args, i]
          ct_args = [x for x, u in zip(args, used_ct) if u]
          ad.backward_pass(
              tangent_jaxpr,
              (),
              False,
              (),
              (*primals_args, *ct_args),
              (),
          )
          return []
  
      jaxpr_trans, _, _ = pe.trace_to_jaxpr_dynamic(
          lu.wrap_init(trans), [v.aval for v in jaxpr.invars]
      )
      return jaxpr_trans
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import operator
  from typing import Any, Callable, Generic, TypeVar
  import jax.numpy
  from jax import lax
  from jax.api_util import flatten_fun_nokwargs
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax.tree_util import tree_flatten, tree_structure, tree_unflatten, treedef_tuple, tree_map, tree_leaves, PyTreeDef
  from jax._src import ad_util
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src.state.types import ReadEffect, AbstractRef, StateEffect
  from jax._src.state import discharge
  from jax._src.state import primitives
  from jax._src.state import utils
  from jax._src.state import types
  from jax._src.typing import Array
  from jax._src.util import partition_list, merge_lists, safe_map, safe_zip, split_list, split_dict, weakref_lru_cache
  from jax._src.lax.control_flow import loops
  from jax._src.lax.control_flow.common import _abstractify, _initial_style_jaxpr
label_desc: Use image node_id 19 for calling a global function with example usage: transpose_jaxpr(jaxpr, which_linear) and returns: jaxpr_trans, list