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()
      ]
      combined_forwards = [
          Forward(f.in_idx + 1, f.out_idx)
          for f in set.intersection(
              *(set(sig.forwards) for sig in signatures)
          )
      ]
      return KeyReuseSignatureWithForwards(
          combined_sinks, combined_sources, combined_forwards
      )
  
imports: from __future__ import annotations
  from collections import defaultdict
  from functools import reduce
  from typing import Any, Callable, NamedTuple
  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: KeyReuseSignatureWithForwards