code: def get_jaxpr_type_signature(
      jaxpr: core.Jaxpr,
      consumed_inputs: list[bool | np.ndarray] | None = None,
      forwarded_inputs: dict[int, int] | None = None,
  ) -> KeyReuseSignatureWithForwards:
      """Parse the jaxpr to determine key reuse signature"""
      consumed: dict[core.Atom, bool | np.ndarray] = {}
      forwards: dict[
          core.Atom, core.Atom
      ] = {}  # map forwarded outputs to inputs.
  
      def resolve_forwards(var: core.Atom) -> core.Atom:
          if not forwards:
              return var
          for _ in range(len(forwards) + 1):
              if isinstance(var, core.Literal):
                  return var
              if var in forwards:
                  var = forwards[var]
              else:
                  return var
          raise ValueError("forwarding cycle detected")
  
      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
          var = resolve_forwards(var)
          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
          var = resolve_forwards(var)
          assert not isinstance(var, core.Literal)
          consumed[var] = mask
  
      def is_consumed(var: core.Atom):
          var = resolve_forwards(var)
          if isinstance(var, core.Literal):
              return False
          return consumed.get(var, False)
  
      if forwarded_inputs:
          for i, j in forwarded_inputs.items():
              forwards[jaxpr.invars[i]] = jaxpr.invars[j]
  
      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 in_idx, out_idx in signature.forwards:
              forwards[eqn.outvars[out_idx]] = eqn.invars[in_idx]
  
          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)
                  and var not in forwards
              ):
                  source(var, True)  # consumed unless in a Source.
          for src in signature.sources:
              source(eqn.outvars[src.idx])
  
      return KeyReuseSignatureWithForwards(
          sinks=[
              Sink(i, consumed[v])
              for i, v in enumerate(jaxpr.invars)
              if is_key(v) and np.any(consumed.get(v, False))
          ],
          sources=[
              Source(i)
              for i, v in enumerate(jaxpr.outvars)
              if is_key(v)
              and resolve_forwards(v) not in jaxpr.invars
              and not consumed.get(v, False)
          ],
          forwards=[
              Forward(jaxpr.invars.index(resolve_forwards(outvar)), idx_out)  # type: ignore[arg-type]
              for idx_out, outvar in enumerate(jaxpr.outvars)
              if is_key(outvar)
              and resolve_forwards(outvar) in jaxpr.invars
          ],
      )
  
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 2 for calling a global function with example usage: get_jaxpr_type_signature(jaxpr, consumed_inputs, forwarded_inputs) and returns: KeyReuseSignatureWithForwards, unknown, consumed, var, True, False, var, var