code: def eval_summary(
      f: Union[
          Callable[..., Any], hk.Transformed, hk.TransformedWithState
      ],
  ) -> Callable[..., Sequence[MethodInvocation]]:
      """Records module method calls performed by ``f``.
  
      >>> f = lambda x: hk.nets.MLP([300, 100, 10])(x)
      >>> x = jnp.ones([8, 28 * 28])
      >>> for i in hk.experimental.eval_summary(f)(x):
      ...   print("mod := {:14} | in := {} out := {}".format(
      ...       i.module_details.module.module_name, i.args_spec[0], i.output_spec))
      mod := mlp            | in := f32[8,784] out := f32[8,10]
      mod := mlp/~/linear_0 | in := f32[8,784] out := f32[8,300]
      mod := mlp/~/linear_1 | in := f32[8,300] out := f32[8,100]
      mod := mlp/~/linear_2 | in := f32[8,100] out := f32[8,10]
  
      Args:
        f: A function or transformed function to trace.
  
      Returns:
        A callable taking the same arguments as the provided function, but returning
        a sequence of :class:`MethodInvocation` instances revealing the methods
        called on each module when applying ``f``.
  
      See Also:
        :func:`tabulate`: Pretty prints a summary of the execution of a function.
      """
      sidechannel = data_structures.ThreadLocalStack()
  
      try:
          f = transform.get_original_fn(f)
      except AttributeError:
          pass
  
      def f_logged(*args, **kwargs):
          used_modules = sidechannel.peek()
          logging_interceptor = functools.partial(
              log_used_modules, used_modules
          )
  
          with hk.intercept_methods(logging_interceptor):
              f(*args, **kwargs)
  
      # We know that we will only evaluate this function once and that inside
      # eval_shape we will re-trace any jitted/pmap-ed code. This allows users to
      # pass in jit/pmap decorated apply functions (e.g. train_step).
      f = make_hk_transform_ignore_jax_transforms(f)
  
      f_orig = hk.transform_with_state(f)
      f_logged = hk.transform_with_state(f_logged)
  
      def init_apply(*args, **kwargs):
          init_rng, apply_rng = jax.random.split(jax.random.PRNGKey(42))
          params, state = f_orig.init(init_rng, *args, **kwargs)
          f_logged.apply(params, state, apply_rng, *args, **kwargs)
  
      def wrapper(*args, **kwargs) -> Sequence[MethodInvocation]:
          used_modules = []
          with sidechannel(used_modules), jax.disable_jit():
              jax.eval_shape(init_apply, *args, **kwargs)
          return used_modules
  
      return wrapper
  
imports: from collections.abc import Mapping, Sequence
  import dataclasses
  import functools
  import pprint
  from typing import Any, Callable, Optional, TypeVar, Union
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import transform
  from haiku._src import utils
  import jax
  import jax.numpy
  import numpy
  import tabulate
label_desc: Use image node_id 5 for calling a global function with example usage: eval_summary(f) and returns: wrapper, used_modules