code: def _cond_lowering(ctx, index, *args, branches, linear):
      del linear  # Unused.
      joined_effects = core.join_effects(
          *(branch.effects for branch in branches)
      )
      ordered_effects = list(
          effects.ordered_effects.filter_in(joined_effects)
      )
      num_tokens = len(ordered_effects)
      tokens_in = ctx.tokens_in.subset(ordered_effects)
      output_token_types = [mlir.token_type() for _ in ordered_effects]
      output_types = [
          *output_token_types,
          *map(mlir.aval_to_ir_types, ctx.avals_out),
      ]
      flat_output_types = util.flatten(output_types)
  
      # CaseOp takes a single argument 'index' and the corresponding blocks
      # have no arguments; the computation within the block uses implicit
      # captures.
      case_op = hlo.CaseOp(
          flat_output_types, index=index, num_branches=len(branches)
      )
      name_stack = ctx.module_context.name_stack.extend("cond")
      for i, jaxpr in enumerate(branches):
          branch = case_op.regions[i].blocks.append()
          with ir.InsertionPoint(branch):
              sub_ctx = ctx.module_context.replace(
                  name_stack=name_stack.extend(f"branch_{i}_fun")
              )
              consts = [
                  mlir.ir_constants(xla.canonicalize_dtype(x))
                  for x in jaxpr.consts
              ]
              out_vals, tokens_out = mlir.jaxpr_subcomp(
                  sub_ctx,
                  jaxpr.jaxpr,
                  tokens_in,
                  consts,
                  *map(mlir.wrap_singleton_ir_values, args),
                  dim_var_values=ctx.dim_var_values,
              )
              out_tokens = [
                  tokens_out.get(eff) for eff in ordered_effects
              ]
              out_vals = [*out_tokens, *out_vals]
              hlo.return_(util.flatten(out_vals))
  
      tokens_and_outputs = util.unflatten(
          case_op.results, map(len, output_types)
      )
      tokens, outputs = util.split_list(
          tokens_and_outputs, [num_tokens]
      )
      ctx.set_tokens_out(mlir.TokenSet(zip(ordered_effects, tokens)))
      return outputs
  
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 23 for calling a global function with example usage: _cond_lowering(ctx, index) and returns: outputs