code: def fast_eval_shape(fun, *args, **kwargs):
      """Equivalent to ``eval_shape`` in JAX.
  
      This utility is equivalent to ``eval_shape`` in JAX except that it avoids
      running Haiku functions whose shapes are trivially known. This can avoid some
      Python overheads in JAX which can accumulate for very large models.
  
      Optimizations:
  
      * All parameter/state initialisers replaced with zeros.
      * ``hk.dropout`` replaced with identity.
      * ``jax.random.fold_in`` replaced with identity.
  
      Args:
        fun: The function to trace.
        *args: Positional arguments to ``fun``.
        **kwargs: Keyword arguments to ``fun``.
  
      Returns:
        The shape produced by ``fun`` for the given args/kwargs.
      """
      with base.custom_creator_unsafe(zeros_creator), mock.patch.object(
          basic, "dropout_impl", noop_dropout
      ), mock.patch.object(
          jax.random, "fold_in", lambda key, data: key
      ):
          if base.inside_transform():
              return stateful.eval_shape(fun, *args, **kwargs)
          else:
              return jax.eval_shape(fun, *args, **kwargs)
  
file_path: dm-haiku/haiku/_src/eval_shape.py
imports: from unittest import mock
  from haiku._src import base
  from haiku._src import basic
  from haiku._src import stateful
  import jax
  import jax.numpy
label_desc: Use image node_id 3 for calling a global function with example usage: fast_eval_shape(fun) and returns: stateful, jax