code: def _while_key_type_signature(eqn, args_consumed):
      cond_jaxpr = eqn.params["cond_jaxpr"].jaxpr
      cond_nconsts = eqn.params["cond_nconsts"]
      body_jaxpr = eqn.params["body_jaxpr"].jaxpr
      body_nconsts = eqn.params["body_nconsts"]
  
      # TODO(jakevdp): pass args_consumed here?
      cond_signature = get_jaxpr_type_signature(cond_jaxpr)
      body_signature = get_jaxpr_type_signature(body_jaxpr)
  
      # Error if there are sinks among consts.
      if any(
          np.any(s.mask)
          for s in cond_signature.sinks
          if s.idx < cond_nconsts
      ):
          raise KeyReuseError(
              "while_loop cond function leads to key reuse when repeatedly executed: "
              f"{cond_signature=}"
          )
      if any(
          np.any(s.mask)
          for s in body_signature.sinks
          if s.idx < body_nconsts
      ):
          raise KeyReuseError(
              "while_loop body function leads to key reuse when repeatedly executed: "
              f"{body_signature=}"
          )
  
      # carry should only consume keys that are sourced on output.
      body_carry_sinks = {
          s.idx - body_nconsts: s.mask
          for s in body_signature.sinks
          if s.idx >= body_nconsts
      }
      cond_carry_sinks = {
          s.idx - cond_nconsts: s.mask
          for s in cond_signature.sinks
          if s.idx >= cond_nconsts
      }
      carry_sources = {s.idx: s.mask for s in body_signature.sources}
      # TODO(jakevdp): check masks at each index?
      if not (cond_carry_sinks.keys() <= carry_sources.keys()):
          raise KeyReuseError(
              "while_loop cond function leads to key reuse when repeatedly executed: "
              f"{cond_signature=}"
          )
      if not (body_carry_sinks.keys() <= carry_sources.keys()):
          raise KeyReuseError(
              "while_loop body function leads to key reuse when repeatedly executed: "
              f"{body_signature=}"
          )
      if body_carry_sinks.keys() & cond_carry_sinks.keys():
          raise KeyReuseError(
              "while_loop cond and body functions both use the same key: "
              f"{cond_signature=} {body_signature=}"
          )
      return body_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 9 for calling a global function with example usage: _while_key_type_signature(eqn, args_consumed) and returns: body_signature