chore: import upstream snapshot with attribution

tensorflow/python/trackable/BUILD

@@ -0,0 +1,291 @@
+# Description:
+#   Trackable class and subclass definitions.
+
+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 = "trackable",
+    strict_deps = True,
+    visibility = [
+        "//tensorflow:internal",
+        "//third_party/py/tf_agents:__subpackages__",
+    ],
+    deps = [
+        ":asset",
+        ":autotrackable",
+        ":base",
+        ":base_delegate",
+        ":constants",
+        ":converter",
+        ":data_structures",
+        ":layer_utils",
+        ":python_state",
+        ":resource",
+        ":trackable_init",
+        ":trackable_utils",
+    ],
+)
+
+py_library(
+    name = "trackable_init",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+)
+
+py_library(
+    name = "base",
+    srcs = ["base.py"],
+    strict_deps = True,
+    deps = [
+        ":constants",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:control_flow_ops_gen",
+        "//tensorflow/python/training/saving:saveable_object",
+        "//tensorflow/python/util:tf_contextlib",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "base_test",
+    srcs = ["base_test.py"],
+    deps = [
+        ":base",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:variable_scope",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "constants",
+    srcs = ["constants.py"],
+    strict_deps = True,
+)
+
+py_library(
+    name = "converter",
+    srcs = ["converter.py"],
+    strict_deps = True,
+    deps = [
+        ":base",
+        ":data_structures",
+        "//tensorflow/python/eager/polymorphic_function:saved_model_utils",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:resource_variable_ops",
+    ],
+)
+
+py_library(
+    name = "trackable_utils",
+    srcs = ["trackable_utils.py"],
+    strict_deps = True,
+)
+
+tf_py_strict_test(
+    name = "trackable_utils_test",
+    srcs = ["trackable_utils_test.py"],
+    deps = [
+        ":trackable_utils",
+        "//tensorflow/python/eager:test",
+    ],
+)
+
+py_library(
+    name = "base_delegate",
+    srcs = ["base_delegate.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "base_delegate_test",
+    srcs = ["base_delegate_test.py"],
+    deps = [
+        ":base",
+        ":base_delegate",
+        "//tensorflow/python:extra_py_tests_deps",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/checkpoint:checkpoint_options",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:save",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "asset",
+    srcs = ["asset.py"],
+    strict_deps = True,
+    deps = [
+        ":base",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_conversion_registry",
+        "//tensorflow/python/lib/io:file_io",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/saved_model:path_helpers",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "autotrackable",
+    srcs = ["autotrackable.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":base",
+        ":data_structures",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:tf_export",
+        "@absl_py//absl/logging",
+    ],
+)
+
+tf_py_strict_test(
+    name = "autotrackable_test",
+    srcs = ["autotrackable_test.py"],
+    deps = [
+        ":autotrackable",
+        ":data_structures",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+py_library(
+    name = "resource",
+    srcs = ["resource.py"],
+    strict_deps = True,
+    visibility = [
+        "//intelligence/climate_foundations/vlm/export:__subpackages__",
+        "//tensorflow:internal",
+    ],
+    deps = [
+        ":base",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/util:tf_contextlib",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "resource_test",
+    srcs = ["resource_test.py"],
+    deps = [
+        ":resource",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:wrap_function",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "layer_utils",
+    srcs = ["layer_utils.py"],
+    strict_deps = True,
+    deps = ["//tensorflow/python/util:object_identity"],
+)
+
+py_library(
+    name = "data_structures",
+    srcs = ["data_structures.py"],
+    strict_deps = True,
+    deps = [
+        ":base",
+        ":layer_utils",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:tf_export",
+        "@pypi//wrapt",
+    ],
+)
+
+tf_py_strict_test(
+    name = "data_structures_test",
+    srcs = ["data_structures_test.py"],
+    tags = [
+        "no_oss",  # Keras is not available in OSS test
+        "no_windows",
+        "nomac",
+    ],
+    deps = [
+        ":autotrackable",
+        ":data_structures",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:serialization",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "python_state",
+    srcs = ["python_state.py"],
+    strict_deps = True,
+    deps = [
+        ":base",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_py_strict_test(
+    name = "python_state_test",
+    srcs = ["python_state_test.py"],
+    deps = [
+        ":python_state",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/module",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/trackable/asset.py

@@ -0,0 +1,116 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Asset-type Trackable object."""
+import os
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_conversion_registry
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.saved_model import path_helpers
+from tensorflow.python.trackable import base
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("saved_model.Asset")
+class Asset(base.Trackable):
+  """Represents a file asset to hermetically include in a SavedModel.
+
+  A SavedModel can include arbitrary files, called assets, that are needed
+  for its use. For example a vocabulary file used initialize a lookup table.
+
+  When a trackable object is exported via `tf.saved_model.save()`, all the
+  `Asset`s reachable from it are copied into the SavedModel assets directory.
+  Upon loading, the assets and the serialized functions that depend on them
+  will refer to the correct filepaths inside the SavedModel directory.
+
+  Example:
+
+  ```
+  filename = tf.saved_model.Asset("file.txt")
+
+  @tf.function(input_signature=[])
+  def func():
+    return tf.io.read_file(filename)
+
+  trackable_obj = tf.train.Checkpoint()
+  trackable_obj.func = func
+  trackable_obj.filename = filename
+  tf.saved_model.save(trackable_obj, "/tmp/saved_model")
+
+  # The created SavedModel is hermetic, it does not depend on
+  # the original file and can be moved to another path.
+  tf.io.gfile.remove("file.txt")
+  tf.io.gfile.rename("/tmp/saved_model", "/tmp/new_location")
+
+  reloaded_obj = tf.saved_model.load("/tmp/new_location")
+  print(reloaded_obj.func())
+  ```
+
+  Attributes:
+    asset_path: A path, or a 0-D `tf.string` tensor with path to the asset.
+  """
+
+  def __init__(self, path):
+    """Record the full path to the asset."""
+    if isinstance(path, os.PathLike):
+      path = os.fspath(path)
+    # The init_scope prevents functions from capturing `path` in an
+    # initialization graph, since it is transient and should not end up in a
+    # serialized function body.
+    with ops.init_scope(), ops.device("CPU"):
+      self._path = ops.convert_to_tensor(
+          path, dtype=dtypes.string, name="asset_path")
+
+  @property
+  def asset_path(self):
+    """Fetch the current asset path."""
+    return self._path
+
+  @classmethod
+  def _deserialize_from_proto(cls, object_proto, export_dir, asset_file_def,
+                              **unused_kwargs):
+    proto = object_proto.asset
+    filename = file_io.join(
+        path_helpers.get_assets_dir(export_dir),
+        asset_file_def[proto.asset_file_def_index].filename)
+    asset = cls(filename)
+    if not context.executing_eagerly():
+      ops.add_to_collection(ops.GraphKeys.ASSET_FILEPATHS, asset.asset_path)
+    return asset
+
+  def _add_trackable_child(self, name, value):
+    setattr(self, name, value)
+
+  def _export_to_saved_model_graph(self, tensor_map, **unused_kwargs):
+    # TODO(b/205008097): Instead of mapping 1-1 between trackable asset
+    # and asset in the graph def consider deduping the assets that
+    # point to the same file.
+    asset_path_initializer = array_ops.placeholder(
+        shape=self.asset_path.shape,
+        dtype=dtypes.string,
+        name="asset_path_initializer")
+    asset_variable = resource_variable_ops.ResourceVariable(
+        asset_path_initializer)
+
+    tensor_map[self.asset_path] = asset_variable
+    return [self.asset_path]
+
+
+tensor_conversion_registry.register_tensor_conversion_function(
+    Asset, lambda asset, **kw: ops.convert_to_tensor(asset.asset_path, **kw))

tensorflow/python/trackable/autotrackable.py

@@ -0,0 +1,152 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Dependency tracking for trackable objects."""
+
+import warnings
+
+from absl import logging
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as defun
+from tensorflow.python.trackable import base
+from tensorflow.python.trackable import data_structures
+from tensorflow.python.types import core as core_types
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("__internal__.tracking.AutoTrackable", v1=[])
+class AutoTrackable(base.Trackable):
+  """Manages dependencies on other objects.
+
+  `Trackable` objects may have dependencies: other `Trackable` objects
+  which should be saved if the object declaring the dependency is saved. A
+  correctly saveable program has a dependency graph such that if changing a
+  global variable affects an object (e.g. changes the behavior of any of its
+  methods) then there is a chain of dependencies from the influenced object to
+  the variable.
+
+  Dependency edges have names, and are created implicitly when a
+  `Trackable` object is assigned to an attribute of another
+  `Trackable` object. For example:
+
+  ```
+  obj = Trackable()
+  obj.v = ResourceVariable(0.)
+  ```
+
+  The `Trackable` object `obj` now has a dependency named "v" on a
+  variable.
+
+  `Trackable` objects may specify `Tensor`s to be saved and restored
+  directly (e.g. a `Variable` indicating how to save itself) rather than through
+  dependencies on other objects. See
+  `Trackable._gather_saveables_for_checkpoint` for details.
+  """
+
+  def __setattr__(self, name, value):
+    """Support self.foo = trackable syntax."""
+    try:
+      if getattr(self, name) is value:
+        # Short circuit for `self.$x = self.$x`.
+        return
+    except AttributeError:
+      pass
+
+    if getattr(self, "_self_setattr_tracking", True):
+      value = data_structures.sticky_attribute_assignment(
+          trackable=self, value=value, name=name)
+    super(AutoTrackable, self).__setattr__(name, value)
+
+  def __delattr__(self, name):
+    self._delete_tracking(name)
+    super(AutoTrackable, self).__delattr__(name)
+
+  def _no_dependency(self, value):
+    """Override to allow TrackableBase to disable dependency tracking."""
+    return data_structures.NoDependency(value)
+
+  def _trackable_children(self, save_type=base.SaveType.CHECKPOINT, **kwargs):
+    """Returns all children of a trackable, including functions."""
+    if save_type != base.SaveType.SAVEDMODEL:
+      return super(AutoTrackable, self)._trackable_children(
+          save_type, **kwargs)
+
+    functions = {}
+    try:
+      # We get the attributes, suppressing warnings and exceptions.
+      logging_verbosity = logging.get_verbosity()
+      logging.set_verbosity(logging.FATAL)
+      for attribute_name in dir(self):
+        try:
+          with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            attribute_value = getattr(self, attribute_name, None)
+        except Exception:  # pylint: disable=broad-except
+          # NOTE: If we make the exception catching here less broad, we might
+          # need to revisit `finally` block below.
+          # We really don't want to throw an exception just because some
+          # object's attribute accessor is broken.
+          attribute_value = None
+        if isinstance(attribute_value, (def_function.Function,
+                                        defun.ConcreteFunction)):
+          functions[attribute_name] = attribute_value
+    finally:
+      logging.set_verbosity(logging_verbosity)
+
+    # Trace concrete functions to force side-effects:
+    #   1. populate the cache for functions that have an input_signature
+    #      and have not been called
+    #   2. force side effects of creation of concrete functions, e.g. create
+    #      variables on first run.
+    for fn in functions.values():
+      if isinstance(fn, def_function.Function):
+        fn._list_all_concrete_functions_for_serialization()  # pylint: disable=protected-access
+
+    # Additional dependencies may have been generated during function tracing
+    # (e.g. captured variables). Make sure we return those too.
+    children = {}
+    for name, child in self._checkpoint_dependencies:
+      if isinstance(child, (core_types.PolymorphicFunction,
+                            core_types.ConcreteFunction)):
+        # Skip "tracked" functions for now since there may be objects that
+        # automatically track functions that should not be saved.
+        # TODO(kathywu): remove once `_list_functions_for_serialization` has
+        # been fully deprecated.
+        continue
+
+      if name in functions and child is not functions[name]:
+        raise ValueError(
+            "Can't save object because it has multiple children with the same "
+            f"name. Object: {self}, attribute name: {name}, child 1: "
+            f"{child}, child 2: {functions[name]}")
+
+      children[name] = child
+
+    children.update(functions)
+    return children
+
+  def _delete_tracking(self, name):
+    """Removes the tracking of name."""
+    self._maybe_initialize_trackable()
+    if name in self._unconditional_dependency_names:
+      del self._unconditional_dependency_names[name]
+      for index, (dep_name, _) in enumerate(
+          self._unconditional_checkpoint_dependencies):
+        if dep_name == name:
+          del self._unconditional_checkpoint_dependencies[index]
+          break
+
+  def _add_trackable_child(self, name, value):
+    self.__setattr__(name, value)

tensorflow/python/trackable/autotrackable_test.py

@@ -0,0 +1,145 @@
+# Copyright 2017 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.
+# ==============================================================================
+import os
+
+import numpy as np
+from tensorflow.python.checkpoint import checkpoint as util
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+from tensorflow.python.trackable import autotrackable
+from tensorflow.python.trackable import data_structures
+from tensorflow.python.util import nest
+
+
+class InterfaceTests(test.TestCase):
+
+  def testMultipleAssignment(self):
+    root = autotrackable.AutoTrackable()
+    root.leaf = autotrackable.AutoTrackable()
+    root.leaf = root.leaf
+    duplicate_name_dep = autotrackable.AutoTrackable()
+    with self.assertRaisesRegex(ValueError, "already declared"):
+      root._track_trackable(duplicate_name_dep, name="leaf")
+    # No error; we're overriding __setattr__, so we can't really stop people
+    # from doing this while maintaining backward compatibility.
+    root.leaf = duplicate_name_dep
+    root._track_trackable(duplicate_name_dep, name="leaf", overwrite=True)
+    self.assertIs(duplicate_name_dep, root._lookup_dependency("leaf"))
+    self.assertIs(duplicate_name_dep, root._trackable_children()["leaf"])
+
+  def testRemoveDependency(self):
+    root = autotrackable.AutoTrackable()
+    root.a = autotrackable.AutoTrackable()
+    self.assertEqual(1, len(root._trackable_children()))
+    self.assertEqual(1, len(root._unconditional_checkpoint_dependencies))
+    self.assertIs(root.a, root._trackable_children()["a"])
+    del root.a
+    self.assertFalse(hasattr(root, "a"))
+    self.assertEqual(0, len(root._trackable_children()))
+    self.assertEqual(0, len(root._unconditional_checkpoint_dependencies))
+    root.a = autotrackable.AutoTrackable()
+    self.assertEqual(1, len(root._trackable_children()))
+    self.assertEqual(1, len(root._unconditional_checkpoint_dependencies))
+    self.assertIs(root.a, root._trackable_children()["a"])
+
+  def testListBasic(self):
+    a = autotrackable.AutoTrackable()
+    b = autotrackable.AutoTrackable()
+    a.l = [b]
+    c = autotrackable.AutoTrackable()
+    a.l.append(c)
+    a_deps = util.list_objects(a)
+    self.assertIn(b, a_deps)
+    self.assertIn(c, a_deps)
+    self.assertIn("l", a._trackable_children())
+    direct_a_dep = a._trackable_children()["l"]
+    self.assertIn(b, direct_a_dep)
+    self.assertIn(c, direct_a_dep)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testMutationDirtiesList(self):
+    a = autotrackable.AutoTrackable()
+    b = autotrackable.AutoTrackable()
+    a.l = [b]
+    c = autotrackable.AutoTrackable()
+    a.l.insert(0, c)
+    checkpoint = util.Checkpoint(a=a)
+    with self.assertRaisesRegex(ValueError, "A list element was replaced"):
+      checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt"))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testOutOfBandEditDirtiesList(self):
+    a = autotrackable.AutoTrackable()
+    b = autotrackable.AutoTrackable()
+    held_reference = [b]
+    a.l = held_reference
+    c = autotrackable.AutoTrackable()
+    held_reference.append(c)
+    checkpoint = util.Checkpoint(a=a)
+    with self.assertRaisesRegex(ValueError, "The wrapped list was modified"):
+      checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt"))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testNestedLists(self):
+    a = autotrackable.AutoTrackable()
+    a.l = []
+    b = autotrackable.AutoTrackable()
+    a.l.append([b])
+    c = autotrackable.AutoTrackable()
+    a.l[0].append(c)
+    a_deps = util.list_objects(a)
+    self.assertIn(b, a_deps)
+    self.assertIn(c, a_deps)
+    a.l[0].append(1)
+    d = autotrackable.AutoTrackable()
+    a.l[0].append(d)
+    a_deps = util.list_objects(a)
+    self.assertIn(d, a_deps)
+    self.assertIn(b, a_deps)
+    self.assertIn(c, a_deps)
+    self.assertNotIn(1, a_deps)
+    e = autotrackable.AutoTrackable()
+    f = autotrackable.AutoTrackable()
+    a.l1 = [[], [e]]
+    a.l1[0].append(f)
+    a_deps = util.list_objects(a)
+    self.assertIn(e, a_deps)
+    self.assertIn(f, a_deps)
+    checkpoint = util.Checkpoint(a=a)
+    checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt"))
+    a.l[0].append(data_structures.NoDependency([]))
+    a.l[0][-1].append(5)
+    checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt"))
+    # Dirtying the inner list means the root object is unsaveable.
+    a.l[0][1] = 2
+    with self.assertRaisesRegex(ValueError, "A list element was replaced"):
+      checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt"))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testAssertions(self):
+    a = autotrackable.AutoTrackable()
+    a.l = {"k": [np.zeros([2, 2])]}
+    self.assertAllEqual(nest.flatten({"k": [np.zeros([2, 2])]}),
+                        nest.flatten(a.l))
+    self.assertAllClose({"k": [np.zeros([2, 2])]}, a.l)
+    nest.map_structure(self.assertAllClose, a.l, {"k": [np.zeros([2, 2])]})
+    a.tensors = {"k": [array_ops.ones([2, 2]), array_ops.zeros([3, 3])]}
+    self.assertAllClose({"k": [np.ones([2, 2]), np.zeros([3, 3])]},
+                        self.evaluate(a.tensors))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/trackable/base_delegate.py

@@ -0,0 +1,146 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""A mixin class that delegates another Trackable to be used when saving.
+
+This is intended to be used with wrapper classes that cannot directly proxy the
+wrapped object (e.g. with wrapt.ObjectProxy), because there are inner attributes
+that cannot be exposed.
+
+The Wrapper class itself cannot contain any Trackable children, as only the
+delegated Trackable will be saved to checkpoint and SavedModel.
+
+This class will "disappear" and be replaced with the wrapped inner Trackable
+after a cycle of SavedModel saving and loading, unless the object is registered
+and loaded with Keras.
+"""
+
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("__internal__.tracking.DelegatingTrackableMixin", v1=[])
+class DelegatingTrackableMixin(object):
+  """A mixin that delegates all Trackable methods to another trackable object.
+
+  DO NOT USE THIS UNLESS YOU ARE THE KERAS LOSS SCALE OPTIMIZER.
+
+  This class must be used with multiple inheritance. A class that subclasses
+  Trackable can also subclass this class, which causes all Trackable methods to
+  be delegated to the trackable object passed in the constructor.
+
+  A subclass can use this mixin to appear as if it were the trackable passed to
+  the constructor, from a Checkpoint's perspective. LossScaleOptimizer uses this
+  mixin, so that the checkpoint format for a LossScaleOptimizer is identical to
+  the checkpoint format for a normal optimizer. This allows a model to be saved
+  with a normal Optimizer and restored with a LossScaleOptimizer, or vice versa.
+  The only difference in checkpoint format is that the loss scale is also saved
+  with a LossScaleOptimizer.
+  """
+
+  def __init__(self, trackable_obj):
+    self._trackable = trackable_obj
+
+  # pylint: disable=protected-access
+  @property
+  def _setattr_tracking(self):
+    return self._trackable._setattr_tracking
+
+  @_setattr_tracking.setter
+  def _setattr_tracking(self, value):
+    self._trackable._setattr_tracking = value
+
+  @property
+  def _update_uid(self):
+    return self._trackable._update_uid
+
+  @_update_uid.setter
+  def _update_uid(self, value):
+    self._trackable._update_uid = value
+
+  @property
+  def _unconditional_checkpoint_dependencies(self):
+    return self._trackable._unconditional_checkpoint_dependencies
+
+  @property
+  def _unconditional_dependency_names(self):
+    return self._trackable._unconditional_dependency_names
+
+  @property
+  def _name_based_restores(self):
+    return self._trackable._name_based_restores
+
+  def _maybe_initialize_trackable(self):
+    return self._trackable._maybe_initialize_trackable()
+
+  @property
+  def _object_identifier(self):
+    return self._trackable._object_identifier
+
+  @property
+  def _tracking_metadata(self):
+    return self._trackable._tracking_metadata
+
+  def _no_dependency(self, *args, **kwargs):
+    return self._trackable._no_dependency(*args, **kwargs)
+
+  def _name_based_attribute_restore(self, *args, **kwargs):
+    return self._trackable._name_based_attribute_restore(*args, **kwargs)
+
+  @property
+  def _checkpoint_dependencies(self):
+    return self._trackable._checkpoint_dependencies
+
+  @property
+  def _deferred_dependencies(self):
+    return self._trackable._deferred_dependencies
+
+  def _lookup_dependency(self, *args, **kwargs):
+    return self._trackable._lookup_dependency(*args, **kwargs)
+
+  def _add_variable_with_custom_getter(self, *args, **kwargs):
+    return self._trackable._add_variable_with_custom_getter(*args, **kwargs)
+
+  def _preload_simple_restoration(self, *args, **kwargs):
+    return self._trackable._preload_simple_restoration(*args, **kwargs)
+
+  def _track_trackable(self, *args, **kwargs):  # pylint: disable=redefined-outer-name
+    return self._trackable._track_trackable(*args, **kwargs)
+
+  def _handle_deferred_dependencies(self, name, trackable):  # pylint: disable=redefined-outer-name
+    return self._trackable._handle_deferred_dependencies(name, trackable)
+
+  def _gather_saveables_for_checkpoint(self, *args, **kwargs):
+    return self._trackable._gather_saveables_for_checkpoint(*args, **kwargs)
+
+  def _trackable_children(self, *args, **kwargs):
+    return self._trackable._trackable_children(*args, **kwargs)
+
+  def _deserialization_dependencies(self, *args, **kwargs):
+    return self._trackable._deserialization_dependencies(*args, **kwargs)
+
+  def _export_to_saved_model_graph(self, *args, **kwargs):
+    return self._trackable._export_to_saved_model_graph(*args, **kwargs)
+
+  def _serialize_to_tensors(self, *args, **kwargs):
+    return self._trackable._serialize_to_tensors(*args, **kwargs)
+
+  def _restore_from_tensors(self, *args, **kwargs):
+    return self._trackable._restore_from_tensors(*args, **kwargs)
+
+  def _copy_trackable_to_cpu(self, object_map):
+    self._trackable._copy_trackable_to_cpu(object_map)
+    if self not in object_map:
+      object_map[self] = DelegatingTrackableMixin(object_map[self._trackable])
+  # pylint: enable=protected-access
+

tensorflow/python/trackable/base_delegate_test.py

@@ -0,0 +1,100 @@
+# Copyright 2021 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 base_delegate."""
+import os
+
+from absl.testing import parameterized
+
+from tensorflow.python.checkpoint import checkpoint as util
+from tensorflow.python.checkpoint import checkpoint_options
+from tensorflow.python.eager import test
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import variables as variables_lib
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import save
+from tensorflow.python.trackable import base
+from tensorflow.python.trackable import base_delegate
+
+
+class Inner(base.Trackable):
+
+  def __init__(self, v):
+    self.v = v
+    self._track_trackable(v, "v")
+
+  def _copy_trackable_to_cpu(self, object_map):
+    if self not in object_map:
+      object_map[self] = Inner(self.v)
+    self.v._copy_trackable_to_cpu(object_map)
+
+
+class Wrapper(base_delegate.DelegatingTrackableMixin, base.Trackable):
+
+  def __init__(self, inner):
+    self.inner = inner
+    super(Wrapper, self).__init__(inner)
+
+  @property
+  def v(self):
+    return self.inner.v
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class BaseDelegateTest(parameterized.TestCase, test.TestCase):
+
+  @parameterized.named_parameters(
+      ("_enable_async_ckpt", True),
+      ("_disable_async_ckpt", False)
+  )
+  def test_checkpoint(self, enable_async_ckpt):
+    a = Wrapper(Inner(variables_lib.Variable(15.0)))
+    b = Wrapper(Inner(variables_lib.Variable(-15.0)))
+    self.evaluate([a.v.initializer, b.v.initializer])
+
+    test_dir = self.get_temp_dir()
+    prefix = os.path.join(test_dir, "ckpt")
+    ckpt = util.Checkpoint(a=a, b=b)
+    ckpt_options = checkpoint_options.CheckpointOptions(
+        experimental_enable_async_checkpoint=enable_async_ckpt)
+    prefix_tensor = ckpt.save(prefix, options=ckpt_options)
+
+    self.assertEqual([15, -15], self.evaluate([a.v, b.v]))
+    self.evaluate(a.v.assign(-3))
+    self.evaluate(b.v.assign(12))
+    self.assertEqual([-3, 12], self.evaluate([a.v, b.v]))
+
+    # Test that the model can be saved with the wrapper and loaded without it.
+    ckpt2 = util.Checkpoint(a=a.inner, b=b.inner)
+    if enable_async_ckpt:
+      ckpt.sync()
+    ckpt2.restore(prefix_tensor).assert_consumed().run_restore_ops()
+    self.assertEqual([15, -15], self.evaluate([a.v, b.v]))
+
+  def test_saved_model(self):
+    a = Wrapper(Inner(variables_lib.Variable(-15.0)))
+    self.evaluate([a.v.initializer])
+    self.assertEqual([-15], self.evaluate([a.v]))
+
+    test_dir = self.get_temp_dir()
+    saved_model_path = os.path.join(test_dir, "saved_model")
+    save.save(a, saved_model_path)
+
+    loaded = load.load(saved_model_path)
+    self.evaluate([loaded.v.initializer])
+    self.assertEqual([-15], self.evaluate([loaded.v]))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/trackable/base_test.py

@@ -0,0 +1,81 @@
+# Copyright 2017 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.
+# ==============================================================================
+import os
+from tensorflow.python.checkpoint import checkpoint as util
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.platform import test
+from tensorflow.python.trackable import base
+
+
+class InterfaceTests(test.TestCase):
+
+  def testOverwrite(self):
+    root = base.Trackable()
+    leaf = base.Trackable()
+    root._track_trackable(leaf, name="leaf")
+    (current_name, current_dependency), = root._trackable_children().items()
+    self.assertIs(leaf, current_dependency)
+    self.assertEqual("leaf", current_name)
+    duplicate_name_dep = base.Trackable()
+    with self.assertRaises(ValueError):
+      root._track_trackable(duplicate_name_dep, name="leaf")
+    root._track_trackable(duplicate_name_dep, name="leaf", overwrite=True)
+    (current_name, current_dependency), = root._trackable_children().items()
+    self.assertIs(duplicate_name_dep, current_dependency)
+    self.assertEqual("leaf", current_name)
+
+  def testAddVariableOverwrite(self):
+    root = base.Trackable()
+    a = root._add_variable_with_custom_getter(
+        name="v", shape=[], getter=variable_scope.get_variable)
+    self.assertEqual([root, a], util.list_objects(root))
+    with ops.Graph().as_default():
+      b = root._add_variable_with_custom_getter(
+          name="v", shape=[], overwrite=True,
+          getter=variable_scope.get_variable)
+      self.assertEqual([root, b], util.list_objects(root))
+    with ops.Graph().as_default():
+      with self.assertRaisesRegex(ValueError,
+                                  "already declared as a dependency"):
+        root._add_variable_with_custom_getter(
+            name="v", shape=[], overwrite=False,
+            getter=variable_scope.get_variable)
+
+  def testAssertConsumedWithUnusedPythonState(self):
+    has_config = base.Trackable()
+    has_config.get_config = lambda: {}
+    saved = util.Checkpoint(obj=has_config)
+    save_path = saved.save(os.path.join(self.get_temp_dir(), "ckpt"))
+    restored = util.Checkpoint(obj=base.Trackable())
+    restored.restore(save_path).assert_consumed()
+
+  def testBuggyGetConfig(self):
+
+    class NotSerializable(object):
+      pass
+
+    class GetConfigRaisesError(base.Trackable):
+
+      def get_config(self):
+        return NotSerializable()
+
+    util.Checkpoint(obj=GetConfigRaisesError()).save(
+        os.path.join(self.get_temp_dir(), "ckpt"))
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/trackable/constants.py

@@ -0,0 +1,34 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Constants used in Trackable for checkpointing and serialization."""
+
+import enum
+
+
+# Key where the object graph proto is saved in a TensorBundle
+OBJECT_GRAPH_PROTO_KEY = "_CHECKPOINTABLE_OBJECT_GRAPH"
+
+# A key indicating a variable's value in an object's checkpointed Tensors
+# (Trackable._gather_saveables_for_checkpoint). If this is the only key and
+# the object has no dependencies, then its value may be restored on object
+# creation (avoiding double assignment when executing eagerly).
+VARIABLE_VALUE_KEY = "VARIABLE_VALUE"
+OBJECT_CONFIG_JSON_KEY = "OBJECT_CONFIG_JSON"
+
+
+@enum.unique
+class SaveType(str, enum.Enum):
+  SAVEDMODEL = "savedmodel"
+  CHECKPOINT = "checkpoint"

tensorflow/python/trackable/converter.py

@@ -0,0 +1,37 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Util for converting a Python object to a Trackable."""
+
+
+from tensorflow.python.eager.polymorphic_function import saved_model_utils
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.trackable import base
+from tensorflow.python.trackable import data_structures
+
+
+def convert_to_trackable(obj, parent=None):
+  """Converts `obj` to `Trackable`."""
+  if isinstance(obj, base.Trackable):
+    return obj
+  obj = data_structures.wrap_or_unwrap(obj)
+  if (tensor_util.is_tf_type(obj) and
+      obj.dtype not in (dtypes.variant, dtypes.resource) and
+      not resource_variable_ops.is_resource_variable(obj)):
+    return saved_model_utils.TrackableConstant(obj, parent)
+  if not isinstance(obj, base.Trackable):
+    raise ValueError(f"Cannot convert {obj} to Trackable.")
+  return obj

tensorflow/python/trackable/data_structures_test.py

@@ -0,0 +1,762 @@
+# Copyright 2018 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.
+# ==============================================================================
+import collections
+import copy
+import json
+import os
+import pickle
+
+from absl.testing import parameterized
+
+from tensorflow.python.checkpoint import checkpoint as util
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.module import module
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.trackable import autotrackable
+from tensorflow.python.trackable import data_structures
+from tensorflow.python.util import nest
+from tensorflow.python.util import serialization
+
+
+class ListTests(test.TestCase):
+
+  def testJSONSerialization(self):
+    obj = autotrackable.AutoTrackable()
+    obj.l = [1]
+    json.dumps(obj.l, default=serialization.get_json_type)
+
+  def testNotTrackable(self):
+    class NotTrackable(object):
+      pass
+
+    with self.assertRaises(ValueError):
+      data_structures.List([NotTrackable()])
+
+  def testCallNotImplemented(self):
+    with self.assertRaisesRegex(TypeError, "not callable"):
+      data_structures.List()(1.)  # pylint: disable=not-callable
+
+  def testNoPop(self):
+    with self.assertRaises(AttributeError):
+      data_structures.List().pop()
+
+  def testNesting(self):
+    with context.graph_mode():
+      inner = data_structures.List()
+      outer = data_structures.List([inner])
+
+      class TestDense(module.Module):
+        def __init__(self):
+          self.w = variables.Variable(1.0)
+          self.b = variables.Variable(0.0)
+
+        def __call__(self, inputs):
+          return self.w * inputs + self.b
+
+      inner.append(TestDense())
+      inner[0](array_ops.ones([2, 3]))
+      self.assertLen(outer.variables, 2)
+      self.assertIsInstance(
+          outer.variables[0],
+          resource_variable_ops.ResourceVariable)
+
+  def testNonLayerVariables(self):
+    v = resource_variable_ops.ResourceVariable([1.])
+    l = data_structures.List([v])
+    self.assertTrue(l.trainable)
+    self.assertEqual([], l.layers)
+    self.assertEqual([v], l.variables)
+    self.assertEqual([v], l.trainable_weights)
+    self.assertEqual([], l.non_trainable_variables)
+    l.trainable = False
+    self.assertEqual([v], l.variables)
+    self.assertEqual([], l.trainable_variables)
+    self.assertEqual([v], l.non_trainable_variables)
+    l.trainable = True
+    v2 = resource_variable_ops.ResourceVariable(1., trainable=False)
+    l.append(v2)
+    self.assertEqual([v, v2], l.weights)
+    self.assertEqual([v], l.trainable_weights)
+    self.assertEqual([v2], l.non_trainable_weights)
+
+  def testCopy(self):
+    v1 = resource_variable_ops.ResourceVariable(1.)
+    v2 = resource_variable_ops.ResourceVariable(1.)
+    v3 = resource_variable_ops.ResourceVariable(1.)
+
+    l1 = data_structures.List([v1, v2])
+    l2 = l1.copy()
+    l2.append(v3)
+    self.assertEqual(list(l1), [v1, v2])
+    self.assertEqual(list(l2), [v1, v2, v3])
+
+  def testSlicing(self):
+    v1 = resource_variable_ops.ResourceVariable(1.)
+    v2 = resource_variable_ops.ResourceVariable(1.)
+    v3 = resource_variable_ops.ResourceVariable(1.)
+    v4 = resource_variable_ops.ResourceVariable(1.)
+
+    l = data_structures.List([v1, v2, v3, v4])
+    self.assertEqual(l[1:], [v2, v3, v4])
+    self.assertEqual(l[1:-1], [v2, v3])
+    self.assertEqual(l[:-1], [v1, v2, v3])
+
+  def testHash(self):
+    has_sequences = {data_structures.List(), data_structures.List()}
+    self.assertEqual(2, len(has_sequences))
+    self.assertNotIn(data_structures.List(), has_sequences)
+
+  def testIMul_zero(self):
+    l = data_structures.List([])
+    with self.assertRaisesRegex(ValueError, "List only supports append"):
+      l *= 0
+
+  def testIMul(self):
+    v = resource_variable_ops.ResourceVariable(1.)
+    l = data_structures.List([v])
+    l *= 2
+    self.assertEqual(list(l), [v] * 2)
+
+  def testMul(self):
+    v = resource_variable_ops.ResourceVariable(1.)
+    l = data_structures.List([v, v, v])
+    self.assertEqual(list(l * 2), [v, v, v] * 2)
+
+  def testRMul(self):
+    v = resource_variable_ops.ResourceVariable(1.)
+    l = data_structures.List([v, v, v])
+    self.assertEqual(list(2 * l), [v, v, v] * 2)
+
+
+class ListWrapperTest(test.TestCase):
+
+  IGNORED = ("__new__", "__init__", "__subclasshook__", "__getattribute__")
+
+  def test_overrides_all_list_methods(self):
+    not_overridden = []
+
+    for name in dir(list):
+      if name in ListWrapperTest.IGNORED:
+        continue
+
+      list_method = getattr(list, name)
+
+      if not callable(list_method):
+        continue
+
+      object_method = getattr(object, name, None)
+      if object_method is not None and object_method == list_method:
+        # Skip methods that aren't overridden from object.
+        continue
+
+      if list_method == getattr(data_structures.ListWrapper, name):
+        not_overridden.append(name)
+
+    if not_overridden:
+      self.fail("ListWrapper does not override %s" % (not_overridden))
+
+  def testPickle(self):
+    original = data_structures.ListWrapper([1, 2])
+    serialized = pickle.dumps(original)
+    del original
+    deserialized = pickle.loads(serialized)
+    self.assertEqual([1, 2], deserialized)
+
+  def testSameStructure(self):
+    l = [1]
+    nest.assert_same_structure(l, data_structures.ListWrapper(copy.copy(l)))
+
+  def testMutateWithoutTrackableComponents(self):
+    m = module.Module()
+    m.l = [1, 2]
+    m.l.insert(0, 0)
+    self.assertEqual(m.l, [0, 1, 2])
+    self.assertEqual(m.l._trackable_children(), {})
+
+  def testFunctionCaching(self):
+    @def_function.function
+    def f(list_input):
+      return list_input[0] + constant_op.constant(1.)
+
+    first_trace = f.get_concrete_function([constant_op.constant(2.)])
+    second_trace = f.get_concrete_function(
+        data_structures.ListWrapper([constant_op.constant(3.)]))
+    self.assertIs(first_trace, second_trace)
+
+  def testListWrapperBasic(self):
+    # ListWrapper, unlike List, compares like the built-in list type (since it
+    # is used to automatically replace lists).
+    a = autotrackable.AutoTrackable()
+    b = autotrackable.AutoTrackable()
+    self.assertEqual([a, a],
+                     [a, a])
+    self.assertEqual(data_structures.ListWrapper([a, a]),
+                     data_structures.ListWrapper([a, a]))
+    self.assertEqual([a, a],
+                     data_structures.ListWrapper([a, a]))
+    self.assertEqual(data_structures.ListWrapper([a, a]),
+                     [a, a])
+    self.assertNotEqual([a, a],
+                        [b, a])
+    self.assertNotEqual(data_structures.ListWrapper([a, a]),
+                        data_structures.ListWrapper([b, a]))
+    self.assertNotEqual([a, a],
+                        data_structures.ListWrapper([b, a]))
+    self.assertLess([a], [a, b])
+    self.assertLess(data_structures.ListWrapper([a]),
+                    data_structures.ListWrapper([a, b]))
+    self.assertLessEqual([a], [a, b])
+    self.assertLessEqual(data_structures.ListWrapper([a]),
+                         data_structures.ListWrapper([a, b]))
+    self.assertGreater([a, b], [a])
+    self.assertGreater(data_structures.ListWrapper([a, b]),
+                       data_structures.ListWrapper([a]))
+    self.assertGreaterEqual([a, b], [a])
+    self.assertGreaterEqual(data_structures.ListWrapper([a, b]),
+                            data_structures.ListWrapper([a]))
+    self.assertEqual([a], data_structures.ListWrapper([a]))
+    self.assertEqual([a], list(data_structures.List([a])))
+    self.assertEqual([a, a], data_structures.ListWrapper([a]) + [a])
+    self.assertEqual([a, a], [a] + data_structures.ListWrapper([a]))
+    self.assertIsInstance(data_structures.ListWrapper([a]), list)
+    self.assertEqual(
+        tensor_shape.TensorShape([None, 2]).as_list(),
+        (data_structures.ListWrapper([None])
+         + tensor_shape.TensorShape([2])).as_list())
+
+  def testAcceptsNonTrackableContent(self):
+    l = data_structures.ListWrapper([1, 2, 3])
+    self.assertEqual(l, [1, 2, 3])
+
+  def testWrapperChangesList(self):
+    l = []
+    l_wrapper = data_structures.ListWrapper(l)
+    l_wrapper.append(1)
+    self.assertEqual([1], l)
+
+  def testListChangesWrapper(self):
+    l = []
+    l_wrapper = data_structures.ListWrapper(l)
+    l.append(1)
+    self.assertEqual([1], l_wrapper)
+
+  def testNotHashable(self):
+    with self.assertRaises(TypeError):
+      hash(data_structures.ListWrapper())  # pylint: disable=no-value-for-parameter
+
+  def testDelItem(self):
+    l = data_structures.ListWrapper([1, 2, 3, [4]])
+    del l[0]
+    self.assertEqual(l, [2, 3, [4]])
+    self.assertUnableToSave(l, "Unable to save .*__delitem__")
+
+  def testDelSlice(self):
+    l = data_structures.ListWrapper([1, 2, 3, [4]])
+    del l[2:3]
+    self.assertEqual(l, [1, 2, [4]])
+    self.assertUnableToSave(l, "Unable to save .*__delslice__")
+
+  def testSetSlice_canSaveForNonTrackableItems(self):
+    l = data_structures.ListWrapper([1, 2, 3, 4])
+    l[:] = 2, 8, 9, 0
+    self.assertEqual(l, [2, 8, 9, 0])
+    l._maybe_initialize_trackable()  # pylint: disable=protected-access
+    self.assertEqual(len(l._trackable_children()), 0)  # pylint: disable=protected-access
+
+  def testSetSlice_cannotSaveIfTrackableModified(self):
+    v1 = resource_variable_ops.ResourceVariable(1.)
+    v2 = resource_variable_ops.ResourceVariable(1.)
+    l = data_structures.ListWrapper([1, 2, v1, v2])
+    l[:] = 2, 8, 9, v2
+    self.assertEqual(l, [2, 8, 9, v2])
+    self.assertUnableToSave(l, "Unable to save .*__setslice__")
+
+  def testSetSlice_truncate(self):
+    l = data_structures.ListWrapper([1, 2, 3, 4])
+    l[:] = []
+    self.assertEqual(l, [])
+
+  def testSetSlice_extend(self):
+    l = data_structures.ListWrapper([1, 2, 3, 4])
+    l[2:] = 1, 2, 3, 4
+    self.assertEqual(l, [1, 2, 1, 2, 3, 4])
+
+  def testIMulNegative(self):
+    l = data_structures.ListWrapper([1, 2, 3, [4]])
+    l *= -1
+    self.assertEqual(l, [1, 2, 3, [4]] * -1)
+    self.assertUnableToSave(l, "Unable to save")
+
+  def testIMulPositive(self):
+    v = variables.Variable(1.)
+    l = data_structures.ListWrapper([1, 2, 3, 4, v])
+    self.assertDictEqual({"4": v}, l._trackable_children())
+    root = util.Checkpoint(l=l)
+    prefix = os.path.join(self.get_temp_dir(), "ckpt")
+    path = root.save(prefix)
+    v.assign(5.)
+    l *= 2
+    self.assertEqual(l, [1, 2, 3, 4, v, 1, 2, 3, 4, v])
+    self.assertDictEqual({"4": v, "9": v}, l._trackable_children())
+    root.restore(path)
+    self.assertAllClose(1., v.numpy())
+
+  def testSort(self):
+    l = data_structures.ListWrapper([[1], [2], [3], [4]])
+    l.sort()
+    self.assertAllEqual(l, [[1], [2], [3], [4]])
+    # Regardless of being a no-op for the input list, we still refuse to save.
+    # This is intentional since otherwise we would end up with a hard to debug
+    # case for users (e.g. sometimes sort on a ListWrapper is trackable and
+    # other times it is not).
+    self.assertUnableToSave(l, "Unable to save .*sort")
+
+  def assertUnableToSave(self, l, msg):
+    l._maybe_initialize_trackable()  # pylint: disable=protected-access
+    with self.assertRaisesRegex(ValueError, msg):
+      return l._trackable_children()  # pylint: disable=protected-access
+
+
+class MappingTests(test.TestCase):
+
+  def testJSONSerialization(self):
+    obj = autotrackable.AutoTrackable()
+    obj.d = {"a": 2}
+    json.dumps(obj.d, default=serialization.get_json_type)
+
+  def testNoOverwrite(self):
+    mapping = data_structures.Mapping()
+    original = data_structures.List()
+    mapping["a"] = original
+    with self.assertRaises(ValueError):
+      mapping["a"] = data_structures.List()
+    self.assertIs(original, mapping["a"])
+    with self.assertRaises(AttributeError):
+      del mapping["a"]  # pylint: disable=unsupported-delete-operation
+    mapping.update(b=data_structures.Mapping())
+    with self.assertRaises(ValueError):
+      mapping.update({"b": data_structures.Mapping()})
+
+  def testNonStringKeys(self):
+    mapping = data_structures.Mapping()
+    with self.assertRaises(TypeError):
+      mapping[1] = data_structures.List()
+
+  def testHashing(self):
+    has_mappings = set([data_structures.Mapping(),
+                        data_structures.Mapping()])
+    self.assertEqual(2, len(has_mappings))
+    self.assertNotIn(data_structures.Mapping(), has_mappings)
+    # In contrast to Mapping, dict wrappers are not hashable
+    a = autotrackable.AutoTrackable()
+    a.d = {}
+    self.assertEqual({}, a.d)
+    self.assertFalse({} != a.d)  # pylint: disable=g-explicit-bool-comparison
+    self.assertNotEqual({1: 2}, a.d)
+    with self.assertRaisesRegex(TypeError, "unhashable"):
+      set([a.d])
+
+  def testListShallowCopy(self):
+    root = autotrackable.AutoTrackable()
+    orig_list = [[1.]]
+    root.a = orig_list
+    copied = copy.copy(root.a)
+    self.assertAllEqual([[1.]], copied)
+    self.assertIsNot(root.a, copied)
+    self.assertIs(root.a[0], copied[0])
+
+    # Dirtiness should be inherited
+    util.list_objects(root.a)
+    orig_list.append(1.)
+    with self.assertRaises(ValueError):
+      util.list_objects(root.a)
+    with self.assertRaises(ValueError):
+      util.list_objects(copy.copy(root.a))
+
+  def testListDeepCopy(self):
+    root = autotrackable.AutoTrackable()
+    orig_list = [[1.]]
+    root.a = orig_list
+    copied = copy.deepcopy(root.a)
+    self.assertAllEqual([[1.]], copied)
+    self.assertIsNot(root.a, copied)
+    self.assertIsNot(root.a[0], copied[0])
+
+    # Dirtiness should be inherited
+    util.list_objects(root.a)
+    orig_list.append(1.)
+    with self.assertRaises(ValueError):
+      util.list_objects(root.a)
+    with self.assertRaises(ValueError):
+      util.list_objects(copy.deepcopy(root.a))
+
+  def testDictShallowCopy(self):
+    root = autotrackable.AutoTrackable()
+    orig_dict = {"a": [1.]}
+    root.a = orig_dict
+    copied = copy.copy(root.a)
+    self.assertAllEqual([1.], copied["a"])
+    self.assertIsNot(root.a, copied)
+    self.assertIs(root.a["a"], copied["a"])
+
+    copied = root.a.copy()
+    self.assertAllEqual([1.], copied["a"])
+    self.assertIsNot(root.a, copied)
+    self.assertIs(root.a["a"], copied["a"])
+
+    # Dirtiness should be inherited
+    util.list_objects(root.a)
+    orig_dict["b"] = []
+    with self.assertRaises(ValueError):
+      util.list_objects(root.a)
+    with self.assertRaises(ValueError):
+      util.list_objects(copy.copy(root.a))
+
+  def testDictDeepCopy(self):
+    root = autotrackable.AutoTrackable()
+    orig_dict = {"a": [1.]}
+    root.a = orig_dict
+    copied = copy.deepcopy(root.a)
+    self.assertAllEqual([1.], copied["a"])
+    self.assertIsNot(root.a, copied)
+    self.assertIsNot(root.a["a"], copied["a"])
+
+    # Dirtiness should be inherited
+    util.list_objects(root.a)
+    orig_dict["b"] = []
+    with self.assertRaises(ValueError):
+      util.list_objects(root.a)
+    with self.assertRaises(ValueError):
+      util.list_objects(copy.deepcopy(root.a))
+
+  def testShallowCopyTrackable(self):
+    original = autotrackable.AutoTrackable()
+    original_sub = autotrackable.AutoTrackable()
+    original.a = [[1.]]
+    original.b = {"a": original_sub}
+    shallow_copied = copy.copy(original)
+    self.assertIs(original_sub, shallow_copied.b["a"])
+    self.assertIsNot(original, shallow_copied)
+    self.assertEqual([[1.]], shallow_copied.a)
+    shallow_deps = util.list_objects(shallow_copied)
+    self.assertIn(shallow_copied.a, shallow_deps)
+    self.assertIn(shallow_copied.b, shallow_deps)
+    self.assertIn(shallow_copied.b["a"], shallow_deps)
+
+  def testDeepCopyTrackable(self):
+    original = autotrackable.AutoTrackable()
+    original_sub = autotrackable.AutoTrackable()
+    original.a = [[1.]]
+    original.b = {"a": original_sub}
+    self.assertIsInstance(original.b, dict)
+    deep_copied = copy.deepcopy(original)
+    self.assertIsInstance(deep_copied.b, dict)
+    self.assertIsNot(original, deep_copied)
+    self.assertIsNot(original_sub, deep_copied.b["a"])
+    self.assertEqual([[1.]], deep_copied.a)
+    self.assertIsInstance(deep_copied.b["a"], autotrackable.AutoTrackable)
+    deps = util.list_objects(deep_copied)
+    self.assertIn(deep_copied.a, deps)
+    self.assertIn(deep_copied.b, deps)
+    self.assertIn(deep_copied.b["a"], deps)
+    self.assertNotIn(original_sub, deps)
+
+  def testConstructableFromSequence(self):
+    result = data_structures._DictWrapper([(1, 2), (3, 4)])
+    self.assertIsInstance(result, dict)
+    self.assertEqual({1: 2, 3: 4}, result)
+
+  def testPickle(self):
+    original = data_structures._DictWrapper(dict(a=1, b=2))
+    serialized = pickle.dumps(original)
+    del original
+    deserialized = pickle.loads(serialized)
+    self.assertEqual(dict(a=1, b=2), deserialized)
+
+  def testListAddOrder(self):
+    self.assertEqual([1., 2.],
+                     data_structures.ListWrapper([1.])
+                     + data_structures.ListWrapper([2.]))
+    self.assertEqual([1., 2.],
+                     data_structures.ListWrapper([1.])
+                     + [2.])
+    self.assertEqual([1., 2.],
+                     [1.]
+                     + data_structures.ListWrapper([2.]))
+
+  def testSameStructure(self):
+    d = {1: "a"}
+    nest.assert_same_structure(d, data_structures._DictWrapper(d.copy()))
+
+  def testFunctionCaching(self):
+    @def_function.function
+    def f(dict_input):
+      return dict_input["x"] + constant_op.constant(1.)
+
+    first_trace = f.get_concrete_function({"x": constant_op.constant(2.)})
+    second_trace = f.get_concrete_function(
+        data_structures._DictWrapper({"x": constant_op.constant(3.)}))
+    self.assertIs(first_trace, second_trace)
+
+
+class TupleTests(test.TestCase, parameterized.TestCase):
+
+  def testJSONSerialization(self):
+    obj = autotrackable.AutoTrackable()
+    obj.l = (1,)
+    json.dumps(obj.l, default=serialization.get_json_type)
+
+  def testNonLayerVariables(self):
+    v = resource_variable_ops.ResourceVariable([1.])
+    l = data_structures._TupleWrapper((v,))
+    self.assertEqual([], l.layers)
+    self.assertEqual([v], l.variables)
+    self.assertEqual([v], l.trainable_weights)
+    self.assertEqual([], l.non_trainable_variables)
+
+  def testCopy(self):
+    v1 = resource_variable_ops.ResourceVariable(1.)
+    v2 = resource_variable_ops.ResourceVariable(1.)
+
+    l1 = data_structures._TupleWrapper((v1, v2))
+    l2 = copy.copy(l1)
+    self.assertEqual(l1, (v1, v2))
+    self.assertEqual(l2, (v1, v2))
+    self.assertIs(l1[0], l2[0])
+    l2_deep = copy.deepcopy(l1)
+    self.assertIsNot(l1[0], l2_deep[0])
+    with self.assertRaises(AttributeError):
+      l2.append(v1)
+
+  def testSlicing(self):
+    v1 = resource_variable_ops.ResourceVariable(1.)
+    v2 = resource_variable_ops.ResourceVariable(1.)
+    v3 = resource_variable_ops.ResourceVariable(1.)
+    v4 = resource_variable_ops.ResourceVariable(1.)
+
+    l = data_structures._TupleWrapper((v1, v2, v3, v4))
+    self.assertEqual(l[1:], (v2, v3, v4))
+    self.assertEqual(l[1:-1], (v2, v3))
+    self.assertEqual(l[:-1], (v1, v2, v3))
+
+  def testHash(self):
+    has_sequences = set([data_structures._TupleWrapper(),
+                         data_structures._TupleWrapper()])
+    self.assertLen(has_sequences, 1)
+    self.assertIn(data_structures._TupleWrapper(), has_sequences)
+
+  def testIMul_zero(self):
+    l = data_structures._TupleWrapper((1,))
+    l *= 0
+    self.assertEqual((), l)
+
+  def testIMul(self):
+    # Note: tuple behavior differs from list behavior. Lists are mutated by
+    # imul/iadd, tuples assign a new object to the left hand side of the
+    # expression.
+    v = resource_variable_ops.ResourceVariable(1.)
+    l = data_structures._TupleWrapper((v,))
+    original = l
+    l *= 2
+    self.assertEqual(l, (v,) * 2)
+    self.assertNotEqual(original, (v,) * 2)
+
+  def testIAdd(self):
+    v = resource_variable_ops.ResourceVariable(1.)
+    l = data_structures._TupleWrapper((v,))
+    original = l
+    l += (1,)
+    self.assertEqual(l, (v, 1))
+    self.assertNotEqual(original, (v, 1))
+    self.assertEqual(original, (v,))
+
+  def testMul(self):
+    v = resource_variable_ops.ResourceVariable(1.)
+    l = data_structures._TupleWrapper((v, v, v))
+    self.assertEqual(l * 2, (v, v, v) * 2)
+
+  def testRMul(self):
+    v = resource_variable_ops.ResourceVariable(1.)
+    l = data_structures._TupleWrapper((v, v, v))
+    self.assertEqual(2 * l, (v, v, v) * 2)
+
+  def testPickle(self):
+    original = data_structures._TupleWrapper((1, 2))
+    serialized = pickle.dumps(original)
+    del original
+    deserialized = pickle.loads(serialized)
+    self.assertEqual((1, 2), deserialized)
+
+  def testNamedTuple(self):
+    named = collections.namedtuple("Named", ("x", "y"))
+    v = variables.Variable(2)
+    nt = named(x=v, y=2)
+    m = module.Module()
+    m.nt = nt
+    self.assertIs(v, m.nt.x)
+    self.assertIs(v, m.nt[0])
+    self.assertIs(
+        v, m._trackable_children()["nt"]._trackable_children()["x"])
+    self.assertEqual(2, m.nt.y)
+
+  def testNamedTupleConflictingAttributes(self):
+    named = collections.namedtuple("Named", ("x", "weights"))
+    v = variables.Variable(2)
+    nt = named(x=v, weights=3)
+    m = module.Module()
+    m.nt = nt
+    self.assertEqual(3, m.nt.weights)
+
+  def testNamedSubclassing(self):
+    named = collections.namedtuple("Named", ("x", "y"))
+    v = variables.Variable(2)
+
+    class NamedSubclass(named):
+
+      def __new__(cls, x, y):
+        del y  # unused
+        return super(NamedSubclass, cls).__new__(cls, x, 3)
+
+      @property
+      def summed(self):
+        return self.x + self.y
+
+    nt = NamedSubclass(x=v, y=2)
+    m = module.Module()
+    m.nt = nt
+    self.assertEqual(3, m.nt.y)
+    self.assertIs(v, m.nt.x)
+    self.assertIn(v,
+                  m._trackable_children()["nt"]._trackable_children().values())
+    self.assertIn("x", m.nt._trackable_children())
+    self.assertIn("0", m.nt._trackable_children())
+    self.assertEqual(5, self.evaluate(m.nt.summed))
+
+  def testUnnamedSubclassing(self):
+    v = variables.Variable(2)
+
+    class UnnamedSubclass(tuple):
+
+      @property
+      def summed(self):
+        return self[0] + self[1]
+
+    unt = UnnamedSubclass([v, 2])
+    m = module.Module()
+    m.unt = unt
+    self.assertIn("0", m.unt._trackable_children())
+    self.assertLen(m.unt._trackable_children(), 1)
+    self.assertEqual(4, self.evaluate(m.unt.summed))
+    nest.assert_same_structure(
+        [m.unt], nest.map_structure(lambda x: x, [m.unt]))
+
+  def testNamedtupleSubclassWithCustomNew(self):
+    class SubclassWithDifferentArgs(collections.namedtuple("A", ["x"])):
+
+      def __new__(cls):
+        return super(SubclassWithDifferentArgs, cls).__new__(cls, [])
+
+    nt = SubclassWithDifferentArgs()
+    m = module.Module()
+    m.nt = nt
+    m.nt.x.append(variables.Variable(1.))
+    prefix = os.path.join(self.get_temp_dir(), "ckpt")
+    ckpt = util.Checkpoint(m=m)
+    with self.assertRaises(ValueError):
+      ckpt.save(prefix)
+
+  def testSameStructure(self):
+    t = (variables.Variable(1.),)
+    m = module.Module()
+    m.t = t
+    nest.assert_same_structure(t, m.t)
+    nest.assert_same_structure(m.t, t)
+
+    nt_type = collections.namedtuple("nt", ["x", "y"])
+    nt = nt_type(x=1, y=2)
+    m.nt = nt
+    nest.assert_same_structure(m.nt, nt)
+    with self.assertRaises(TypeError):  # pylint: disable=g-error-prone-assert-raises
+      nest.assert_same_structure(m.nt, m.t)
+
+  def testFlatten(self):
+    t = data_structures._TupleWrapper((1, data_structures._TupleWrapper((2,))))
+    self.assertEqual([1, 2], nest.flatten(t))
+    self.assertEqual(
+        nest.flatten_with_tuple_paths((1, (2,))),
+        nest.flatten_with_tuple_paths(t))
+    self.assertEqual((3, (4,)),
+                     nest.pack_sequence_as(t, [3, 4]))
+    nt_type = collections.namedtuple("nt", ["x", "y"])
+    nt = nt_type(1., 2.)
+    wrapped_nt = data_structures._TupleWrapper(nt)
+    self.assertEqual(
+        nest.flatten_with_tuple_paths(nt),
+        nest.flatten_with_tuple_paths(wrapped_nt))
+    self.assertEqual((3, 4,),
+                     nest.pack_sequence_as(wrapped_nt, [3, 4]))
+    self.assertEqual(3, nest.pack_sequence_as(wrapped_nt, [3, 4]).x)
+
+  def testFunctionCaching(self):
+    @def_function.function
+    def f(tuple_input):
+      return tuple_input[0] + constant_op.constant(1.)
+
+    first_trace = f.get_concrete_function((constant_op.constant(2.),))
+    second_trace = f.get_concrete_function(
+        data_structures._TupleWrapper((constant_op.constant(3.),)))
+    self.assertIs(first_trace, second_trace)
+
+  def testPythonMapImpl(self):
+    t = data_structures._TupleWrapper((1, data_structures._TupleWrapper((2,))))
+    self.assertEqual(
+        (4, (5,)),
+        nest.map_structure_up_to((None, (None,)), lambda x: x + 3, t,
+                                 check_types=True))
+    nest.assert_shallow_structure((None, None), t)
+
+  def testDatasetMap(self):
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        constant_op.constant([1, 2, 3]))
+    dataset = dataset.map(lambda x: data_structures._TupleWrapper((x,)))
+    for index, element in enumerate(dataset):
+      self.assertEqual((index + 1,), self.evaluate(element))
+
+  def testDatasetMapNamed(self):
+    nt_type = collections.namedtuple("A", ["x"])
+    dataset = dataset_ops.Dataset.from_tensor_slices(
+        constant_op.constant([1, 2, 3]))
+    dataset = dataset.map(lambda x: data_structures._TupleWrapper(nt_type(x)))
+    for index, element in enumerate(dataset):
+      self.assertEqual((index + 1,), self.evaluate(element))
+
+  def testLoopAssignedModule(self):
+    m = module.Module()
+    m.s = (m,)
+    self.assertLen(m._trackable_children(), 1)
+    self.assertIn("s", m._trackable_children())
+    self.assertIs(m.s, m._trackable_children()["s"])
+    self.assertEqual((), m.trainable_variables)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/trackable/layer_utils.py

@@ -0,0 +1,141 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Utilities related to layer/model functionality."""
+
+# TODO(b/110718070): Move these functions back to tensorflow/python/keras/utils
+# once __init__ files no longer require all of tf.keras to be imported together.
+
+import collections
+import weakref
+
+from tensorflow.python.util import object_identity
+
+try:
+  # typing module is only used for comment type annotations.
+  import typing  # pylint: disable=g-import-not-at-top, unused-import
+except ImportError:
+  pass
+
+
+def is_layer(obj):
+  """Implicit check for Layer-like objects."""
+  # TODO(b/110718070): Replace with isinstance(obj, base_layer.Layer).
+  return hasattr(obj, "_is_layer") and not isinstance(obj, type)
+
+
+def has_weights(obj):
+  """Implicit check for Layer-like objects."""
+  # TODO(b/110718070): Replace with isinstance(obj, base_layer.Layer).
+  has_weight = (hasattr(type(obj), "trainable_weights")
+                and hasattr(type(obj), "non_trainable_weights"))
+
+  return has_weight and not isinstance(obj, type)
+
+
+class MutationSentinel(object):
+  """Container for tracking whether a property is in a cached state."""
+  _in_cached_state = False
+
+  def mark_as(self, value):  # type: (MutationSentinel, bool) -> bool
+    may_affect_upstream = (value != self._in_cached_state)
+    self._in_cached_state = value
+    return may_affect_upstream
+
+  @property
+  def in_cached_state(self):
+    return self._in_cached_state
+
+
+class AttributeSentinel(object):
+  """Container for managing attribute cache state within a Layer.
+
+  The cache can be invalidated either on an individual basis (for instance when
+  an attribute is mutated) or a layer-wide basis (such as when a new dependency
+  is added).
+  """
+
+  def __init__(self, always_propagate=False):
+    self._parents = weakref.WeakSet()
+    self.attributes = collections.defaultdict(MutationSentinel)
+
+    # The trackable data structure containers are simple pass throughs. They
+    # don't know or care about particular attributes. As a result, they will
+    # consider themselves to be in a cached state, so it's up to the Layer
+    # which contains them to terminate propagation.
+    self.always_propagate = always_propagate
+
+  def __repr__(self):
+    return "{}\n  {}".format(
+        super(AttributeSentinel, self).__repr__(),
+        {k: v.in_cached_state for k, v in self.attributes.items()})
+
+  def add_parent(self, node):
+    # type: (AttributeSentinel, AttributeSentinel) -> None
+
+    # Properly tracking removal is quite challenging; however since this is only
+    # used to invalidate a cache it's alright to be overly conservative. We need
+    # to invalidate the cache of `node` (since it has implicitly gained a child)
+    # but we don't need to invalidate self since attributes should not depend on
+    # parent Layers.
+    self._parents.add(node)
+    node.invalidate_all()
+
+  def get(self, key):
+    # type: (AttributeSentinel, str) -> bool
+    return self.attributes[key].in_cached_state
+
+  def _set(self, key, value):
+    # type: (AttributeSentinel, str, bool) -> None
+    may_affect_upstream = self.attributes[key].mark_as(value)
+    if may_affect_upstream or self.always_propagate:
+      for node in self._parents:  # type: AttributeSentinel
+        node.invalidate(key)
+
+  def mark_cached(self, key):
+    # type: (AttributeSentinel, str) -> None
+    self._set(key, True)
+
+  def invalidate(self, key):
+    # type: (AttributeSentinel, str) -> None
+    self._set(key, False)
+
+  def invalidate_all(self):
+    # Parents may have different keys than their children, so we locally
+    # invalidate but use the `invalidate_all` method of parents.
+    for key in self.attributes.keys():
+      self.attributes[key].mark_as(False)
+
+    for node in self._parents:
+      node.invalidate_all()
+
+
+def filter_empty_layer_containers(layer_list):
+  """Filter out empty Layer-like containers and uniquify."""
+  # TODO(b/130381733): Make this an attribute in base_layer.Layer.
+  existing = object_identity.ObjectIdentitySet()
+  to_visit = layer_list[::-1]
+  while to_visit:
+    obj = to_visit.pop()
+    if obj in existing:
+      continue
+    existing.add(obj)
+    if is_layer(obj):
+      yield obj
+    else:
+      sub_layers = getattr(obj, "layers", None) or []
+
+      # Trackable data structures will not show up in ".layers" lists, but
+      # the layers they contain will.
+      to_visit.extend(sub_layers[::-1])

tensorflow/python/trackable/python_state.py

@@ -0,0 +1,87 @@
+"""Utilities for including Python state in TensorFlow checkpoints."""
+# Copyright 2018 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.
+# ==============================================================================
+import abc
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.trackable import base
+from tensorflow.python.util.tf_export import tf_export
+
+
+PYTHON_STATE = "py_state"
+
+
+@tf_export("train.experimental.PythonState")
+class PythonState(base.Trackable, metaclass=abc.ABCMeta):
+  """A mixin for putting Python state in an object-based checkpoint.
+
+  This is an abstract class which allows extensions to TensorFlow's object-based
+  checkpointing (see `tf.train.Checkpoint`). For example a wrapper for NumPy
+  arrays:
+
+  ```python
+  import io
+  import numpy
+
+  class NumpyWrapper(tf.train.experimental.PythonState):
+
+    def __init__(self, array):
+      self.array = array
+
+    def serialize(self):
+      string_file = io.BytesIO()
+      try:
+        numpy.save(string_file, self.array, allow_pickle=False)
+        serialized = string_file.getvalue()
+      finally:
+        string_file.close()
+      return serialized
+
+    def deserialize(self, string_value):
+      string_file = io.BytesIO(string_value)
+      try:
+        self.array = numpy.load(string_file, allow_pickle=False)
+      finally:
+        string_file.close()
+  ```
+
+  Instances of `NumpyWrapper` are checkpointable objects, and will be saved and
+  restored from checkpoints along with TensorFlow state like variables.
+
+  ```python
+  root = tf.train.Checkpoint(numpy=NumpyWrapper(numpy.array([1.])))
+  save_path = root.save(prefix)
+  root.numpy.array *= 2.
+  assert [2.] == root.numpy.array
+  root.restore(save_path)
+  assert [1.] == root.numpy.array
+  ```
+  """
+
+  @abc.abstractmethod
+  def serialize(self):
+    """Callback to serialize the object. Returns a string."""
+
+  @abc.abstractmethod
+  def deserialize(self, string_value):
+    """Callback to deserialize the object."""
+
+  def _serialize_to_tensors(self):
+    """Implements Trackable._serialize_to_tensors."""
+    with ops.init_scope():
+      value = constant_op.constant(self.serialize(), dtype=dtypes.string)
+    return {PYTHON_STATE: value}

tensorflow/python/trackable/python_state_test.py

@@ -0,0 +1,209 @@
+# Copyright 2018 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.
+# ==============================================================================
+import io
+import os
+
+import numpy
+from tensorflow.python.checkpoint import checkpoint as util
+from tensorflow.python.client import session
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.module import module
+from tensorflow.python.platform import test
+from tensorflow.python.trackable import python_state
+
+
+class _NumpyState(module.Module):
+  """A checkpointable object whose NumPy array attributes are saved/restored.
+
+  Example usage:
+
+  ```python
+  arrays = _NumpyState()
+  checkpoint = tf.train.Checkpoint(numpy_arrays=arrays)
+  arrays.x = numpy.zeros([3, 4])
+  save_path = checkpoint.save("/tmp/ckpt")
+  arrays.x[1, 1] = 4.
+  checkpoint.restore(save_path)
+  assert (arrays.x == numpy.zeros([3, 4])).all()
+
+  second_checkpoint = tf.train.Checkpoint(
+      numpy_arrays=_NumpyState())
+  # Attributes of NumpyState objects are created automatically by restore()
+  second_checkpoint.restore(save_path)
+  assert (second_checkpoint.numpy_arrays.x == numpy.zeros([3, 4])).all()
+  ```
+
+  Note that `NumpyState` objects re-create the attributes of the previously
+  saved object on `restore()`. This is in contrast to TensorFlow variables, for
+  which a `Variable` object must be created and assigned to an attribute.
+
+  This snippet works both when graph building and when executing eagerly. On
+  save, the NumPy array(s) are fed as strings to be saved in the checkpoint (via
+  a placeholder when graph building, or as a string constant when executing
+  eagerly). When restoring they skip the TensorFlow graph entirely, and so no
+  restore ops need be run. This means that restoration always happens eagerly,
+  rather than waiting for `checkpoint.restore(...).run_restore_ops()` like
+  TensorFlow variables when graph building.
+  """
+
+  def __init__(self):
+    super(_NumpyState, self).__setattr__("_arrays", module.Module())
+
+  def __getattribute__(self, name):
+    """Un-wrap `_NumpyWrapper` objects when accessing attributes."""
+    try:
+      arrays = super(_NumpyState, self).__getattribute__("_arrays")
+    except AttributeError:
+      # _arrays hasn't been assigned yet
+      return super(_NumpyState, self).__getattribute__(name)
+    try:
+      value = getattr(arrays, name)
+    except AttributeError:
+      dummy_array = numpy.array([])
+      setattr(arrays, name, _NumpyWrapper(dummy_array))
+      value = getattr(arrays, name)
+      if value.array is dummy_array:
+        # No set or restored attribute with this name
+        delattr(arrays, name)
+        return super(_NumpyState, self).__getattribute__(name)
+
+    if isinstance(value, _NumpyWrapper):
+      return value.array
+    return super(_NumpyState, self).__getattribute__(name)
+
+  def __setattr__(self, name, value):
+    """Automatically wrap NumPy arrays assigned to attributes."""
+    if isinstance(value, (numpy.ndarray, numpy.generic)):
+      try:
+        existing = getattr(self._arrays, name)
+        existing.array = value
+        return
+      except AttributeError:
+        value = _NumpyWrapper(value)
+      setattr(self._arrays, name, value)
+      return
+    super(_NumpyState, self).__setattr__(name, value)
+
+
+class _NumpyWrapper(python_state.PythonState):
+  """Wraps a NumPy array for storage in an object-based checkpoint."""
+
+  def __init__(self, array):
+    """Specify a NumPy array to wrap.
+
+    Args:
+      array: The NumPy array to save and restore (may be overwritten).
+    """
+    self.array = array
+
+  def serialize(self):
+    """Callback to serialize the array."""
+    string_file = io.BytesIO()
+    try:
+      numpy.save(string_file, self.array, allow_pickle=False)
+      serialized = string_file.getvalue()
+    finally:
+      string_file.close()
+    return serialized
+
+  def deserialize(self, string_value):
+    """Callback to deserialize the array."""
+    string_file = io.BytesIO(string_value)
+    try:
+      self.array = numpy.load(string_file, allow_pickle=False)  # pylint: disable=unexpected-keyword-arg
+    finally:
+      string_file.close()
+
+
+class NumpyStateTests(test.TestCase):
+
+  def testWrapper(self):
+    directory = self.get_temp_dir()
+    prefix = os.path.join(directory, "ckpt")
+    root = util.Checkpoint(numpy=_NumpyWrapper(numpy.array([1.])))
+    save_path = root.save(prefix)
+    root.numpy.array *= 2.
+    self.assertEqual([2.], root.numpy.array)
+    root.restore(save_path)
+    self.assertEqual([1.], root.numpy.array)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testSaveRestoreNumpyState(self):
+    directory = self.get_temp_dir()
+    prefix = os.path.join(directory, "ckpt")
+    save_state = _NumpyState()
+    saver = util.Checkpoint(numpy=save_state)
+    save_state.a = numpy.ones([2, 2])
+    save_state.b = numpy.ones([2, 2])
+    save_state.b = numpy.zeros([2, 2])
+    save_state.c = numpy.int64(3)
+    self.assertAllEqual(numpy.ones([2, 2]), save_state.a)
+    self.assertAllEqual(numpy.zeros([2, 2]), save_state.b)
+    self.assertEqual(3, save_state.c)
+    first_save_path = saver.save(prefix)
+    save_state.a[1, 1] = 2.
+    save_state.c = numpy.int64(4)
+    second_save_path = saver.save(prefix)
+
+    load_state = _NumpyState()
+    loader = util.Checkpoint(numpy=load_state)
+    loader.restore(first_save_path).initialize_or_restore()
+    self.assertAllEqual(numpy.ones([2, 2]), load_state.a)
+    self.assertAllEqual(numpy.zeros([2, 2]), load_state.b)
+    self.assertEqual(3, load_state.c)
+    load_state.a[0, 0] = 42.
+    self.assertAllEqual([[42., 1.], [1., 1.]], load_state.a)
+    loader.restore(first_save_path).run_restore_ops()
+    self.assertAllEqual(numpy.ones([2, 2]), load_state.a)
+    loader.restore(second_save_path).run_restore_ops()
+    self.assertAllEqual([[1., 1.], [1., 2.]], load_state.a)
+    self.assertAllEqual(numpy.zeros([2, 2]), load_state.b)
+    self.assertEqual(4, load_state.c)
+
+  def testNoGraphPollution(self):
+    graph = ops.Graph()
+    with graph.as_default(), session.Session():
+      directory = self.get_temp_dir()
+      prefix = os.path.join(directory, "ckpt")
+      save_state = _NumpyState()
+      saver = util.Checkpoint(numpy=save_state)
+      save_state.a = numpy.ones([2, 2])
+      save_path = saver.save(prefix)
+      saver.restore(save_path)
+      graph.finalize()
+      saver.save(prefix)
+      save_state.a = numpy.zeros([2, 2])
+      saver.save(prefix)
+      saver.restore(save_path)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testDocstringExample(self):
+    arrays = _NumpyState()
+    checkpoint = util.Checkpoint(numpy_arrays=arrays)
+    arrays.x = numpy.zeros([3, 4])
+    save_path = checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt"))
+    arrays.x[1, 1] = 4.
+    checkpoint.restore(save_path)
+    self.assertAllEqual(numpy.zeros([3, 4]), arrays.x)
+
+    second_checkpoint = util.Checkpoint(numpy_arrays=_NumpyState())
+    second_checkpoint.restore(save_path)
+    self.assertAllEqual(numpy.zeros([3, 4]), second_checkpoint.numpy_arrays.x)
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/python/trackable/resource.py

@@ -0,0 +1,308 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Definitions for resource-type trackable object classes."""
+
+import contextlib
+import copy
+import weakref
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.trackable import base
+from tensorflow.python.util import tf_contextlib
+from tensorflow.python.util.tf_export import tf_export
+
+# global _RESOURCE_TRACKER_STACK
+_RESOURCE_TRACKER_STACK = []
+
+
+class ResourceTracker:
+  """An object that tracks a list of resources."""
+
+  __slots__ = ["_resources"]
+
+  def __init__(self):
+    self._resources = []
+
+  @property
+  def resources(self):
+    return self._resources
+
+  def add_resource(self, resource):
+    self._resources.append(resource)
+
+
+@tf_contextlib.contextmanager
+def resource_tracker_scope(resource_tracker):
+  """A context to manage resource trackers.
+
+  Use this in order to collect up all resources created within a block of code.
+  Example usage:
+
+  ```python
+  resource_tracker = ResourceTracker()
+  with resource_tracker_scope(resource_tracker):
+    resource = TrackableResource()
+
+  assert resource_tracker.resources == [resource]
+
+  Args:
+    resource_tracker: The passed in ResourceTracker object
+
+  Yields:
+    A scope in which the resource_tracker is active.
+  """
+  global _RESOURCE_TRACKER_STACK
+  old = list(_RESOURCE_TRACKER_STACK)
+  _RESOURCE_TRACKER_STACK.append(resource_tracker)
+  try:
+    yield
+  finally:
+    _RESOURCE_TRACKER_STACK = old
+
+
+def _make_getter(captured_getter, captured_previous):
+  """To avoid capturing loop variables."""
+
+  def getter(*args, **kwargs):
+    return captured_getter(captured_previous, *args, **kwargs)
+
+  return getter
+
+
+class _ResourceMetaclass(type):
+  """Metaclass for CapturableResource."""
+
+  def __call__(cls, *args, **kwargs):
+
+    def default_resource_creator(next_creator, *a, **kw):
+      assert next_creator is None
+      obj = cls.__new__(cls, *a, **kw)
+      obj.__init__(*a, **kw)
+      return obj
+
+    previous_getter = lambda *a, **kw: default_resource_creator(None, *a, **kw)
+    resource_creator_stack = ops.get_default_graph()._resource_creator_stack
+    for getter in resource_creator_stack[cls._resource_type()]:
+      previous_getter = _make_getter(getter, previous_getter)
+
+    return previous_getter(*args, **kwargs)
+
+
+class CapturableResource(base.Trackable, metaclass=_ResourceMetaclass):
+  """Holds a Tensor which a tf.function can capture.
+
+  `CapturableResource`s are discovered by traversing the graph of object
+  attributes, e.g. during `tf.saved_model.save`. They are excluded from the
+  scope-based tracking of `TrackableResource`; generally things that require
+  initialization should inherit from `TrackableResource` instead of
+  `CapturableResource` directly.
+  """
+
+  def __init__(self, device=""):
+    """Initialize the `CapturableResource`.
+
+    Args:
+      device: A string indicating a required placement for this resource,
+        e.g. "CPU" if this resource must be created on a CPU device. A blank
+        device allows the user to place resource creation, so generally this
+        should be blank unless the resource only makes sense on one device.
+    """
+    self._resource_handle_value = None
+    self._resource_device = device
+    self._self_destruction_context = (
+        context.eager_mode if context.executing_eagerly()
+        else ops.get_default_graph().as_default)
+
+  @classmethod
+  def _resource_type(cls):
+    return cls.__name__
+
+  @property
+  def _destruction_context(self):
+    return getattr(self, "_self_destruction_context",
+                   # no-op context
+                   contextlib.suppress)
+
+  @_destruction_context.setter
+  def _destruction_context(self, destruction_context):
+    self._self_destruction_context = destruction_context
+
+  def _create_resource(self):
+    """A function that creates a resource handle."""
+    raise NotImplementedError("TrackableResource._create_resource not "
+                              "implemented.")
+
+  @property
+  def _resource_handle(self):
+    return self._resource_handle_value
+
+  @_resource_handle.setter
+  def _resource_handle(self, value):
+    if isinstance(value, (tensor.Tensor, ops.EagerTensor)):
+      value._parent_trackable = weakref.ref(self)  # pylint: disable=protected-access
+    self._resource_handle_value = value
+
+  def _initialize(self):
+    """A function that initializes the resource. Optional."""
+    pass
+
+  def _destroy_resource(self):
+    """A function that destroys the resource. Optional."""
+    pass
+
+  @property
+  def resource_handle(self):
+    """Returns the resource handle associated with this Resource."""
+    if self._resource_handle is None:
+      with ops.device(self._resource_device):
+        self._resource_handle = self._create_resource()
+    return self._resource_handle
+
+  def _export_to_saved_model_graph(
+      self, object_map, tensor_map, **unused_kwargs):
+    """For implementing `Trackable`."""
+    new_obj = copy.copy(self)
+    # pylint: disable=protected-access
+    with ops.device(self._resource_device):
+      new_resource = new_obj._create_resource()
+    new_obj._resource_handle = new_resource
+    # pylint: enable=protected-access
+    object_map[self] = new_obj
+    tensor_map[self.resource_handle] = new_resource
+    return [self.resource_handle]
+
+  def _trackable_children(self, save_type=base.SaveType.CHECKPOINT, **kwargs):
+    children = super()._trackable_children(save_type, **kwargs)
+    if save_type == "savedmodel":
+      @def_function.function(input_signature=[], autograph=False)
+      def _creator():
+        resource = self._create_resource()
+        return resource
+
+      @def_function.function(input_signature=[], autograph=False)
+      def _initializer():
+        self._initialize()
+        return 1  # Dummy return
+
+      @def_function.function(input_signature=[], autograph=False)
+      def _destroyer():
+        self._destroy_resource()
+        return 1  # Dummy return
+
+      children.update({
+          "_create_resource": _creator,
+          "_initialize": _initializer,
+          "_destroy_resource": _destroyer,
+      })
+    return children
+
+  def __del__(self):
+    try:
+      # Outer race condition: on program exit, the destruction context may be
+      # deleted before this __del__ is called. At this point we can safely
+      # exit without calling _destroy_resource() and let Python handle things.
+      with self._destruction_context():
+        # Inner race condition: possible between this and `ScopedTFFunction`
+        # whereby if an entire garbage collection chain containing both
+        # objects is moved to unreachable during the same garbage collection
+        # cycle, the __del__ for `ScopedTFFunction` can be collected before
+        # this method is called. In that case, we can't do much but
+        # continue.
+        self._destroy_resource()
+    except Exception:  # pylint: disable=broad-except
+      # Silence all error logs that occur when attempting to destroy this
+      # resource.
+      pass
+
+
+@tf_export("saved_model.experimental.TrackableResource")
+class TrackableResource(CapturableResource):
+  """Holds a Tensor which a tf.function can capture.
+
+  A TrackableResource is most useful for stateful Tensors that require
+  initialization, such as `tf.lookup.StaticHashTable`. `TrackableResource`s
+  are discovered by traversing the graph of object attributes, e.g. during
+  `tf.saved_model.save`.
+
+  A TrackableResource has three methods to override:
+
+  * `_create_resource` should create the resource tensor handle.
+  * `_initialize` should initialize the resource held at `self.resource_handle`.
+  * `_destroy_resource` is called upon a `TrackableResource`'s destruction
+    and should decrement the resource's ref count. For most resources, this
+    should be done with a call to `tf.raw_ops.DestroyResourceOp`.
+
+  Example usage:
+
+  >>> class DemoResource(tf.saved_model.experimental.TrackableResource):
+  ...   def __init__(self):
+  ...     super().__init__()
+  ...     self._initialize()
+  ...   def _create_resource(self):
+  ...     return tf.raw_ops.VarHandleOp(dtype=tf.float32, shape=[2])
+  ...   def _initialize(self):
+  ...     tf.raw_ops.AssignVariableOp(
+  ...         resource=self.resource_handle, value=tf.ones([2]))
+  ...   def _destroy_resource(self):
+  ...     tf.raw_ops.DestroyResourceOp(resource=self.resource_handle)
+  >>> class DemoModule(tf.Module):
+  ...   def __init__(self):
+  ...     self.resource = DemoResource()
+  ...   def increment(self, tensor):
+  ...     return tensor + tf.raw_ops.ReadVariableOp(
+  ...         resource=self.resource.resource_handle, dtype=tf.float32)
+  >>> demo = DemoModule()
+  >>> demo.increment([5, 1])
+  <tf.Tensor: shape=(2,), dtype=float32, numpy=array([6., 2.], dtype=float32)>
+  """
+
+  def __init__(self, device=""):
+    """Initialize the `TrackableResource`.
+
+    Args:
+      device: A string indicating a required placement for this resource,
+        e.g. "CPU" if this resource must be created on a CPU device. A blank
+        device allows the user to place resource creation, so generally this
+        should be blank unless the resource only makes sense on one device.
+    """
+    global _RESOURCE_TRACKER_STACK
+    for resource_tracker in _RESOURCE_TRACKER_STACK:
+      resource_tracker.add_resource(self)
+    super().__init__(device=device)
+
+
+# TODO(b/124205571,b/124092991): Solve destruction of resources.
+class RestoredResource(TrackableResource):
+  """Restored SavedResource."""
+
+  def __init__(self, device=""):
+    super().__init__(device=device)
+
+  @classmethod
+  def _deserialize_from_proto(cls, object_proto, dependencies, **unused_kwargs):
+    obj = cls(device=object_proto.resource.device)
+    resource_creator = dependencies.get("_create_resource")
+    if resource_creator is not None:
+      obj._create_resource = resource_creator  # pylint: disable=protected-access
+    return obj
+
+  def _add_trackable_child(self, name, value):
+    setattr(self, name, value)
+    if (isinstance(value, base.Trackable) and
+        not isinstance(value, def_function.Function)):
+      self._track_trackable(value, name)

tensorflow/python/trackable/resource_test.py

@@ -0,0 +1,176 @@
+# Copyright 2017 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.
+# ==============================================================================
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import wrap_function
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+from tensorflow.python.trackable import resource
+
+
+def run_inside_wrap_function_in_eager_mode(graph_function):
+  """Decorator to execute the same graph code in eager and graph modes.
+
+  In graph mode, we just execute the graph_function passed as argument. In eager
+  mode, we wrap the function using wrap_function and then execute the wrapped
+  result.
+
+  Args:
+    graph_function: python function containing graph code to be wrapped
+
+  Returns:
+    decorated function
+  """
+  def wrap_and_execute(self):
+    if context.executing_eagerly():
+      wrapped = wrap_function.wrap_function(graph_function, [self])
+      # use the wrapped graph function
+      wrapped()
+    else:
+      # use the original function
+      graph_function(self)
+  return wrap_and_execute
+
+
+class _DummyResource(resource.TrackableResource):
+
+  def __init__(self, handle_name):
+    self._handle_name = handle_name
+    super(_DummyResource, self).__init__()
+
+  def _create_resource(self):
+    return self._handle_name
+
+
+class _DummyResource1(resource.TrackableResource):
+
+  def __init__(self, handle_name):
+    self._handle_name = handle_name
+    self._value = 0
+    super(_DummyResource1, self).__init__()
+
+  def _create_resource(self):
+    return self._handle_name
+
+
+class ResourceTrackerTest(test.TestCase):
+
+  def testBasic(self):
+    resource_tracker = resource.ResourceTracker()
+    with resource.resource_tracker_scope(resource_tracker):
+      dummy_resource1 = _DummyResource("test1")
+      dummy_resource2 = _DummyResource("test2")
+
+    self.assertEqual(2, len(resource_tracker.resources))
+    self.assertEqual("test1", resource_tracker.resources[0].resource_handle)
+    self.assertEqual("test2", resource_tracker.resources[1].resource_handle)
+
+  def testTwoScopes(self):
+    resource_tracker1 = resource.ResourceTracker()
+    with resource.resource_tracker_scope(resource_tracker1):
+      dummy_resource1 = _DummyResource("test1")
+
+    resource_tracker2 = resource.ResourceTracker()
+    with resource.resource_tracker_scope(resource_tracker2):
+      dummy_resource2 = _DummyResource("test2")
+
+    self.assertEqual(1, len(resource_tracker1.resources))
+    self.assertEqual("test1", resource_tracker1.resources[0].resource_handle)
+    self.assertEqual(1, len(resource_tracker2.resources))
+    self.assertEqual("test2", resource_tracker2.resources[0].resource_handle)
+
+  def testNestedScopesScopes(self):
+    resource_tracker = resource.ResourceTracker()
+    with resource.resource_tracker_scope(resource_tracker):
+      resource_tracker1 = resource.ResourceTracker()
+      with resource.resource_tracker_scope(resource_tracker1):
+        dummy_resource1 = _DummyResource("test1")
+
+      resource_tracker2 = resource.ResourceTracker()
+      with resource.resource_tracker_scope(resource_tracker2):
+        dummy_resource2 = _DummyResource("test2")
+
+    self.assertEqual(1, len(resource_tracker1.resources))
+    self.assertEqual("test1", resource_tracker1.resources[0].resource_handle)
+    self.assertEqual(1, len(resource_tracker2.resources))
+    self.assertEqual("test2", resource_tracker2.resources[0].resource_handle)
+    self.assertEqual(2, len(resource_tracker.resources))
+    self.assertEqual("test1", resource_tracker.resources[0].resource_handle)
+    self.assertEqual("test2", resource_tracker.resources[1].resource_handle)
+
+
+class ResourceCreatorScopeTest(test.TestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  @run_inside_wrap_function_in_eager_mode
+  def testResourceCreator(self):
+    def resource_creator_fn(next_creator, *a, **kwargs):
+      kwargs["handle_name"] = "forced_name"
+      return next_creator(*a, **kwargs)
+
+    # test that two resource classes use the same creator function
+    with ops.resource_creator_scope(["_DummyResource", "_DummyResource1"],
+                                    resource_creator_fn):
+      dummy_0 = _DummyResource(handle_name="fake_name_0")
+      dummy_1 = _DummyResource1(handle_name="fake_name_1")
+
+    self.assertEqual(dummy_0._handle_name, "forced_name")
+    self.assertEqual(dummy_1._handle_name, "forced_name")
+
+  @test_util.run_in_graph_and_eager_modes
+  @run_inside_wrap_function_in_eager_mode
+  def testResourceCreatorNestingError(self):
+
+    def creator(next_creator, *a, **kwargs):
+      return next_creator(*a, **kwargs)
+
+    # Save the state so we can clean up at the end.
+    graph = ops.get_default_graph()
+    old_creator_stack = graph._resource_creator_stack["_DummyResource"]
+
+    try:
+      scope = ops.resource_creator_scope(creator, "_DummyResource")
+      scope.__enter__()
+      with ops.resource_creator_scope(creator, "_DummyResource"):
+        with self.assertRaises(RuntimeError):
+          scope.__exit__(None, None, None)
+    finally:
+      graph._resource_creator_stack["_DummyResource"] = old_creator_stack
+
+  @test_util.run_in_graph_and_eager_modes
+  @run_inside_wrap_function_in_eager_mode
+  def testResourceCreatorNesting(self):
+
+    def resource_creator_fn_0(next_creator, *a, **kwargs):
+      instance = next_creator(*a, **kwargs)
+      instance._value = 1
+      return instance
+
+    def resource_creator_fn_1(next_creator, *a, **kwargs):
+      kwargs["handle_name"] = "forced_name1"
+      return next_creator(*a, **kwargs)
+
+    with ops.resource_creator_scope(["_DummyResource1"], resource_creator_fn_0):
+      with ops.resource_creator_scope(["_DummyResource1"],
+                                      resource_creator_fn_1):
+        dummy_0 = _DummyResource1(handle_name="fake_name")
+
+    self.assertEqual(dummy_0._handle_name, "forced_name1")
+    self.assertEqual(dummy_0._value, 1)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/trackable/trackable_utils.py

@@ -0,0 +1,178 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Utility methods for the trackable dependencies."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+
+
+def pretty_print_node_path(path):
+  if not path:
+    return "root object"
+  else:
+    return "root." + ".".join([p.name for p in path])
+
+
+class CyclicDependencyError(Exception):
+
+  def __init__(self, leftover_dependency_map):
+    """Creates a CyclicDependencyException."""
+    # Leftover edges that were not able to be topologically sorted.
+    self.leftover_dependency_map = leftover_dependency_map
+    super(CyclicDependencyError, self).__init__()
+
+
+def order_by_dependency(dependency_map):
+  """Topologically sorts the keys of a map so that dependencies appear first.
+
+  Uses Kahn's algorithm:
+  https://en.wikipedia.org/wiki/Topological_sorting#Kahn's_algorithm
+
+  Args:
+    dependency_map: a dict mapping values to a list of dependencies (other keys
+      in the map). All keys and dependencies must be hashable types.
+
+  Returns:
+    A sorted array of keys from dependency_map.
+
+  Raises:
+    CyclicDependencyError: if there is a cycle in the graph.
+    ValueError: If there are values in the dependency map that are not keys in
+      the map.
+  """
+  # Maps trackables -> trackables that depend on them. These are the edges used
+  # in Kahn's algorithm.
+  reverse_dependency_map = collections.defaultdict(set)
+  for x, deps in dependency_map.items():
+    for dep in deps:
+      reverse_dependency_map[dep].add(x)
+
+  # Validate that all values in the dependency map are also keys.
+  unknown_keys = reverse_dependency_map.keys() - dependency_map.keys()
+  if unknown_keys:
+    raise ValueError("Found values in the dependency map which are not keys: "
+                     f"{unknown_keys}")
+
+  # Generate the list sorted by objects without dependencies -> dependencies.
+  # The returned list will reverse this.
+  reversed_dependency_arr = []
+
+  # Prefill `to_visit` with all nodes that do not have other objects depending
+  # on them.
+  to_visit = [x for x in dependency_map if x not in reverse_dependency_map]
+
+  while to_visit:
+    x = to_visit.pop(0)
+    reversed_dependency_arr.append(x)
+    for dep in set(dependency_map[x]):
+      edges = reverse_dependency_map[dep]
+      edges.remove(x)
+      if not edges:
+        to_visit.append(dep)
+        reverse_dependency_map.pop(dep)
+
+  if reverse_dependency_map:
+    leftover_dependency_map = collections.defaultdict(list)
+    for dep, xs in reverse_dependency_map.items():
+      for x in xs:
+        leftover_dependency_map[x].append(dep)
+    raise CyclicDependencyError(leftover_dependency_map)
+
+  return reversed(reversed_dependency_arr)
+
+
+_ESCAPE_CHAR = "."  # For avoiding conflicts with user-specified names.
+
+# Keyword for identifying that the next bit of a checkpoint variable name is a
+# slot name. Checkpoint names for slot variables look like:
+#
+#   <path to variable>/<_OPTIMIZER_SLOTS_NAME>/<path to optimizer>/<slot name>
+#
+# Where <path to variable> is a full path from the checkpoint root to the
+# variable being slotted for.
+_OPTIMIZER_SLOTS_NAME = _ESCAPE_CHAR + "OPTIMIZER_SLOT"
+# Keyword for separating the path to an object from the name of an
+# attribute in checkpoint names. Used like:
+#   <path to variable>/<_OBJECT_ATTRIBUTES_NAME>/<name of attribute>
+OBJECT_ATTRIBUTES_NAME = _ESCAPE_CHAR + "ATTRIBUTES"
+
+# A constant string that is used to reference the save and restore functions of
+#  Trackable objects that define `_serialize_to_tensors` and
+# `_restore_from_tensors`. This is written as the key in the
+# `SavedObject.saveable_objects<string, SaveableObject>` map in the SavedModel.
+SERIALIZE_TO_TENSORS_NAME = _ESCAPE_CHAR + "TENSORS"
+
+
+def escape_local_name(name):
+  # We need to support slashes in local names for compatibility, since this
+  # naming scheme is being patched in to things like Layer.add_variable where
+  # slashes were previously accepted. We also want to use slashes to indicate
+  # edges traversed to reach the variable, so we escape forward slashes in
+  # names.
+  return (name.replace(_ESCAPE_CHAR, _ESCAPE_CHAR + _ESCAPE_CHAR).replace(
+      r"/", _ESCAPE_CHAR + "S"))
+
+
+def object_path_to_string(node_path_arr):
+  """Converts a list of nodes to a string."""
+  return "/".join(
+      (escape_local_name(trackable.name) for trackable in node_path_arr))
+
+
+def checkpoint_key(object_path, local_name):
+  """Returns the checkpoint key for a local attribute of an object."""
+  key_suffix = escape_local_name(local_name)
+  if local_name == SERIALIZE_TO_TENSORS_NAME:
+    # In the case that Trackable uses the _serialize_to_tensor API for defining
+    # tensors to save to the checkpoint, the suffix should be the key(s)
+    # returned by `_serialize_to_tensor`. The suffix used here is empty.
+    key_suffix = ""
+
+  return f"{object_path}/{OBJECT_ATTRIBUTES_NAME}/{key_suffix}"
+
+
+def extract_object_name(key):
+  """Substrings the checkpoint key to the start of "/.ATTRIBUTES"."""
+  search_key = "/" + OBJECT_ATTRIBUTES_NAME
+  return key[:key.index(search_key)]
+
+
+def extract_local_name(key, prefix=None):
+  """Returns the substring after the "/.ATTIBUTES/" in the checkpoint key."""
+  # "local name" refers to the keys of `Trackable._serialize_to_tensors.`
+  prefix = prefix or ""
+  search_key = OBJECT_ATTRIBUTES_NAME + "/" + prefix
+  # If checkpoint is saved from TF1, return key as is.
+  try:
+    return key[key.index(search_key) + len(search_key):]
+  except ValueError:
+    return key
+
+
+def slot_variable_key(variable_path, optimizer_path, slot_name):
+  """Returns checkpoint key for a slot variable."""
+  # Name slot variables:
+  #
+  #   <variable name>/<_OPTIMIZER_SLOTS_NAME>/<optimizer path>/<slot name>
+  #
+  # where <variable name> is exactly the checkpoint name used for the original
+  # variable, including the path from the checkpoint root and the local name in
+  # the object which owns it. Note that we only save slot variables if the
+  # variable it's slotting for is also being saved.
+
+  return (f"{variable_path}/{_OPTIMIZER_SLOTS_NAME}/{optimizer_path}/"
+          f"{escape_local_name(slot_name)}")

tensorflow/python/trackable/trackable_utils_test.py

@@ -0,0 +1,54 @@
+# Copyright 2021 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 trackable_utils."""
+
+from tensorflow.python.eager import test
+from tensorflow.python.trackable import trackable_utils
+
+
+class TrackableUtilsTest(test.TestCase):
+
+  def test_order_by_dependency(self):
+    """Tests order_by_dependency correctness."""
+
+    # Visual graph (vertical lines point down, so 1 depends on 2):
+    #    1
+    #  /   \
+    # 2 --> 3 <-- 4
+    #       |
+    #       5
+    # One possible order: [5, 3, 4, 2, 1]
+    dependencies = {1: [2, 3], 2: [3], 3: [5], 4: [3], 5: []}
+
+    sorted_arr = list(trackable_utils.order_by_dependency(dependencies))
+    indices = {x: sorted_arr.index(x) for x in range(1, 6)}
+    self.assertEqual(indices[5], 0)
+    self.assertEqual(indices[3], 1)
+    self.assertGreater(indices[1], indices[2])  # 2 must appear before 1
+
+  def test_order_by_no_dependency(self):
+    sorted_arr = list(trackable_utils.order_by_dependency(
+        {x: [] for x in range(15)}))
+    self.assertEqual(set(sorted_arr), set(range(15)))
+
+  def test_order_by_dependency_invalid_map(self):
+    with self.assertRaisesRegex(
+        ValueError, "Found values in the dependency map which are not keys"):
+      trackable_utils.order_by_dependency({1: [2]})
+
+
+if __name__ == "__main__":
+  test.main()
+