code: def _cond_key_type_signature(eqn, args_consumed):
      signatures = [
          get_jaxpr_type_signature(
              branch.jaxpr, consumed_inputs=args_consumed[1:]
          )
          for branch in eqn.params["branches"]
      ]
      sinks = defaultdict(list)
      sources = defaultdict(list)
      for sig in signatures:
          for sink in sig.sinks:
              sinks[sink.idx].append(sink.mask)
          for source in sig.sources:
              sources[source.idx].append(source.mask)
  
      combined_sinks = [
          Sink(i + 1, reduce(np.logical_or, m))
          for i, m in sinks.items()
      ]
      combined_sources = [
          Source(i + 1, reduce(np.logical_and, m))
          for i, m in sources.items()
      ]
      return KeyReuseSignature(combined_sinks, combined_sources)
  
imports: from __future__ import annotations
  from collections import defaultdict
  from functools import reduce
  from typing import Any, Callable
  import jax
  from jax import core
  from jax import lax
  from jax import tree_util
  from jax._src import api_util
  from jax._src import linear_util
  from jax._src import pjit
  from jax._src import prng
  from jax._src import random
  from jax._src import util
  from jax._src.debugging import debug_callback_p
  from jax._src.interpreters import partial_eval
  from jax.experimental.key_reuse._common import consume_p, unconsumed_copy_p, assert_consumed_value_p, KeyReuseError, Sink, Source, KeyReuseSignature
  import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: _cond_key_type_signature(eqn, args_consumed) and returns: KeyReuseSignature