code: def intercept_methods(interceptor: MethodGetter):
      """Register a new method interceptor.
  
      Method interceptors allow you to (at a distance) intercept method calls to
      modules and modify args/kwargs before calling the underlying method. After the
      underlying method is called you can modify its result before it is passed back
      to the user.
  
      For example you could intercept method calls to :class:`~haiku.BatchNorm` and
      ensure it is always computed in full precision:
  
      >>> def my_interceptor(next_f, args, kwargs, context):
      ...   if (type(context.module) is not hk.BatchNorm
      ...       or context.method_name != "__call__"):
      ...     # We ignore methods other than BatchNorm.__call__.
      ...     return next_f(*args, **kwargs)
      ...
      ...   def cast_if_array(x):
      ...     if isinstance(x, jax.Array):
      ...       x = x.astype(jnp.float32)
      ...     return x
      ...
      ...   args, kwargs = jax.tree_util.tree_map(cast_if_array, (args, kwargs))
      ...   out = next_f(*args, **kwargs)
      ...   return out
  
      We can create and use our module in the usual way, we just need to wrap any
      method calls we want to intercept in the context manager:
  
      >>> mod = hk.BatchNorm(decay_rate=0.9, create_scale=True, create_offset=True)
      >>> x = jnp.ones([], jnp.bfloat16)
      >>> with hk.intercept_methods(my_interceptor):
      ...   out = mod(x, is_training=True)
      >>> assert out.dtype == jnp.float32
  
      Without the interceptor BatchNorm would compute in bf16, however since we
      cast `x` before the underlying method is called we compute in f32.
  
      Args:
        interceptor: A method interceptor.
  
      Returns:
        Context manager under which the interceptor is active.
      """
      return interceptor_stack(interceptor)
  
imports: from collections.abc import Mapping
  import contextlib
  import functools
  import inspect
  import re
  from typing import Any, Callable, ContextManager, NamedTuple, Optional, Protocol, TypeVar
  from haiku._src import base
  from haiku._src import config
  from haiku._src import data_structures
  from haiku._src import utils
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: intercept_methods(interceptor) and returns: interceptor_stack