code: def _pjit_key_type_signature(eqn, args_consumed):
      jaxpr = eqn.params["jaxpr"]
      non_literal_invars = [
          v for v in eqn.invars if not isinstance(v, core.Literal)
      ]
      if len(set(non_literal_invars)) != len(non_literal_invars):
          raise ValueError(f"pjit with duplicate inputs: {eqn.invars=}")
      return get_jaxpr_type_signature(
          jaxpr.jaxpr, consumed_inputs=args_consumed
      )
  
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 5 for calling a global function with example usage: _pjit_key_type_signature(eqn, args_consumed) and returns: get_jaxpr_type_signature