code: def _cond_transpose(reduce_axes, cts, *args, branches, linear):
      del linear  # could use for error checking, but see #14026
      index, *ops = args
      linear = [type(x) is ad.UndefinedPrimal for x in ops]
      in_avals = map(raise_to_shaped, branches[0].in_avals)
      num_res = len(ops) - sum(linear)
      if any(
          isinstance(eff, RefEffect)
          for branch in branches
          for eff in branch.jaxpr.effects
      ):
          raise NotImplementedError(
              "State effect not supported in cond transpose."
          )
  
      branches_trans = tuple(
          _transpose_cond_jaxpr(jaxpr, num_res, reduce_axes)
          for jaxpr in branches
      )
      lin_in_avals = [
          raise_to_shaped(a, weak_type=False)
          for a, l in zip(in_avals, linear)
          if l
      ]
      assert all(
          core.typematch(out_aval, lin_in_aval)
          for jaxpr in branches_trans
          for out_aval, lin_in_aval in zip(
              jaxpr.out_avals, lin_in_avals
          )
      )
  
      res = ops[:num_res]
      cts = map(ad.instantiate_zeros, cts)
      linear_trans = (False,) * num_res + (True,) * len(cts)
  
      out = cond_p.bind(
          index,
          *res,
          *cts,
          branches=branches_trans,
          linear=linear_trans,
      )
      assert all(map(core.typecheck, lin_in_avals, out))
  
      out_iter = iter(out)
      out = [next(out_iter) if l else None for l in linear]
      assert next(out_iter, None) is None
      return [None] + out
  
imports: from __future__ import annotations
  import collections
  from collections.abc import Sequence
  import functools
  from functools import partial
  import inspect
  import itertools
  import operator
  from typing import Any, Callable, TypeVar
  from jax.tree_util import tree_flatten, tree_unflatten
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import util
  from jax._src.state.discharge import register_discharge_rule, discharge_state
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.core import ConcreteArray, raise_to_shaped, replace_jaxpr_effects
  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._src.interpreters import xla
  from jax._src.lax import lax
  from jax._src.traceback_util import api_boundary
  from jax._src.util import safe_map, split_list, partition_list
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  import numpy
  from jax._src.lax.control_flow.common import _abstractify, _avals_short, _check_tree_and_avals, _initial_style_jaxprs_with_common_consts, _make_closed_jaxpr, _prune_zeros, _typecheck_param
label_desc: Use image node_id 19 for calling a global function with example usage: _cond_transpose(reduce_axes, cts) and returns: unknown