code: def _scan_key_type_signature(eqn, args_consumed):
      jaxpr = eqn.params["jaxpr"].jaxpr
      num_consts = eqn.params["num_consts"]
      num_carry = eqn.params["num_carry"]
      length = eqn.params["length"]
      signature = get_jaxpr_type_signature(jaxpr, args_consumed)
  
      # scan body should not consume key in constants
      if any(
          np.any(s.mask) for s in signature.sinks if s.idx < num_consts
      ):
          raise KeyReuseError(
              f"scan body function leads to key reuse when repeatedly executed: {signature=}"
          )
  
      # scan carry should only consume keys that are sourced on output.
      carry_sinks = {
          s.idx - num_consts: s.mask
          for s in signature.sinks
          if 0 <= s.idx - num_consts < num_carry
      }
      carry_sources = {
          s.idx: s.mask
          for s in signature.sources
          if 0 <= s.idx < num_carry
      }
      if (
          carry_sinks.keys() != carry_sources.keys()
      ):  # TODO(jakevdp): check that masks match
          raise KeyReuseError(
              f"scan body function leads to key reuse when repeatedly executed: {signature=}"
          )
      return signature
  
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 8 for calling a global function with example usage: _scan_key_type_signature(eqn, args_consumed) and returns: signature