code: def _cond_with_per_branch_args(
      pred,
      true_operand,
      true_fun: Callable,
      false_operand,
      false_fun: Callable,
  ):
      """Conditionally apply ``true_fun`` or ``false_fun``.
  
      Has equivalent semantics to this Python implementation::
  
        def cond(pred, true_operand, true_fun, false_operand, false_fun):
          if pred:
            return true_fun(true_operand)
          else:
            return false_fun(false_operand)
  
      Pred has to be a scalar type, collection types (list, tuple) are not supported
      """
      if not (callable(true_fun) and callable(false_fun)):
          raise TypeError(
              "lax.cond: true_fun and false_fun arguments should be callable."
          )
      return _cond(
          pred,
          lambda op: true_fun(op[0]),
          lambda op: false_fun(op[1]),
          (true_operand, false_operand),
      )
  
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 4 for calling a global function with example usage: _cond_with_per_branch_args(pred, true_operand, true_fun, false_operand, false_fun) and returns: _cond