code: def _for_transpose(
      in_cts, *args, jaxpr, nsteps, reverse, which_linear, unroll
  ):
      # if any in_ct is nonzero, we definitely want it in args_ (and the
      # corresponding x in args could be an undefined primal, but doesn't have to be)
      # for non-res stuff:
      #   getting and setting => (nonzero ct, UndefinedPrimal arg)
      #   just setting =>        (nonzero ct, not UndefinedPrimal, dummy value)
      #   just getting =>        (zero ct   , UndefinedPrimal arg)
      # for res stuff:
      #                          (zero ct   , not UndefinedPrimal)
      args_ = []
      which_linear_transpose = []
      for x, ct in zip(args, in_cts):
          if type(ct) is ad_util.Zero and not ad.is_undefined_primal(x):
              # this is a residual, take x!
              args_.append(x)
              which_linear_transpose.append(False)
          elif type(ct) is ad_util.Zero and ad.is_undefined_primal(x):
              # the loop was 'just getting', plug in a zero
              args_.append(ad_util.zeros_like_aval(x.aval))
              which_linear_transpose.append(False)
          elif type(
              ct
          ) is not ad_util.Zero and not ad.is_undefined_primal(x):
              # the loop was 'just setting', grab that cotangent! x is dummy
              args_.append(ct)
              which_linear_transpose.append(False)
          elif type(ct) is not ad_util.Zero and ad.is_undefined_primal(
              x
          ):
              # the loop was 'getting and setting', grab that cotangent!
              args_.append(ct)
              which_linear_transpose.append(True)
  
      jaxpr_transpose = transpose_jaxpr(jaxpr, which_linear)
      assert len(args_) == len(jaxpr_transpose.invars) - 1
      all_outs = for_p.bind(
          *args_,
          jaxpr=jaxpr_transpose,
          nsteps=nsteps,
          reverse=not reverse,
          which_linear=tuple(which_linear_transpose),
          unroll=unroll,
      )
      ct_outs = [
          ct if ad.is_undefined_primal(x) else None
          for x, ct in zip(args, all_outs)
      ]
      return ct_outs
  
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 20 for calling a global function with example usage: _for_transpose(in_cts) and returns: ct_outs