chore: import upstream snapshot with attribution

tensorflow/python/module/BUILD

@@ -0,0 +1,46 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "module",
+    srcs = ["module.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/trackable:autotrackable",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "module_test",
+    srcs = ["module_test.py"],
+    deps = [
+        ":module",
+        "//tensorflow/python:extra_py_tests_deps",
+        "//tensorflow/python:tf2",
+        "//tensorflow/python/distribute:ps_values",
+        "//tensorflow/python/distribute:tpu_values",
+        "//tensorflow/python/distribute:values",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:composite_tensor",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/module/module.py

@@ -0,0 +1,467 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Modules encapsulate building stateful components."""
+
+import re
+
+from tensorflow.python import tf2
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variables
+from tensorflow.python.trackable import autotrackable
+from tensorflow.python.util import nest
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("Module")
+class Module(autotrackable.AutoTrackable):
+  """Base neural network module class.
+
+  A module is a named container for `tf.Variable`s, other `tf.Module`s and
+  functions which apply to user input. For example a dense layer in a neural
+  network might be implemented as a `tf.Module`:
+
+  >>> class Dense(tf.Module):
+  ...   def __init__(self, input_dim, output_size, name=None):
+  ...     super().__init__(name=name)
+  ...     self.w = tf.Variable(
+  ...       tf.random.normal([input_dim, output_size]), name='w')
+  ...     self.b = tf.Variable(tf.zeros([output_size]), name='b')
+  ...   def __call__(self, x):
+  ...     y = tf.matmul(x, self.w) + self.b
+  ...     return tf.nn.relu(y)
+
+  You can use the Dense layer as you would expect:
+
+  >>> d = Dense(input_dim=3, output_size=2)
+  >>> d(tf.ones([1, 3]))
+  <tf.Tensor: shape=(1, 2), dtype=float32, numpy=..., dtype=float32)>
+
+
+  By subclassing `tf.Module` instead of `object` any `tf.Variable` or
+  `tf.Module` instances assigned to object properties can be collected using
+  the `variables`, `trainable_variables` or `submodules` property:
+
+  >>> d.variables
+      (<tf.Variable 'b:0' shape=(2,) dtype=float32, numpy=...,
+      dtype=float32)>,
+      <tf.Variable 'w:0' shape=(3, 2) dtype=float32, numpy=..., dtype=float32)>)
+
+
+  Subclasses of `tf.Module` can also take advantage of the `_flatten` method
+  which can be used to implement tracking of any other types.
+
+  All `tf.Module` classes have an associated `tf.name_scope` which can be used
+  to group operations in TensorBoard and create hierarchies for variable names
+  which can help with debugging. We suggest using the name scope when creating
+  nested submodules/parameters or for forward methods whose graph you might want
+  to inspect in TensorBoard. You can enter the name scope explicitly using
+  `with self.name_scope:` or you can annotate methods (apart from `__init__`)
+  with `@tf.Module.with_name_scope`.
+
+  >>> class MLP(tf.Module):
+  ...   def __init__(self, input_size, sizes, name=None):
+  ...     super().__init__(name=name)
+  ...     self.layers = []
+  ...     with self.name_scope:
+  ...       for size in sizes:
+  ...         self.layers.append(Dense(input_dim=input_size, output_size=size))
+  ...         input_size = size
+  ...   @tf.Module.with_name_scope
+  ...   def __call__(self, x):
+  ...     for layer in self.layers:
+  ...       x = layer(x)
+  ...     return x
+
+  >>> module = MLP(input_size=5, sizes=[5, 5])
+  >>> module.variables
+  (<tf.Variable 'mlp/b:0' shape=(5,) dtype=float32, numpy=..., dtype=float32)>,
+  <tf.Variable 'mlp/w:0' shape=(5, 5) dtype=float32, numpy=...,
+     dtype=float32)>,
+  <tf.Variable 'mlp/b:0' shape=(5,) dtype=float32, numpy=..., dtype=float32)>,
+  <tf.Variable 'mlp/w:0' shape=(5, 5) dtype=float32, numpy=...,
+     dtype=float32)>)
+  """
+
+  # AutoTrackable adds object attributes that users will not expect us to
+  # include when flattening (these reference dependencies reachable via other
+  # object attributes).
+  _TF_MODULE_IGNORED_PROPERTIES = frozenset((
+      "_self_unconditional_checkpoint_dependencies",
+      "_self_unconditional_dependency_names"
+  ))
+
+  def __init__(self, name=None):
+    if name is None:
+      name = camel_to_snake(type(self).__name__)
+    else:
+      if not valid_identifier(name):
+        raise ValueError(
+            "%r is not a valid module name. Module names must be valid Python "
+            "identifiers (e.g. a valid class name)." % name)
+
+    self._name = name
+    if tf2.enabled():
+      with ops.name_scope_v2(name) as scope_name:
+        self._name_scope = ops.name_scope_v2(scope_name)
+    else:
+      with ops.name_scope(name, skip_on_eager=False) as scope_name:
+        self._scope_name = scope_name
+
+  @property
+  def name(self):
+    """Returns the name of this module as passed or determined in the ctor.
+
+    NOTE: This is not the same as the `self.name_scope.name` which includes
+    parent module names.
+    """
+    return self._name
+
+  @property
+  def name_scope(self):
+    """Returns a `tf.name_scope` instance for this class."""
+    if tf2.enabled():
+      return self._name_scope
+    else:
+      # In TF1 name_scope is not re-entrant in eager so we cannot memoize it.
+      return ops.name_scope(self._scope_name, skip_on_eager=False)
+
+  @property
+  def variables(self):
+    """Sequence of variables owned by this module and its submodules.
+
+    Note: this method uses reflection to find variables on the current instance
+    and submodules. For performance reasons you may wish to cache the result
+    of calling this method if you don't expect the return value to change.
+
+    Returns:
+      A sequence of variables for the current module (sorted by attribute
+      name) followed by variables from all submodules recursively (breadth
+      first).
+    """
+    return tuple(self._flatten(predicate=_is_variable, expand_composites=True))
+
+  @property
+  def trainable_variables(self):
+    """Sequence of trainable variables owned by this module and its submodules.
+
+    Note: this method uses reflection to find variables on the current instance
+    and submodules. For performance reasons you may wish to cache the result
+    of calling this method if you don't expect the return value to change.
+
+    Returns:
+      A sequence of variables for the current module (sorted by attribute
+      name) followed by variables from all submodules recursively (breadth
+      first).
+    """
+    return tuple(
+        self._flatten(predicate=_is_trainable_variable, expand_composites=True))
+
+  @property
+  def non_trainable_variables(self):
+    """Sequence of non-trainable variables owned by this module and its submodules.
+
+    Note: this method uses reflection to find variables on the current instance
+    and submodules. For performance reasons you may wish to cache the result
+    of calling this method if you don't expect the return value to change.
+
+    Returns:
+      A sequence of variables for the current module (sorted by attribute
+      name) followed by variables from all submodules recursively (breadth
+      first).
+    """
+    return tuple(self._flatten(
+        predicate=_is_non_trainable_variable, expand_composites=True))
+
+  @property
+  def submodules(self):
+    """Sequence of all sub-modules.
+
+    Submodules are modules which are properties of this module, or found as
+    properties of modules which are properties of this module (and so on).
+
+    >>> a = tf.Module()
+    >>> b = tf.Module()
+    >>> c = tf.Module()
+    >>> a.b = b
+    >>> b.c = c
+    >>> list(a.submodules) == [b, c]
+    True
+    >>> list(b.submodules) == [c]
+    True
+    >>> list(c.submodules) == []
+    True
+
+    Returns:
+      A sequence of all submodules.
+    """
+    return tuple(self._flatten(predicate=_is_module))
+
+  def _flatten(self,
+               recursive=True,
+               predicate=None,
+               attribute_traversal_key=None,
+               with_path=False,
+               expand_composites=False):
+    """Flattened attribute values in sorted order by attribute name.
+
+    Modules are flattened by first walking their attributes in name order.
+    Each attribute value is then flattened to find leaf values. If flatten is
+    applied `recursive`ly and if the leaf is a `Module` it will also be
+    flattened to find leaves. Finally every leaf value is optionally tested
+    against the given `predicate` and finally yielded.
+
+    ```
+    class Foo(tf.Module):
+      def __init__(self):
+        super().__init__()
+        self.x = [tf.constant('a'), tf.constant('b')]
+        self.y = {'i': tf.constant('c'), 'j': tf.constant('d')}
+        self.z = tf.constant('e')
+
+      @property
+      def tensors(self):
+        return tuple(self._flatten(predicate=is_tensor, with_path=True))
+
+    foo = Foo()
+    foo.tensors
+    # ==> ((('x', 0),   <tf.Tensor: ...'a'>),
+    #     (('x', 1),   <tf.Tensor: ...'b'>),
+    #     (('y', 'i'), <tf.Tensor: ...'c'>),
+    #     (('y', 'j'), <tf.Tensor: ...'d'>),
+    #     (('z',),     <tf.Tensor: ...'e'>))
+    ```
+
+    `attribute_traversal_key` controls the order object properties are visited.
+    If not set objects are visited in ascending order by name.
+
+    Args:
+      recursive: Whether to recurse into child modules or not.
+      predicate: (Optional) If set then only values matching predicate are
+        yielded. A value of `None` (the default) means no items will be
+        filtered.
+      attribute_traversal_key: (Optional) Method to rekey object attributes
+        before they are sorted. Contract is the same as `key` argument to
+        builtin `sorted` and only applies to object properties.
+      with_path: (Optional) Whether to include the path to the object as well
+        as the object itself. If `with_path` is `True` then leaves will not be
+        de-duplicated (e.g. if the same leaf instance is reachable via multiple
+        modules then it will be yielded multiple times with different paths).
+      expand_composites: If true, then composite tensors are expanded into their
+        component tensors.
+
+    Returns:
+      Flat generator for leaves of the current module and optionally all
+      submodules.
+    """
+    if predicate is None:
+      predicate = lambda _: True
+
+    return _flatten_module(
+        self,
+        recursive=recursive,
+        predicate=predicate,
+        attributes_to_ignore=self._TF_MODULE_IGNORED_PROPERTIES,
+        attribute_traversal_key=attribute_traversal_key,
+        with_path=with_path,
+        expand_composites=expand_composites)
+
+  @classmethod
+  def with_name_scope(cls, method):
+    """Decorator to automatically enter the module name scope.
+
+    >>> class MyModule(tf.Module):
+    ...   @tf.Module.with_name_scope
+    ...   def __call__(self, x):
+    ...     if not hasattr(self, 'w'):
+    ...       self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
+    ...     return tf.matmul(x, self.w)
+
+    Using the above module would produce `tf.Variable`s and `tf.Tensor`s whose
+    names included the module name:
+
+    >>> mod = MyModule()
+    >>> mod(tf.ones([1, 2]))
+    <tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
+    >>> mod.w
+    <tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
+    numpy=..., dtype=float32)>
+
+    Args:
+      method: The method to wrap.
+
+    Returns:
+      The original method wrapped such that it enters the module's name scope.
+    """
+    def method_with_name_scope(self, *args, **kwargs):
+      with self.name_scope:
+        return method(self, *args, **kwargs)
+
+    return tf_decorator.make_decorator(method, method_with_name_scope)
+
+
+def _is_variable(obj):
+  return isinstance(obj, variables.Variable)
+
+
+def _is_trainable_variable(obj):
+  return _is_variable(obj) and getattr(obj, "trainable", False)
+
+
+def _is_non_trainable_variable(obj):
+  return _is_variable(obj) and not getattr(obj, "trainable", False)
+
+
+def _is_module(obj):
+  return isinstance(obj, Module)
+
+_CAMEL_TO_SNAKE_R = re.compile(r"((?<=[a-z0-9])[A-Z]|(?!^)[A-Z](?=[a-z]))")
+_VALID_IDENTIFIER = re.compile(r"^[a-zA-Z_]([a-zA-Z0-9_])*$")
+
+
+def valid_identifier(name):
+  return bool(_VALID_IDENTIFIER.match(name))
+
+
+def camel_to_snake(value):
+  return _CAMEL_TO_SNAKE_R.sub(r"_\1", value).lower()
+
+
+def _flatten_non_variable_composites_with_tuple_path(structure, path_prefix=()):
+  """Flattens composite tensors with tuple path expect variables."""
+  for path, child in nest.flatten_with_tuple_paths(structure):
+    if (isinstance(child, composite_tensor.CompositeTensor) and
+        not _is_variable(child)):
+      # pylint: disable=protected-access
+      spec = child._type_spec
+      yield from _flatten_non_variable_composites_with_tuple_path(
+          spec._to_components(child),
+          path_prefix + path + (spec.value_type.__name__,))
+      # pylint: enable=protected-access
+    else:
+      yield path_prefix + path, child
+
+
+def _flatten_module(module,
+                    recursive,
+                    predicate,
+                    attribute_traversal_key,
+                    attributes_to_ignore,
+                    with_path,
+                    expand_composites,
+                    module_path=(),
+                    seen=None,
+                    recursion_stack=None):
+  """Implementation of `flatten`.
+
+  Args:
+    module: Current module to process.
+    recursive: Whether to recurse into child modules or not.
+    predicate: (Optional) If set then only values matching predicate are
+      yielded. A value of `None` (the default) means no items will be
+      filtered.
+    attribute_traversal_key: (Optional) Method to rekey object attributes
+      before they are sorted. Contract is the same as `key` argument to
+      builtin `sorted` and only applies to object properties.
+    attributes_to_ignore: object attributes to ignored.
+    with_path: (Optional) Whether to include the path to the object as well
+      as the object itself. If `with_path` is `True` then leaves will not be
+      de-duplicated (e.g. if the same leaf instance is reachable via multiple
+      modules then it will be yielded multiple times with different paths).
+    expand_composites: If true, then composite tensors are expanded into their
+      component tensors.
+    module_path: The path to the current module as a tuple.
+    seen: A set containing all leaf IDs seen so far.
+    recursion_stack: A list containing all module IDs associated with the
+      current call stack.
+
+  Yields:
+    Matched leaves with the optional corresponding paths of the current module
+    and optionally all its submodules.
+  """
+  module_id = id(module)
+  if seen is None:
+    seen = set([module_id])
+
+  module_dict = vars(module)
+  submodules = []
+
+  if recursion_stack is None:
+    recursion_stack = []
+
+  # When calling `_flatten_module` with `with_path=False`, the global lookup
+  # table `seen` guarantees the uniqueness of the matched objects.
+  # In the case of `with_path=True`, there might be multiple paths associated
+  # with the same predicate, so we don't stop traversing according to `seen`
+  # to make sure all these paths are returned.
+  # When there are cycles connecting submodules, we break cycles by avoiding
+  # following back edges (links pointing to a node in `recursion_stack`).
+  if module_id in recursion_stack:
+    recursive = False
+
+  for key in sorted(module_dict, key=attribute_traversal_key):
+    if key in attributes_to_ignore:
+      continue
+
+    prop = module_dict[key]
+    try:
+      if expand_composites:
+        leaves = list(_flatten_non_variable_composites_with_tuple_path(prop))
+      else:
+        leaves = nest.flatten_with_tuple_paths(prop)
+    except Exception as cause:  # pylint: disable=broad-except
+      raise ValueError("Error processing property {!r} of {!r}".format(
+          key, prop)) from cause
+
+    for leaf_path, leaf in leaves:
+      leaf_path = (key,) + leaf_path
+
+      if not with_path:
+        leaf_id = id(leaf)
+        if leaf_id in seen:
+          continue
+        seen.add(leaf_id)
+
+      if predicate(leaf):
+        if with_path:
+          yield module_path + leaf_path, leaf
+        else:
+          yield leaf
+
+      if recursive and _is_module(leaf):
+        # Walk direct properties first then recurse.
+        submodules.append((module_path + leaf_path, leaf))
+
+  recursion_stack.append(module_id)
+
+  for submodule_path, submodule in submodules:
+    subvalues = _flatten_module(
+        submodule,
+        recursive=recursive,
+        predicate=predicate,
+        attribute_traversal_key=attribute_traversal_key,
+        attributes_to_ignore=submodule._TF_MODULE_IGNORED_PROPERTIES,  # pylint: disable=protected-access
+        with_path=with_path,
+        expand_composites=expand_composites,
+        module_path=submodule_path,
+        seen=seen,
+        recursion_stack=recursion_stack)
+
+    for subvalue in subvalues:
+      # Predicate is already tested for these values.
+      yield subvalue
+
+  recursion_stack.pop()

tensorflow/python/module/module_test.py

@@ -0,0 +1,613 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for `tf.Module`."""
+
+import abc
+import collections
+import itertools
+import sys
+import unittest
+
+from absl.testing import parameterized
+
+from tensorflow.python import tf2
+from tensorflow.python.distribute import ps_values
+from tensorflow.python.distribute import tpu_values
+from tensorflow.python.distribute import values as distributed_values
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework import type_spec
+from tensorflow.python.module import module
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class TestModuleNaming(test_util.TensorFlowTestCase):
+
+  def test_single_name(self):
+    mod = module.Module(name="simple")
+    self.assertEqual(mod.name, "simple")
+    self.assertEqual(mod.name_scope.name, "simple/")
+
+  def test_construct_in_scope(self):
+    with ops.name_scope("foo", skip_on_eager=False):
+      mod = module.Module(name="bar")
+    self.assertEqual(mod.name, "bar")
+    self.assertEqual(mod.name_scope.name, "foo/bar/")
+
+  def test_enters_name_scope_in_call(self):
+    mod = ReturnsNameScopeModule()
+    for _ in range(3):
+      self.assertEqual(mod(), mod.name_scope.name)
+
+  def test_enters_name_scope_in_other_method(self):
+    mod = ReturnsNameScopeModule()
+    for _ in range(3):
+      self.assertEqual(mod.alternative_forward(), mod.name_scope.name)
+
+  def test_subclassed_module(self):
+    mod = SubclassedReturnsNameScopeModule()
+    for _ in range(3):
+      self.assertEqual(mod.alternative_forward(), mod.name_scope.name)
+      self.assertEqual(mod.alternative_alternative_forward(),
+                       mod.name_scope.name)
+
+  def test_submodule_created_late(self):
+    m = TreeModule()
+    self.assertEqual(m.name, "tree_module")
+    self.assertEqual(m.name_scope.name, "tree_module/")
+    leaf1 = m.new_leaf()
+    self.assertEqual(leaf1.name, "tree_module")
+    self.assertEqual(leaf1.name_scope.name, "tree_module/tree_module/")
+
+  def test_does_not_evaluate_property_methods(self):
+    mod = PropertyThrowsWhenCalledModule()
+    with self.assertRaises(AssertionError):
+      mod.raise_assertion_error  # pylint: disable=pointless-statement
+
+  def test_overridden_name_scope(self):
+    mod = ModuleOverridingNameScope()
+    self.assertEqual(mod(), mod.name_scope.name)
+    self.assertEqual(mod.alternative_forward(), mod.name_scope.name)
+
+  def test_patched_callable(self):
+    with ops.name_scope("foo", skip_on_eager=False):
+      mod = module.Module(name="bar")
+    mod.foo = get_name_scope
+    # `foo` is not a method so we do not re-enter the name scope.
+    self.assertEqual(mod.foo(), "")
+
+  def test_property(self):
+    mod = PropertyModule()
+    mod.some_property = None, None  # None, None for the linter.
+    getter_scope_name, setter_scope_name = mod.some_property
+    self.assertEqual(getter_scope_name, "property_module/")
+    self.assertEqual(setter_scope_name, "property_module/")
+
+  def test_property_no_name_scope(self):
+    mod = PropertyModule()
+    mod.no_name_scope_property = None, None  # None, None for the linter.
+    getter_scope_name, setter_scope_name = mod.no_name_scope_property
+    self.assertEqual(getter_scope_name, "")
+    self.assertEqual(setter_scope_name, "")
+
+  def test_invalid_name(self):
+    msg = ".* is not a valid module name"
+    with self.assertRaisesRegex(ValueError, msg):
+      module.Module(name="$Foo")
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_modules_not_numbered_in_eager(self):
+    if not context.executing_eagerly():
+      self.skipTest("Eager specific")
+
+    mod = RecursiveModule(2)
+    self.assertEqual(mod.name_scope.name, "badger/")
+    self.assertEqual(mod.child.name_scope.name, "badger/badger/")
+
+    mod = RecursiveModule(2)
+    self.assertEqual(mod.name_scope.name, "badger/")
+    self.assertEqual(mod.child.name_scope.name, "badger/badger/")
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_module_numbering_in_graph(self):
+    if context.executing_eagerly():
+      self.skipTest("Graph specific")
+
+    mod = RecursiveModule(2)
+    self.assertEqual(mod.name_scope.name, "badger/")
+    self.assertEqual(mod.child.name_scope.name, "badger/badger/")
+
+    mod = RecursiveModule(2)
+    self.assertEqual(mod.name_scope.name, "badger_1/")
+    self.assertEqual(mod.child.name_scope.name, "badger_1/badger/")
+
+  def test_ctor_error_closes_name_scope(self):
+    with self.assertRaises(ErrorModuleError):
+      # If super constructor is called then a name scope is opened then an error
+      # is thrown. The metaclass should handle this and close the namescope
+      # before re-throwing the exception.
+      ErrorModule(call_super=True)
+
+    self.assertEqual("", get_name_scope())
+
+  def test_ctor_error_handles_ctor_not_opening_name_scope(self):
+    with self.assertRaises(ErrorModuleError):
+      # If super ctor is not called then the name scope isn't opened. We need to
+      # ensure that this doesn't trigger an exception (e.g. the metaclass trying
+      # to __exit__ a non-existent name scope).
+      ErrorModule(call_super=False)
+
+    self.assertEqual("", get_name_scope())
+
+  def test_forward_method_closes_name_scope(self):
+    mod = ErrorModule(call_super=True, raise_in_constructor=False)
+    with self.assertRaises(ErrorModuleError):
+      mod()
+
+    self.assertEqual("", get_name_scope())
+
+  def test_get_attr_doesnt_enter_name_scope(self):
+    scope_names = []
+
+    class GetAttrModule(module.Module):
+
+      def __getattr__(self, name):
+        scope_names.append((name, get_name_scope()))
+        return super().__getattr__(name)
+
+    mod = GetAttrModule()
+    with self.assertRaises(AttributeError):
+      mod.does_not_exist  # pylint: disable=pointless-statement
+    self.assertIn(("does_not_exist", ""), scope_names)
+
+  def test_get_attribute_doesnt_enter_name_scope(self):
+    scope_names = []
+
+    class GetAttributeModule(module.Module):
+
+      def __getattribute__(self, name):
+        scope_names.append((name, get_name_scope()))
+        return super().__getattribute__(name)
+
+    mod = GetAttributeModule()
+    with self.assertRaises(AttributeError):
+      mod.does_not_exist  # pylint: disable=pointless-statement
+    self.assertIn(("does_not_exist", ""), scope_names)
+
+
+class VariableNamingTest(test_util.TensorFlowTestCase):
+
+  def test_variable_names(self):
+    mod = RecursiveModule(3)
+    self.assertEqual(mod.w.name, "badger/mushroom:0")
+    self.assertEqual(mod.child.w.name, "badger/badger/mushroom:0")
+    self.assertEqual(mod.child.child.w.name, "badger/badger/badger/mushroom:0")
+
+
+class NameScopeTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_deprecated_v1
+  def test_not_memoized_in_tf1(self):
+    if tf2.enabled():
+      self.skipTest("Requires TF1")
+
+    mod = module.Module(name="name")
+    name_scope_1 = mod.name_scope
+    name_scope_2 = mod.name_scope
+    self.assertIsNot(name_scope_1, name_scope_2)
+    self.assertEqual(name_scope_1.name, name_scope_2.name)
+
+  def test_memoized_in_tf2(self):
+    if not tf2.enabled():
+      self.skipTest("Requires TF2")
+
+    mod = module.Module(name="name")
+    name_scope_1 = mod.name_scope
+    name_scope_2 = mod.name_scope
+    self.assertIs(name_scope_1, name_scope_2)
+
+
+class VariableTrackingTest(test_util.TensorFlowTestCase):
+
+  def test_variables(self):
+    m = RecursiveModule(3)
+    self.assertEqual(m.variables, (m.w, m.child.w, m.child.child.w))
+    self.assertEqual(m.child.variables, (m.child.w, m.child.child.w))
+    self.assertEqual(m.child.child.variables, (m.child.child.w,))
+
+  def test_trainable_variables(self):
+    m = RecursiveModule(3)
+    self.assertEqual(m.trainable_variables,
+                     (m.w, m.child.w, m.child.child.w))
+    self.assertEqual(m.child.trainable_variables,
+                     (m.child.w, m.child.child.w))
+    self.assertEqual(m.child.child.trainable_variables, (m.child.child.w,))
+
+  def test_trainable_variables_ignores_non_trainable(self):
+    m = RecursiveModule(3, trainable=False)
+    self.assertEqual(len(m.trainable_variables), 0)
+    self.assertEqual(len(m.child.trainable_variables), 0)
+    self.assertEqual(len(m.child.child.trainable_variables), 0)
+
+  def test_supports_distributed_variables(self):
+    mirrored = distributed_values.MirroredVariable(
+        None, [variables.Variable(1.)], variables.VariableAggregation.SUM)
+    tpu = tpu_values.TPUMirroredVariable(
+        strategy=None, values=[variables.Variable(42.)], aggregation=None)
+    aggregating = ps_values.AggregatingVariable(
+        strategy=None, v=variables.Variable(1.), aggregation=None)
+
+    m = module.Module()
+    m.a = mirrored
+    m.b = tpu
+    m.c = aggregating
+    self.assertEqual(m.variables, (mirrored, tpu, aggregating))
+
+  def test_composite_variable(self):
+
+    class Spec(type_spec.TypeSpec):
+
+      value_type = property(lambda self: CompositeVariable)
+
+      def _component_specs(self):
+        pass
+
+      def _serialize(self):
+        pass
+
+      def _to_components(self, value):
+        return value._variables
+
+      def _from_components(self, variable_list):
+        return CompositeVariable(variable_list)
+
+    class CompositeVariable(composite_tensor.CompositeTensor):
+
+      def __init__(self, variable_list):
+        self._variables = variable_list
+
+      @property
+      def _type_spec(self):
+        return Spec()
+
+    m = module.Module()
+    m.a = CompositeVariable([variables.Variable(1.), variables.Variable(2.)])
+    self.assertEqual(list(m.variables), list(m.a._variables))
+
+
+class ModuleTrackingTest(test_util.TensorFlowTestCase):
+
+  def test_submodules(self):
+    m = RecursiveModule(3)
+    self.assertEqual(list(m.submodules), [m.child, m.child.child])
+    self.assertEqual(list(m.child.submodules), [m.child.child])
+    self.assertEqual(list(m.child.child.submodules), [])
+
+  def test_non_ctor_submodule(self):
+    m = TreeModule()
+    leaf1 = m.new_leaf()
+    self.assertEqual(set(m.submodules), {leaf1})
+    leaf2 = m.new_leaf()
+    self.assertEqual(set(m.submodules), {leaf1, leaf2})
+
+
+class ForwardMethodsTest(test_util.TensorFlowTestCase):
+
+  def testFunctionType(self):
+    mod = ModuleWithFunctionAnnotatedCall()
+    self.assertIsInstance(mod.forward, def_function.Function)
+    self.assertIsInstance(mod.forward_ag, def_function.Function)
+
+  def testEntersNameScope_call(self):
+    mod = ModuleWithFunctionAnnotatedCall()
+    self.assertEqual(self.evaluate(mod.forward()),
+                     b"module_with_function_annotated_call/")
+    self.assertEqual(self.evaluate(mod.forward_ag()),
+                     b"module_with_function_annotated_call/")
+
+  def testEntersNameScope_concreteFunction(self):
+    mod = ModuleWithFunctionAnnotatedCall()
+    self.assertEqual(self.evaluate(mod.forward.get_concrete_function()()),
+                     b"module_with_function_annotated_call/")
+    self.assertEqual(self.evaluate(mod.forward_ag.get_concrete_function()()),
+                     b"module_with_function_annotated_call/")
+
+
+class AbcTest(test_util.TensorFlowTestCase):
+
+  def testAbstract(self):
+    msg = "Can't instantiate.*abstract"
+    with self.assertRaisesRegex(TypeError, msg):
+      AbstractModule()  # pylint: disable=abstract-class-instantiated
+
+  def testConcrete(self):
+    mod = ConcreteModule()
+    x, scope_name = mod(2.)
+    self.assertEqual(x, 4.)
+    self.assertEqual(scope_name, "concrete_module/")
+    self.assertEqual(get_name_scope(), "")
+
+
+def get_name_scope():
+  with ops.name_scope("x", skip_on_eager=False) as ns:
+    ns = "/".join(ns.split("/")[:-2])
+    return ns + "/" if ns else ""
+
+
+class ErrorModuleError(Exception):
+  pass
+
+
+class ErrorModule(module.Module):
+
+  def __init__(self, call_super, raise_in_constructor=True):
+    if call_super:
+      super().__init__()
+    if raise_in_constructor:
+      raise ErrorModuleError("Deliberate error!")
+
+  def __call__(self):
+    raise ErrorModuleError("Deliberate error!")
+
+
+class RecursiveModule(module.Module):
+
+  def __init__(self, depth, trainable=True):
+    super().__init__(name="badger")
+    with self.name_scope:
+      self.child = None
+      if depth > 1:
+        self.child = RecursiveModule(depth - 1, trainable=trainable)
+      self.w = variables.Variable(1.0, trainable=trainable, name="mushroom")
+
+
+class AbstractModule(module.Module, metaclass=abc.ABCMeta):
+
+  @abc.abstractmethod
+  def __call__(self, x):
+    pass
+
+
+class ConcreteModule(AbstractModule):
+
+  @module.Module.with_name_scope
+  def __call__(self, x):
+    return x ** 2, get_name_scope()
+
+
+class TreeModule(module.Module):
+
+  def __init__(self, name=None):
+    super().__init__(name=name)
+    self._leaves = []
+
+  @module.Module.with_name_scope
+  def new_leaf(self, name=None):
+    leaf = TreeModule(name=name)
+    self._leaves.append(leaf)
+    return leaf
+
+
+class ReturnsNameScopeModule(module.Module):
+
+  @module.Module.with_name_scope
+  def alternative_forward(self):
+    return get_name_scope()
+
+  @module.Module.with_name_scope
+  def __call__(self):
+    return get_name_scope()
+
+
+class SubclassedReturnsNameScopeModule(ReturnsNameScopeModule):
+
+  @module.Module.with_name_scope
+  def alternative_alternative_forward(self):
+    return get_name_scope()
+
+
+class PropertyThrowsWhenCalledModule(module.Module):
+
+  @property
+  def raise_assertion_error(self):
+    raise AssertionError
+
+
+class ModuleOverridingNameScope(ReturnsNameScopeModule):
+
+  @property
+  def name_scope(self):
+    return ops.name_scope("yolo/", skip_on_eager=False)
+
+
+class ModuleWithFunctionAnnotatedCall(module.Module):
+
+  @def_function.function(autograph=False)
+  @module.Module.with_name_scope
+  def forward(self):
+    return get_name_scope()
+
+  @def_function.function(autograph=True)
+  @module.Module.with_name_scope
+  def forward_ag(self):
+    return get_name_scope()
+
+
+class PropertyModule(module.Module):
+
+  def __init__(self):
+    super().__init__()
+    self._setter_scope_name = None
+
+  @property
+  @module.Module.with_name_scope
+  def some_property(self):
+    getter_scope_name = get_name_scope()
+    return getter_scope_name, self._setter_scope_name
+
+  @some_property.setter
+  @module.Module.with_name_scope
+  def some_property(self, my_property):
+    self._setter_scope_name = get_name_scope()
+
+  @property
+  def no_name_scope_property(self):
+    getter_scope_name = get_name_scope()
+    return getter_scope_name, self._setter_scope_name
+
+  @no_name_scope_property.setter
+  def no_name_scope_property(self, my_property):
+    self._setter_scope_name = get_name_scope()
+
+NamedPair = collections.namedtuple("NamedPair", ("first", "second"))
+mk_index_dict = lambda v: dict(enumerate(v))
+
+
+class FlattenTest(parameterized.TestCase, test_util.TensorFlowTestCase):
+
+  @parameterized.parameters(lambda v: NamedPair(*v), list, tuple, mk_index_dict)
+  def test_flatten(self, container_type):
+    parent = SimpleModule(container_type=container_type)
+    child = parent.c
+
+    self.assertEqual(
+        list(parent._flatten(recursive=False, predicate=is_member)),
+        [parent.a[0], parent.a[1], parent.z])
+
+    self.assertEqual(
+        list(parent._flatten(predicate=is_member)),
+        [parent.a[0], parent.a[1], parent.z, child.a[0], child.a[1], child.z])
+
+  def test_attribute_traversal_key(self):
+    mod = LayerModule()
+    self.assertEqual(
+        mod.variables,
+        mod._trainable_variables + mod._non_trainable_variables + [mod._bonus])
+
+  def test_attributes_to_ignore(self):
+    class DangerousModule(module.Module):
+      _TF_MODULE_IGNORED_PROPERTIES = frozenset(itertools.chain(
+          ("dangerous_submodule", "dangerous_variable"),
+          module.Module._TF_MODULE_IGNORED_PROPERTIES
+      ))
+
+    mod = DangerousModule()
+    mod.dangerous_submodule = module.Module()
+    mod.dangerous_variable = variables.Variable(1.)
+    mod.normal_variable = variables.Variable(2.)
+
+    self.assertEmpty(mod.submodules)
+    self.assertLen(mod.variables, 1)
+    self.assertEqual(mod.variables[0], mod.normal_variable)
+
+  @unittest.skipIf(sys.version_info.major == 3 and sys.version_info.minor == 12,
+                   reason="b/313658911: _TupleWrapper __dict__ attribute error")
+  def test_with_path(self):
+    mod = module.Module()
+    mod.w = variables.Variable(1.)
+    mod.encoder = module.Module()
+    mod.encoder.w = [({"k": mod.w}, {"k": mod.w})]
+    mod.decoder = mod.encoder
+
+    state_dict = dict(
+        mod._flatten(with_path=True, predicate=module._is_variable))
+
+    self.assertEqual(state_dict,
+                     {("w",): mod.w,
+                      ("encoder", "w", 0, 0, "k"): mod.encoder.w[0][0]["k"],
+                      ("encoder", "w", 0, 1, "k"): mod.encoder.w[0][1]["k"],
+                      ("decoder", "w", 0, 0, "k"): mod.decoder.w[0][0]["k"],
+                      ("decoder", "w", 0, 1, "k"): mod.decoder.w[0][1]["k"]},)
+
+  def test_cycles_with_path(self):
+    mod = module.Module()
+    mod.w = variables.Variable(1.)
+    mod.encoder = module.Module()
+    mod.encoder.w = [({"k": mod.w}, {"k": mod.w})]
+    mod.decoder = mod.encoder
+
+    # This introduces two cycles: on mod.encoder.mod and mod.decoder.mod.
+    mod.decoder.mod = mod
+
+    state_dict = dict(
+        mod._flatten(with_path=True, predicate=module._is_variable))
+
+    self.assertEqual(state_dict,
+                     {("w",): mod.w,
+                      ("encoder", "mod", "w"): mod.encoder.mod.w,
+                      ("decoder", "mod", "w"): mod.decoder.mod.w,
+                      ("encoder", "w", 0, 0, "k"): mod.encoder.w[0][0]["k"],
+                      ("encoder", "w", 0, 1, "k"): mod.encoder.w[0][1]["k"],
+                      ("decoder", "w", 0, 0, "k"): mod.decoder.w[0][0]["k"],
+                      ("decoder", "w", 0, 1, "k"): mod.decoder.w[0][1]["k"]},)
+
+  def test_raises_error_with_path(self):
+    non_orderable = object
+
+    m = module.Module()
+    m.layers = {non_orderable(): None, non_orderable(): None}
+    with self.assertRaisesRegex(ValueError,
+                                "Error processing property 'layers'"):
+      m.variables  # pylint: disable=pointless-statement
+
+
+class LayerModule(module.Module):
+
+  def __init__(self):
+    super().__init__()
+    self._trainable_variables = [
+        variables.Variable(1., name="a"),
+        variables.Variable(2., name="b"),
+    ]
+    self._non_trainable_variables = [
+        variables.Variable(3., name="c"),
+        variables.Variable(4., name="d"),
+    ]
+    self._bonus = variables.Variable(5., name="e")
+
+  @property
+  def variables(self):
+    def key_function(name):
+      indexes = {"_trainable_variables": 0, "_non_trainable_variables": 1}
+      return indexes.get(name, 2), name
+
+    return list(
+        self._flatten(
+            predicate=module._is_variable,
+            attribute_traversal_key=key_function))
+
+
+class MemberType:
+  """A simple type to search for."""
+  pass
+
+
+class SimpleModule(module.Module):
+
+  def __init__(self, create_child=True, container_type=list):
+    super().__init__()
+    self.z = MemberType()
+    self.a = container_type([MemberType(), MemberType()])
+    if create_child:
+      self.c = SimpleModule(create_child=False)
+
+is_member = lambda v: isinstance(v, MemberType)
+
+if __name__ == "__main__":
+  test.main()