code: def _cond_typecheck(bind_time, *in_atoms, branches, linear):
      if not bind_time:
          _, *in_atoms = in_atoms
      avals = [x.aval for x in in_atoms]
      tc = partial(_typecheck_param, "cond")
      tc(
          branches,
          "branches",
          "tuple of ClosedJaxpr",
          type(branches) is tuple
          and all(type(x) is core.ClosedJaxpr for x in branches),
      )
      tc(
          linear,
          "linear",
          "tuple of bool",
          type(linear) is tuple
          and all(type(x) is bool for x in linear),
      )
  
      if len(branches) == 0:
          raise core.JaxprTypeError(
              "cond requires at least one branch function"
          )
      if len(linear) + 1 != len(avals):
          raise core.JaxprTypeError(
              f"cond given {len(linear)} linear flags for "
              f"{len(avals) - 1} non-predicate operands"
          )
  
      jaxpr0 = branches[0]
      jaxpr0_in_avals_str = _avals_short(jaxpr0.in_avals)
      jaxpr0_out_avals_str = _avals_short(jaxpr0.out_avals)
      joined_effects = _join_cond_effects(branches)
      disallowed_effects = (
          effects.control_flow_allowed_effects.filter_not_in(
              joined_effects
          )
      )
      if disallowed_effects:
          raise NotImplementedError(
              f"Effects not supported in `cond`: {disallowed_effects}"
          )
  
      for i, jaxpr in enumerate(branches[1:]):
          if len(jaxpr0.in_avals) != len(jaxpr.in_avals):
              raise core.JaxprTypeError(
                  f"cond branch 0 takes {len(jaxpr0.in_avals)} inputs, "
                  f"branch {i+1} takes {len(jaxpr.in_avals)}"
              )
          if len(jaxpr0.out_avals) != len(jaxpr.out_avals):
              raise core.JaxprTypeError(
                  f"cond branch 0 outputs {len(jaxpr0.out_avals)} values, "
                  f"branch {i+1} outputs {len(jaxpr.out_avals)}"
              )
          if not all(
              map(core.typematch, jaxpr0.in_avals, jaxpr.in_avals)
          ):
              raise core.JaxprTypeError(
                  f"cond branches 0 and {i+1} have mismatching input types: "
                  f"{jaxpr0_in_avals_str} vs {_avals_short(jaxpr.in_avals)}"
              )
          if not all(
              map(core.typematch, jaxpr0.out_avals, jaxpr.out_avals)
          ):
              raise core.JaxprTypeError(
                  f"cond branches 0 and {i+1} have mismatching output types: "
                  f"{jaxpr0_out_avals_str} vs {_avals_short(jaxpr.out_avals)}"
              )
  
      if len(avals) != 1 + len(jaxpr0.in_avals):
          raise core.JaxprTypeError(
              f"cond called with {len(avals) - 1} non-predicate operands, "
              f"but branches take {len(jaxpr0.in_avals)} inputs"
          )
  
      index_aval, *op_avals = avals
      if index_aval.dtype != np.int32:
          raise core.JaxprTypeError(
              f"cond called with index of type {index_aval.dtype} instead of int32"
          )
      if not all(map(core.typecompat, jaxpr0.in_avals, op_avals)):
          raise core.JaxprTypeError(
              f"cond branches take input types {jaxpr0_in_avals_str}, "
              f"called with operands of type {_avals_short(op_avals)}"
          )
      return jaxpr0.out_avals, joined_effects
  
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 21 for calling a global function with example usage: _cond_typecheck(bind_time) and returns: jaxpr0,  joined_effects