code: def get_jaxpr_type_signature(
      jaxpr: core.Jaxpr,
      consumed_inputs: list[bool | np.ndarray] | None = None,
  ) -> KeyReuseSignature:
      """Parse the jaxpr to determine key reuse signature"""
      consumed: dict[core.Atom, bool | np.ndarray] = {}
  
      def is_key(var: core.Atom):
          return hasattr(var.aval, "dtype") and jax.dtypes.issubdtype(
              var.aval.dtype, jax.dtypes.prng_key
          )
  
      def sink(var: core.Atom, mask=True):
          if not is_key(var):
              return
          assert not isinstance(var, core.Literal)
          if np.any(np.logical_and(consumed.get(var, False), mask)):
              return True
          consumed[var] = np.logical_or(consumed.get(var, False), mask)
  
      def source(var: core.Atom, mask=False):
          if not is_key(var):
              return
          assert not isinstance(var, core.Literal)
          consumed[var] = mask
  
      def is_consumed(var: core.Atom):
          if isinstance(var, core.Literal):
              return False
          return consumed.get(var, False)
  
      if consumed_inputs:
          for var, mask in util.safe_zip(jaxpr.invars, consumed_inputs):
              if not isinstance(var, core.Literal):
                  source(var, mask)
  
      for eqn in jaxpr.eqns:
          if eqn.primitive in key_reuse_signatures:
              signature = key_reuse_signatures[eqn.primitive]
          elif eqn.primitive in key_reuse_signatures_dynamic:
              args_consumed = [is_consumed(var) for var in eqn.invars]
              signature = key_reuse_signatures_dynamic[eqn.primitive](
                  eqn, args_consumed
              )
          else:
              args_consumed = [is_consumed(var) for var in eqn.invars]
              signature = unknown_signature(eqn, args_consumed)
  
          for snk in signature.sinks:
              if sink(eqn.invars[snk.idx], snk.mask):
                  raise KeyReuseError(
                      f"In {eqn.primitive}, key values {eqn.invars[snk.idx]} are already consumed.\n"
                      f"eqn: {eqn}\njaxpr:\n{jaxpr}"
                  )
          for var in eqn.outvars:
              if not isinstance(var, core.Literal):
                  source(var, True)  # consumed unless in a Source.
          for src in signature.sources:
              source(eqn.outvars[src.idx])
  
      forwards = [
          v
          for v in jaxpr.outvars
          if is_key(v)
          and v in jaxpr.invars
          and not np.any(consumed.get(v, False))
      ]
      sinks = [
          v
          for v in jaxpr.invars
          if is_key(v) and np.any(consumed.get(v, False))
      ]
      sources = [
          v
          for v in jaxpr.outvars
          if is_key(v) and not np.any(consumed.get(v, False))
      ]
      return KeyReuseSignature(
          sinks=[
              Sink(i, True if v in forwards else consumed[v])
              for i, v in enumerate(jaxpr.invars)
              if v in forwards or v in sinks
          ],
          sources=[
              Source(i)
              for i, v in enumerate(jaxpr.outvars)
              if (v in forwards or v in sources)
              and v
              not in jaxpr.outvars[
                  :i
              ]  # Only source the first of duplicate return values
          ],
      )
  
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 2 for calling a global function with example usage: get_jaxpr_type_signature(jaxpr, consumed_inputs) and returns: KeyReuseSignature, unknown, consumed, True, False