code: def _run_state_transpose(
      in_cts, *args, jaxpr: core.Jaxpr, which_linear: tuple[bool, ...]
  ):
      # 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)
      assert any(which_linear)
      transpose_args = []
      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!
              transpose_args.append(x)
          elif type(ct) is ad_util.Zero and ad.is_undefined_primal(x):
              # the loop was 'just getting', plug in a zero
              transpose_args.append(ad_util.zeros_like_aval(x.aval))
          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
              transpose_args.append(ct)
          elif type(ct) is not ad_util.Zero and ad.is_undefined_primal(
              x
          ):
              # the loop was 'getting and setting', grab that cotangent!
              transpose_args.append(ct)
      jaxpr_transpose_, consts = _transpose_jaxpr(jaxpr, which_linear)
      jaxpr_transpose = hoist_consts_to_refs(jaxpr_transpose_)
      which_linear = (*[False] * len(consts), *which_linear)
      const_all_outs = run_state_p.bind(
          *consts,
          *transpose_args,
          jaxpr=jaxpr_transpose,
          which_linear=which_linear,
      )
      _, all_outs = split_list(const_all_outs, [len(consts)])
      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 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 28 for calling a global function with example usage: _run_state_transpose(in_cts) and returns: ct_outs