code: def create_variables_eagerly(f: Callable[..., T]) -> Callable[..., T]:
      """Wraps a function and attempts to create variables using eager mode.
  
      Example usage:
  
      >>> model = snt.Sequential([snt.Linear(1) for _ in range(100)])
  
      >>> @tf.function
      ... @snt.distribute.create_variables_eagerly
      ... def f(x):
      ...   return model(x)
  
      >>> _ = f(tf.ones([1, 1]))
  
      On a CPU only machine ``f`` will run ~4x faster (700ms vs. 3s), the benefits
      are more pronounced in a distributed setup since eager variable creation can
      skip a number of checks that are required in graph mode (e.g. checking whether
      the variable has already been created) which end up ping-ponging RPCs.
  
      Args:
        f: Any function.
  
      Returns:
        A function running `f` in a context where variables are created eagerly.
      """
  
      def wrapper(*args, **kwargs):
          with contextlib.ExitStack() as stack:
              # The two hacks below enable a large speedup when initializing large
              # models on TPU pods.
              # TODO(b/141243467) Remove these workarounds.
              stack.enter_context(_eager_initial_values())
              stack.enter_context(
                  tf.variable_creator_scope(_eager_variable_creator)
              )
              return f(*args, **kwargs)
  
      return wrapper
  
file_path: sonnet/sonnet/src/distribute/replicator.py
imports: from typing import Callable, TypeVar
  from absl import logging
  import contextlib
  from sonnet.src import initializers
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: create_variables_eagerly(f) and returns: wrapper, f