code: def custom_variable_getter(
      getter: Callable[[tf.Variable], Any],
      classes: Optional[Iterable[Type[Any]]] = None,
      instances: Optional[Iterable[Any]] = None,
  ) -> ContextManager[Any]:
      """Applies the given `getter` when getting variables of given `classes`.
  
      If a member is a nested structure containing any variable, `getter` will be
      applied to each variable in the nest.
  
      For example:
      >>> class Times2(snt.Module):
      ...   def __init__(self):
      ...     super(Times2, self).__init__()
      ...     self.v = tf.Variable(2.)
      ...
      ...   def __call__(self, x):
      ...     return x * self.v
  
      >>> x = 42.
      >>> times2 = Times2()
  
      >>> with tf.GradientTape() as tape:
      ...   y = times2(x)
      >>> assert tape.gradient(y, times2.v).numpy() == x
  
      >>> with custom_variable_getter(tf.stop_gradient):
      ...   with tf.GradientTape() as tape:
      ...     y = times2(x)
      >>> assert tape.gradient(y, times2.v) is None
  
      >>> with tf.GradientTape() as tape:
      ...   y = times2(x)
      >>> assert tape.gradient(y, times2.v).numpy() == x
  
      Args:
        getter: A callable to apply to each variable of the class.
        classes: The classes in which the getter is applied. If `None`, defaults to
          `set(o.__class__ for o in instances)`. If `classes and `instances` are
          both `None`, defaults to `[Module]`.
        instances: The instances in which the getter is applied. If `None`, the
          getter will apply in all instances of `classes`.
  
      Returns:
        A context manager in which the custom getter is active.
      """
  
      def wrapped_getter(x):
          x_flat = tree.flatten(x)
          if any(_is_variable(it) for it in x_flat):
              return tree.unflatten_as(
                  x,
                  [
                      getter(it) if _is_variable(it) else it
                      for it in x_flat
                  ],
              )
          else:
              return x
  
      return _custom_getter(
          wrapped_getter, classes=classes, instances=instances
      )
  
file_path: sonnet/sonnet/src/custom_getter.py
imports: from typing import Any, Callable, ContextManager, Iterable, Optional, Type
  import contextlib
  from sonnet.src import base
  import tensorflow
  import tree
label_desc: Use image node_id 3 for calling a global function with example usage: custom_variable_getter(getter, classes, instances) and returns: _custom_getter, tree, x