code: def check_key_reuse(
      fun: Callable[..., Any], /, *args: Any
  ) -> KeyReuseSignature:
      """Function to statically check key reuse."""
      args_flat, in_tree = tree_util.tree_flatten(args)
      in_avals_flat = [core.get_aval(arg) for arg in args_flat]
      wrapped_fun, _ = api_util.flatten_fun_nokwargs(
          lu.wrap_init(fun), in_tree
      )
      jaxpr, _, _ = pe.trace_to_jaxpr_dynamic(
          wrapped_fun, in_avals_flat
      )
      return get_jaxpr_type_signature(jaxpr)
  
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 3 for calling a global function with example usage: check_key_reuse() and returns: get_jaxpr_type_signature