code: def _transpose_jaxpr(
      jaxpr: core.Jaxpr, which_linear: Sequence[bool]
  ) -> tuple[core.Jaxpr, Any]:
      def trans(*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_,
              *_,
              num_res_out,
              num_res_ref,
          ) = pe.partial_eval_jaxpr_stateful(
              jaxpr,
              which_linear,
              in_inst=which_linear,
              ensure_out_inst=[],
              ensure_out_unknowns=[],
              saveable=_save_everything,
          )
  
          num_unknown = sum(which_linear)
          num_known = len(jaxpr.invars) - num_unknown
          res_args, _ = partition_list(which_linear, args)
          res_jaxpr_avals = [v.aval for v in res_jaxpr_.invars]
          _, res_avals = split_list(res_jaxpr_avals, [num_known])
          res_avals = [
              a.inner_aval for a in res_avals
          ]  # pytype: disable=attribute-error
          all_avals = [
              *res_avals,
              *[v.aval for v in res_jaxpr_.outvars],
          ]
          empty_res = map(ad.zeros_like_aval, all_avals)
          res_jaxpr, _ = _convert_outputs_to_writes(res_jaxpr_)
          res = run_state_p.bind(
              *res_args,
              *empty_res,
              jaxpr=res_jaxpr,
              which_linear=(False,) * (len(res_args) + len(empty_res)),
          )
          res = res[len(res_args) :]
          ref_res_, nonref_res_ = split_list(res, [num_res_ref])
  
          # 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_inputs = pe.dce_jaxpr(tangent_jaxpr_, [])
          used_res, used_cts = split_list(used_inputs, [len(res)])
          used_nonref_res, used_ref_res = split_list(
              used_res, [num_res_out]
          )
          _, nonref_res = partition_list(used_nonref_res, nonref_res_)
          _, ref_res = partition_list(used_ref_res, ref_res_)
          primals_args = [*nonref_res, *ref_res]
          _, tangent_args = partition_list(which_linear, args)
          _, ct_args = partition_list(used_cts, tangent_args)
          ad.backward_pass(
              tangent_jaxpr,
              (),
              False,
              (),
              (*primals_args, *ct_args),
              (),
          )
          return []
  
      jaxpr_trans, _, consts = pe.trace_to_jaxpr_dynamic(
          lu.wrap_init(trans), [v.aval for v in jaxpr.invars]
      )
      return jaxpr_trans, consts
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import dataclasses
  from functools import partial
  import operator
  from typing import Any, Callable, Protocol
  import numpy
  from jax._src import api_util
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import tree_util
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.lax import lax
  from jax._src.lax import slicing
  from jax._src.state import indexing
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.state.primitives import get_p, swap_p, addupdate_p
  from jax._src.state.utils import hoist_consts_to_refs
  from jax._src.typing import Array
  from jax._src.util import safe_map, safe_zip, split_list, weakref_lru_cache, partition_list, merge_lists, split_dict
label_desc: Use image node_id 27 for calling a global function with example usage: _transpose_jaxpr(jaxpr, which_linear) and returns: jaxpr_trans,  consts, list