Add with_path to allow modules to be flattened with key paths.

```
>>> class MyModule(tf.Module):
...   @Property
...   def state_dict(self):
...     return dict(self._flatten(
...         predicate=lambda v: isinstance(v, tf.Variable), with_path=True))

>>> mod = MyModule()
>>> mod.encoder = Encoder()
>>> mod.decoder = mod.encoder
>>> mod.state_dict
{('encoder', 'w'), <tf.Variable: ...>,
 ('decoder', 'w'), <tf.Variable: ...>}
```

h/t tensorflow/community#56

PiperOrigin-RevId: 232908045

Author: tomhennigan

tensorflow/python/module/module.py

@@ -207,7 +207,8 @@ def variables(self):
 
     Returns:
       A sequence of variables for the current module (sorted by attribute
-      name) followed by variables from all submodules recursively (depth first).
+      name) followed by variables from all submodules recursively (breadth
+      first).
     """
     return tuple(self._flatten(predicate=_IS_VARIABLE))
 
@@ -221,7 +222,8 @@ def trainable_variables(self):
 
     Returns:
       A sequence of variables for the current module (sorted by attribute
-      name) followed by variables from all submodules recursively (depth first).
+      name) followed by variables from all submodules recursively (breadth
+      first).
     """
     return tuple(self._flatten(predicate=_IS_TRAINABLE_VARIABLE))
 
@@ -249,7 +251,8 @@ def submodules(self):
   def _flatten(self,
                recursive=True,
                predicate=None,
-               attribute_traversal_key=None):
+               attribute_traversal_key=None,
+               with_path=False):
     """Flattened attribute values in sorted order by attribute name.
 
     Modules are flattened by first walking their attributes in name order.
@@ -267,11 +270,15 @@ def _flatten(self,
     ...
     ...   @property
     ...   def tensors(self):
-    ...     return tuple(self._flatten(predicate=is_tensor))
+    ...     return tuple(self._flatten(predicate=is_tensor, with_path=True))
 
     >>> foo = Foo()
     >>> foo.tensors
-    (<tf.Tensor...'a'>, <tf.Tensor...'b'>, ...'c'>, ...'d'>, ...'e'>)
+    ((('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.
@@ -284,6 +291,10 @@ def _flatten(self,
       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).
 
     Returns:
       Flat generator for leaves of the current module and optionally all
@@ -297,7 +308,7 @@ def _flatten(self,
         recursive=recursive,
         predicate=predicate,
         attribute_traversal_key=attribute_traversal_key,
-        seen=set())
+        with_path=with_path)
 
   @classmethod
   def no_name_scope(cls, method):
@@ -337,8 +348,20 @@ def camel_to_snake(value):
   return _CAMEL_TO_SNAKE_R.sub(r"_\1", value).lower()
 
 
-def _flatten_module(module, recursive, predicate, attribute_traversal_key,
-                    seen):
+# AutoCheckpointable adds object attributes that users will not expect us to
+# include when flattening (these reference dependencies reachable via other
+# object attributes).
+AUTO_CHECKPOINTABLE_ATTRS = ("_unconditional_checkpoint_dependencies",
+                             "_unconditional_dependency_names")
+
+
+def _flatten_module(module,
+                    recursive,
+                    predicate,
+                    attribute_traversal_key,
+                    with_path,
+                    module_path=(),
+                    seen=None):
   """Implementation of `flatten`."""
   if seen is None:
     seen = set([id(module)])
@@ -347,25 +370,37 @@ def _flatten_module(module, recursive, predicate, attribute_traversal_key,
   submodules = []
 
   for key in sorted(module_dict, key=attribute_traversal_key):
-    for leaf in nest.flatten(module_dict[key]):
-      leaf_id = id(leaf)
-      if leaf_id in seen:
-        continue
+    if key in AUTO_CHECKPOINTABLE_ATTRS:
+      continue
+
+    for leaf_path, leaf in nest.flatten_with_tuple_paths(module_dict[key]):
+      leaf_path = (key,) + leaf_path
+
+      # TODO(tomhennigan) Handle cycles for `with_path=True` (e.g. `a.a = a`).
+      if not with_path:
+        leaf_id = id(leaf)
+        if leaf_id in seen:
+          continue
+        seen.add(leaf_id)
 
-      seen.add(leaf_id)
       if predicate(leaf):
-        yield leaf
+        if with_path:
+          yield module_path + leaf_path, leaf
+        else:
+          yield leaf
 
       if recursive and isinstance(leaf, Module):
         # Walk direct properties first then recurse.
-        submodules.append(leaf)
+        submodules.append((module_path + leaf_path, leaf))
 
-  for submodule in submodules:
+  for submodule_path, submodule in submodules:
     subvalues = _flatten_module(
         submodule,
         recursive=recursive,
         predicate=predicate,
         attribute_traversal_key=attribute_traversal_key,
+        with_path=with_path,
+        module_path=submodule_path,
         seen=seen)
 
     for subvalue in subvalues:

tensorflow/python/module/module_test.py

@@ -299,10 +299,10 @@ def with_name_scope(self):
 mk_index_dict = lambda v: dict(enumerate(v))
 
 
-class WalkTest(parameterized.TestCase, test.TestCase):
+class FlattenTest(parameterized.TestCase, test.TestCase):
 
   @parameterized.parameters(lambda v: NamedPair(*v), list, tuple, mk_index_dict)
-  def test_walk(self, container_type):
+  def test_flatten(self, container_type):
     parent = SimpleModule(container_type=container_type)
     child = parent.c
 
@@ -320,6 +320,23 @@ def test_attribute_traversal_key(self):
         mod.variables,
         mod._trainable_variables + mod._non_trainable_variables + [mod._bonus])
 
+  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"]},)
+
 
 class LayerModule(module.Module):