code: def _cond_batching_rule(
      spmd_axis_name,
      axis_size,
      axis_name,
      main_type,
      args,
      dims,
      branches,
      linear,
  ):
      index, *ops = args
      index_dim, *op_dims = dims
      # TODO(sharadmv): clean this up by adding a specific blocklist
      if any(
          isinstance(eff, RefEffect)
          for branch in branches
          for eff in branch.jaxpr.effects
      ):
          raise NotImplementedError(
              "State effect not supported in vmap-of-cond."
          )
      from jax._src.callback import _IOEffect, _OrderedIOEffect
  
      if any(
          eff in branch.effects
          for eff in [_IOEffect, _OrderedIOEffect]
          for branch in branches
      ):
          raise NotImplementedError(
              "IO effect not supported in vmap-of-cond."
          )
  
      if index_dim is not batching.not_mapped:
          # Convert to a lax.select. While we could get away with not broadcasting
          # some operands yet, because all outputs must be broadcast together anyway
          # for the select we broadcast the input operands for simplicity and leave
          # optimizations to XLA.
          # TODO(mattjj,frostig): assumes branches are side-effect-free, revise!
          index, *ops = (
              batching.bdim_at_front(x, d, axis_size)
              for x, d in zip(args, dims)
          )
  
          in_batched = [True] * len(branches[0].in_avals)
          out_batched = [True] * len(branches[0].out_avals)
  
          branches_batched = [
              batching.batch_jaxpr(
                  jaxpr,
                  axis_size,
                  in_batched,
                  out_batched,
                  axis_name,
                  spmd_axis_name,
                  main_type,
              )[0]
              for jaxpr in branches
          ]
  
          branch_outs = []
          for i, jaxpr in enumerate(branches_batched):
              # Perform a select on the inputs for safety of reverse-mode autodiff; see
              # https://github.com/google/jax/issues/1052
              predicate = lax.eq(index, lax._const(index, i))
              ops_ = [
                  _bcast_select(predicate, x, lax.stop_gradient(x))
                  for x in ops
              ]
              branch_outs.append(core.jaxpr_as_fun(jaxpr)(*ops_))
          out = [
              _bcast_select_n(index, *outs)
              for outs in zip(*branch_outs)
          ]
          return out, [0 if b else None for b in out_batched]
      else:
          ops_bat = [d is not batching.not_mapped for d in op_dims]
          ops = [
              batching.moveaxis(x, d, 0) if b else x
              for b, x, d in zip(ops_bat, ops, op_dims)
          ]
  
          branches_out_bat = [
              batching.batch_jaxpr(
                  jaxpr,
                  axis_size,
                  ops_bat,
                  False,
                  axis_name,
                  spmd_axis_name,
                  main_type,
              )[1]
              for jaxpr in branches
          ]
          out_bat = [any(bat) for bat in zip(*branches_out_bat)]
          branches_batched = tuple(
              batching.batch_jaxpr(
                  jaxpr,
                  axis_size,
                  ops_bat,
                  out_bat,
                  axis_name,
                  spmd_axis_name,
                  main_type,
              )[0]
              for jaxpr in branches
          )
  
          out_dims = [0 if b else batching.not_mapped for b in out_bat]
          out = cond_p.bind(
              index, *ops, branches=branches_batched, linear=linear
          )
          return out, out_dims
  
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 9 for calling a global function with example usage: _cond_batching_rule(spmd_axis_name, axis_size, axis_name, main_type, args, dims, branches, linear) and returns: out,  unknown, out,  out_dims