docs: def reuse_variables(method):
    """Wraps an arbitrary method so it does variable sharing.
  
    This decorator creates variables the first time it calls `method`, and reuses
    them for subsequent calls. The object that calls `method` provides a
    `tf.VariableScope`, either as a `variable_scope` attribute or as the return
    value of an `_enter_variable_scope()` method.
  
    The first time the wrapped method is invoked, it enters the caller's
    `tf.VariableScope` with `reuse=False`. On all subsequent calls it enters the
    same variable scope with `reuse=True`.
  
    Variables are created in the context of the `tf.VariableScope` provided by the
    caller object. Ops are created with an additional `tf.name_scope()`, which
    adds a scope for the wrapped method name. For example:
  
    ```python
    class MyClass(object):
  
      def __init__(self, name):
        with tf.variable_scope(None, default_name=name) as variable_scope:
          self.variable_scope = variable_scope
  
      @snt.reuse_variables
      def add_x(self, tensor):
        x = tf.get_variable("x", shape=tensor.get_shape())
        return tensor + x
  
    module = MyClass("my_module_name")
    input_tensor = tf.zeros(shape=(5,))
  
    # This creates the variable "my_module_name/x"
    # and op "my_module_name/add_x/add"
    output = module.add_x(input_tensor)
    ```
  
    For performance when executing eagerly it may be desirable to additionally
    annotate these methods using `defun`, such that they are encapsulated as
    graph functions. This is not recommended if your method returns a variable
    since the output of `defun` would be an op that returned the variable's value
    when evaluated (rather than the variable instance).
  
    ```python
    class FooModule(snt.AbstractModule):
      def _build(self, inputs):
        return complex_math(inputs)
  
      @tfe.defun
      @snt.reuse_variables
      def more_complex_stuff(self, inputs):
        return more_complex_math(inputs)
    ```
  
    Args:
      method: The method to wrap.
  
    Returns:
      The wrapped method.
    """
    initialized_variable_scopes_eager = set()
    initialized_variable_scopes_graph = weakref.WeakKeyDictionary()
  
    # Ensure that the argument passed in is really a method by checking that the
    # first positional argument to it is "self".
    arg_spec = inspect.getargspec(method)
    is_method = arg_spec.args and arg_spec.args[0] == "self"
  
    if not is_method:
      raise TypeError("reuse_variables can only be used with methods.")
  
    @wrapt.decorator
    def eager_test(method, obj, args, kwargs):
      """Validates runtime state in eager mode."""
      # If @reuse_variables is combined with @property, obj is passed in args
      # and method is still unbound at this stage.
      if obj is None:
        obj = args[0]
  
      if tf.executing_eagerly() and not hasattr(obj, "_template"):
        raise ValueError(
            "reuse_variables is not supported in eager mode except in Sonnet "
            "modules.")
  
      return method(*args, **kwargs)
  
    @wrapt.decorator
    def call_method(method, obj, args, kwargs):
      """Calls `method` with a variable scope whose reuse flag is set correctly.
  
      The first time the wrapper is called it creates a
      `(tf.Graph, tf.VariableScope)` key and checks it for membership in
      `initialized_variable_scopes`. The check is `False` if and only if this is
      the first time the wrapper has been called with the key, otherwise the
      check is `True`. The result of this check is used as the `reuse` flag for
      entering the provided variable scope before calling `method`.
  
      Here are two examples of how to use the reuse_variables decorator.
  
      1. Decorate an arbitrary instance method with a `variable_scope` attribute:
  
        ```python
        class Reusable(object):
  
          def __init__(self, name):
            with tf.variable_scope(None, default_name=name) as vs:
              self.variable_scope = vs
  
          @snt.reuse_variables
          def add_a(self, input_tensor):
            a = tf.get_variable("a", shape=input_tensor.get_shape())
            return a + input_tensor
  
        obj = Reusable("reusable")
        x = tf.constant(5.0)
        out1 = obj.add_a(x)
        out2 = obj.add_a(x)
        # out1 == out2
        ```
  
      2. Decorating a snt.AbstractModule instance method:
  
        ```python
        class ReusableModule(snt.AbstractModule):
  
          @snt.reuse_variables
          def add_a(self, input_tensor):
            a = tf.get_variable("a", shape=input_tensor.get_shape())
            return a + input_tensor
  
          # We don't need @snt.reuse_variables here because build is
          wrapped by # `tf.make_template` inside `snt.AbstractModule`.
          def _build(self, input_tensor):
            b = tf.get_variable("b", shape=input_tensor.get_shape())
            return b + self.add_a(input_tensor)
  
        obj = Reusable("reusable")
        x = tf.constant(5.0)
        out1 = obj(x)
        out2 = obj(x)
        # out1 == out2
        ```
  
      Args:
        method: The method to wrap.
        obj: The object instance passed to the wrapped method.
        args: The positional arguments (Tensors) passed to the wrapped method.
        kwargs: The keyword arguments passed to the wrapped method.
  
      Returns:
        Output of the wrapped method.
  
      Raises:
        ValueError: If no variable scope is provided or if `method` is a method
                    and a variable_scope keyword argument is also provided.
      """
  
      # If @reuse_variables is combined with @property, obj is passed in args
      # and method is still unbound at this stage.
      if obj is None:
        obj = args[0]
  
      def default_context_manager(reuse=None):
        variable_scope = obj.variable_scope
        return tf.variable_scope(variable_scope, reuse=reuse)
  
      variable_scope_context_manager = getattr(obj, "_enter_variable_scope",
                                               default_context_manager)
  
      with tf.init_scope():
        # We need `init_scope` incase we're running inside a defun. In that case
        # what we want is information about where the function will be called not
        # where the function is being built.
        graph = tf.get_default_graph()
        will_call_in_eager_context = tf.executing_eagerly()
  
      if will_call_in_eager_context:
        initialized_variable_scopes = initialized_variable_scopes_eager
      else:
        if graph not in initialized_variable_scopes_graph:
          initialized_variable_scopes_graph[graph] = set()
        initialized_variable_scopes = initialized_variable_scopes_graph[graph]
  
      # Temporarily enter the variable scope to capture it
      with variable_scope_context_manager() as tmp_variable_scope:
        variable_scope = tmp_variable_scope
  
      reuse = variable_scope.name in initialized_variable_scopes
  
      # Enter the pure variable scope with reuse correctly set
      with variable_scope_ops._pure_variable_scope(  # pylint:disable=protected-access
          variable_scope, reuse=reuse) as pure_variable_scope:
        current_name_scope = tf.get_default_graph().get_name_scope()
        # Force tf.name_scope to treat current_name_scope as an "absolute" scope
        # so we can re-enter it.
        if current_name_scope and current_name_scope[-1] != "/":
          current_name_scope += "/"
        with tf.name_scope(current_name_scope):
          module_name = pure_variable_scope.name
          method_name = to_snake_case(method.__name__)
          method_name_scope = "{}/{}".format(module_name, method_name)
          with tf.name_scope(method_name_scope) as scope:
            if hasattr(obj, "_capture_variables"):
              with obj._capture_variables():  # pylint: disable=protected-access
                out_ops = method(*args, **kwargs)
            else:
              out_ops = method(*args, **kwargs)
        initialized_variable_scopes.add(pure_variable_scope.name)
        try:
          # If `obj` is a Sonnet module, let it know it's been connected
          # to the TF graph.
          obj._is_connected = True  # pylint: disable=protected-access
          if not tf.executing_eagerly():
            obj._add_connected_subgraph(  # pylint: disable=protected-access
                method, out_ops, scope, args, kwargs)
        except AttributeError:
          pass
      return out_ops
  
    return eager_test(call_method(method))