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, 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 9 for calling a global function with example usage: _while_key_type_signature(eqn, args_consumed) and returns: body_signature