chore: import upstream snapshot with attribution

tensorflow/python/training/saving/BUILD

@@ -0,0 +1,106 @@
+# Description:
+#   Low-level utilities for reading and writing checkpoints.
+
+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 = "checkpoint_options",
+    srcs = ["checkpoint_options.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/checkpoint:checkpoint_options",
+        "//tensorflow/python/util:deprecation",
+    ],
+)
+
+py_library(
+    name = "functional_saver",
+    srcs = ["functional_saver.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/checkpoint:functional_saver",
+        "//tensorflow/python/util:deprecation",
+    ],
+)
+
+py_library(
+    name = "saveable_object",
+    srcs = ["saveable_object.py"],
+    strict_deps = True,
+)
+
+py_library(
+    name = "saveable_object_util",
+    srcs = ["saveable_object_util.py"],
+    strict_deps = True,
+    deps = [
+        ":saveable_object",
+        "//tensorflow/python/checkpoint:saveable_compat",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:device",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:control_flow_ops_gen",
+        "//tensorflow/python/ops:ref_variable",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:state_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/trackable:base_delegate",
+        "//tensorflow/python/trackable:python_state",
+        "//tensorflow/python/trackable:trackable_utils",
+        "//tensorflow/python/types:core",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:object_identity",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "trace_saveable_util",
+    srcs = ["trace_saveable_util.py"],
+    strict_deps = True,
+    deps = [
+        ":saveable_object",
+        ":saveable_object_util",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/framework:type_spec",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+tf_py_strict_test(
+    name = "saveable_object_util_test",
+    srcs = ["saveable_object_util_test.py"],
+    deps = [
+        ":saveable_object",
+        ":saveable_object_util",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:resource_variable_ops_gen",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/trackable:base",
+        "//tensorflow/python/trackable:resource",
+    ],
+)

tensorflow/python/training/saving/checkpoint_options.py

@@ -0,0 +1,23 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Options for saving Checkpoints."""
+
+# TODO(kathywu): Delete this file after all imports have been moved to the path
+# below.
+from tensorflow.python.checkpoint import checkpoint_options
+from tensorflow.python.util import deprecation
+
+__getattr__ = deprecation.deprecate_moved_module(
+    __name__, checkpoint_options, "2.11")

tensorflow/python/training/saving/functional_saver.py

@@ -0,0 +1,23 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Saves and restore variables inside traced @tf.functions."""
+
+# TODO(kathywu): Delete this file after all imports have been moved to the path
+# below.
+from tensorflow.python.checkpoint import functional_saver
+from tensorflow.python.util import deprecation
+
+__getattr__ = deprecation.deprecate_moved_module(
+    __name__, functional_saver, "2.11")

tensorflow/python/training/saving/saveable_object.py

@@ -0,0 +1,94 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Types for specifying saving and loading behavior."""
+
+
+class SaveSpec:
+  """Class used to describe tensor slices that need to be saved."""
+
+  def __init__(self, tensor, slice_spec, name, dtype=None, device=None):
+    """Creates a `SaveSpec` object.
+
+    Args:
+      tensor: the tensor to save or callable that produces a tensor to save.
+        If the value is `None`, the `SaveSpec` is ignored.
+      slice_spec: the slice to be saved. See `Variable.SaveSliceInfo`.
+      name: the name to save the tensor under.
+      dtype: The data type of the Tensor. Required if `tensor` is callable.
+        Used for error checking in the restore op.
+      device: The device generating and consuming this tensor. Required if
+        `tensor` is callable. Used to group objects to save by device.
+    """
+    self._tensor = tensor
+    self.slice_spec = slice_spec
+    self.name = name
+    if callable(self._tensor):
+      if dtype is None or device is None:
+        raise AssertionError(
+            "When passing a callable `tensor` to a SaveSpec, an explicit "
+            "dtype and device must be provided.")
+      self.dtype = dtype
+      self.device = device
+    else:
+      self.dtype = tensor.dtype
+      if device is not None:
+        self.device = device
+      else:
+        self.device = tensor.device
+
+  @property
+  def tensor(self):
+    return self._tensor() if callable(self._tensor) else self._tensor
+
+
+class SaveableObject:
+  """Base class for saving and restoring saveable objects."""
+
+  def __init__(self, op, specs, name):
+    """Creates a `SaveableObject` object.
+
+    Args:
+      op: the "producer" object that this class wraps; it produces a list of
+        tensors to save.  E.g., a "Variable" object saving its backing tensor.
+      specs: a list of SaveSpec, each element of which describes one tensor to
+        save under this object. All Tensors must be on the same device.
+      name: the name to save the object under.
+    """
+    self.op = op
+    self.specs = specs
+    self.name = name
+
+  @property
+  def device(self):
+    """The device for SaveSpec Tensors."""
+    return self.specs[0].device
+
+  def restore(self, restored_tensors, restored_shapes):
+    """Restores this object from 'restored_tensors'.
+
+    Args:
+      restored_tensors: the tensors that were loaded from a checkpoint
+      restored_shapes: the shapes this object should conform to after
+        restore, or None.
+
+    Returns:
+      An operation that restores the state of the object.
+
+    Raises:
+      ValueError: If the object cannot be restored using the provided
+        parameters.
+    """
+    # pylint: disable=unused-argument
+    raise ValueError("Calling an abstract method.")

tensorflow/python/training/saving/saveable_object_util.py

@@ -0,0 +1,839 @@
+# Copyright 2015 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 for working with and creating SaveableObjects."""
+import functools
+
+from tensorflow.python.checkpoint import saveable_compat
+from tensorflow.python.client import session
+from tensorflow.python.eager import context
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import device as pydev
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_util
+
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_control_flow_ops
+from tensorflow.python.ops import ref_variable
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.trackable import base as trackable
+from tensorflow.python.trackable import base_delegate
+from tensorflow.python.trackable import python_state
+from tensorflow.python.trackable import trackable_utils
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.types import core
+from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+from tensorflow.python.util import object_identity
+from tensorflow.python.util.tf_export import tf_export
+
+# Op names which identify variable reads which should be saved.
+_VARIABLE_OPS = set(["Variable",
+                     "VariableV2",
+                     "AutoReloadVariable",
+                     "VarHandleOp",
+                     "ReadVariableOp"])
+
+_REF_VARIABLE_OPS = frozenset(["Variable", "VariableV2", "AutoReloadVariable"])
+
+
+def set_cpu0(device_string):
+  """Creates a new device string based on `device_string` but using /CPU:0.
+
+  If the device is already on /CPU:0 or it is a custom device, this is a no-op.
+
+  Args:
+    device_string: A device string.
+
+  Returns:
+    A device string.
+  """
+  if context.is_custom_device(device_string):
+    return device_string
+  parsed_device = pydev.DeviceSpec.from_string(device_string)
+  parsed_device = parsed_device.replace(device_type="CPU", device_index=0)
+  return parsed_device.to_string()
+
+
+class ReferenceVariableSaveable(saveable_object.SaveableObject):
+  """SaveableObject implementation that handles reference variables."""
+
+  def __init__(self, var, slice_spec, name):
+    spec = saveable_object.SaveSpec(var, slice_spec, name, dtype=var.dtype)
+    super(ReferenceVariableSaveable, self).__init__(var, [spec], name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    restored_tensor = restored_tensors[0]
+    if restored_shapes is not None:
+      restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
+    return state_ops.assign(
+        self.op,
+        restored_tensor,
+        validate_shape=restored_shapes is None and
+        self.op.get_shape().is_fully_defined())
+
+
+class ResourceVariableSaveable(saveable_object.SaveableObject):
+  """SaveableObject implementation that handles ResourceVariables."""
+
+  def __init__(self, var, slice_spec, name):
+    self._var_device = var.device
+    self._var_shape = var.shape
+    if isinstance(var, tensor_lib.Tensor):
+      self.handle_op = var.op.inputs[0]
+      tensor = var
+    elif resource_variable_ops.is_resource_variable(var):
+
+      def _read_variable_closure(v):
+        def f():
+          with ops.device(v.device):
+            if context.executing_eagerly() and not v.is_initialized():
+              # A SaveSpec tensor value of `None` indicates that the variable is
+              # uninitialized.
+              return None
+            # Read the variable without making a copy to limit memory usage.
+            x = v.read_value_no_copy()
+            # To allow variables placed on non-CPU devices to be checkpointed,
+            # we copy them to CPU on the same machine first.
+            with ops.device("/device:CPU:0"):
+              return array_ops.identity(x)
+
+        return f
+
+      self.handle_op = var.handle
+      tensor = _read_variable_closure(var)
+    else:
+      raise ValueError(
+          "Saveable is neither a resource variable nor a read operation."
+          f" Got: {repr(var)}")
+    spec = saveable_object.SaveSpec(tensor, slice_spec, name,
+                                    dtype=var.dtype, device=var.device)
+    super(ResourceVariableSaveable, self).__init__(var, [spec], name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    """Restores tensors. Raises ValueError if incompatible shape found."""
+    restored_tensor = restored_tensors[0]
+    if restored_shapes is not None:
+      restored_tensor = array_ops.reshape(restored_tensor, restored_shapes[0])
+    # Copy the restored tensor to the variable's device.
+    with ops.device(self._var_device):
+      restored_tensor = array_ops.identity(restored_tensor)
+      try:
+        assigned_variable = resource_variable_ops.shape_safe_assign_variable_handle(
+            self.handle_op, self._var_shape, restored_tensor)
+      except ValueError as e:
+        raise ValueError(
+            f"Received incompatible tensor with shape {restored_tensor.shape} "
+            f"when attempting to restore variable with shape {self._var_shape} "
+            f"and name {self.name}.") from e
+      return assigned_variable
+
+
+def _tensor_comes_from_variable(v):
+  return isinstance(v, tensor_lib.Tensor) and v.op.type in _VARIABLE_OPS
+
+
+def saveable_objects_for_op(op, name):
+  """Create `SaveableObject`s from an operation.
+
+  Args:
+    op: A variable, operation, or SaveableObject to coerce into a
+      SaveableObject.
+    name: A string name for the SaveableObject.
+
+  Yields:
+    `SaveableObject`s which together save/restore `op`.
+
+  Raises:
+    TypeError: If `name` is not a string.
+    ValueError: For operations with no known conversion to SaveableObject.
+  """
+  if not isinstance(name, str):
+    raise TypeError(
+        "names_to_saveables must be a dict mapping string names to "
+        f"trackable operations. Name is not a string: {name}")
+  if isinstance(op, saveable_object.SaveableObject):
+    yield op
+  elif isinstance(op, (list, tuple, variables.PartitionedVariable)):
+    if isinstance(op, variables.PartitionedVariable):
+      op = list(op)
+    # A set of slices.
+    slice_name = None
+    # pylint: disable=protected-access
+    for variable in op:
+      if isinstance(variable, saveable_object.SaveableObject):
+        yield variable
+        continue
+      if not isinstance(variable, variables.Variable):
+        raise ValueError(f"Slices must all be Variables: {variable}")
+      if not variable._save_slice_info:
+        raise ValueError(f"Slices must all be slices: {variable}")
+      if slice_name is None:
+        slice_name = variable._save_slice_info.full_name
+      elif slice_name != variable._save_slice_info.full_name:
+        raise ValueError(
+            f"Slices must all be from the same tensor: {slice_name} != "
+            f"{variable._save_slice_info.full_name}")
+      if variable.op.type in _REF_VARIABLE_OPS:
+        yield ReferenceVariableSaveable(
+            variable, variable._save_slice_info.spec, name)
+      else:
+        yield ResourceVariableSaveable(variable, variable._save_slice_info.spec,
+                                       name)
+    # pylint: enable=protected-access
+  elif isinstance(op, trackable.Trackable) and not isinstance(
+      op, variables.Variable):
+    # pylint: disable=protected-access
+    for attr, factory in saveable_objects_from_trackable(
+        op, tf1_saver=True).items():
+      if attr == trackable.VARIABLE_VALUE_KEY:
+        # Keep original name for classes masquerading as variables and
+        # Trackables that define _serialize_to_tensors.
+        full_name = name
+      elif attr == trackable_utils.SERIALIZE_TO_TENSORS_NAME:
+        full_name = name
+      else:
+        full_name = name + "_" + attr
+      op = (factory(full_name) if callable(factory) else factory)
+      for op in saveable_objects_for_op(op, op.name):
+        yield op
+    # pylint: enable=protected-access
+  else:
+    # A variable or tensor.
+    if isinstance(op, resource_variable_ops.BaseResourceVariable):
+      if op._in_graph_mode:  # pylint: disable=protected-access
+        variable = op._graph_element  # pylint: disable=protected-access
+      else:
+        variable = op
+      yield ResourceVariableSaveable(variable, "", name)
+    else:
+      if context.executing_eagerly():
+        raise ValueError("Can only save/restore ResourceVariables when "
+                         f"executing eagerly, got type: {type(op)}.")
+
+      variable = ops.convert_to_tensor(op, as_ref=True)
+      if not _tensor_comes_from_variable(variable):
+        raise TypeError(
+            "names_to_saveables must be a dict mapping string "
+            f"names to Tensors/Variables. Not a variable: {variable}")
+      if variable.op.type in _REF_VARIABLE_OPS:
+        yield ReferenceVariableSaveable(variable, "", name)
+      else:
+        yield ResourceVariableSaveable(variable, "", name)
+
+
+def op_list_to_dict(op_list, convert_variable_to_tensor=True):
+  """Create a dictionary of names to operation lists.
+
+  This method is only used when the variable name matters (e.g. when saving
+  or restoring from a TF1 name-based checkpoint). In TF2, this can be called
+  from `tf.train.Checkpoint.restore` when loading from a name-based checkpoint.
+
+  Args:
+    op_list: A (nested) list, tuple, or set of Variables or SaveableObjects.
+    convert_variable_to_tensor: Whether or not to convert single Variables
+      with no slice info into Tensors.
+
+  Returns:
+    A dictionary of names to the operations that must be saved under
+    that name.  Variables with save_slice_info are grouped together under the
+    same key in no particular order.
+
+  Raises:
+    TypeError: If the type of op_list or its elements is not supported.
+    ValueError: If at least two saveables share the same name.
+  """
+  if not isinstance(op_list, (list, tuple, set)):
+    raise TypeError("Variables to save should be passed in a dict or a "
+                    f"list. Got {op_list}")
+  # List casting is necessary to support sets.
+  op_list = nest.flatten(list(op_list))
+  # When ResourceVariables are converted to Tensors, read ops are added to the
+  # graph. Sorting the op_list ensures that the resulting graph is always
+  # constructed in a deterministic way:
+  op_list = sorted(op_list, key=lambda x: x.name)
+  names_to_saveables = {}
+  # pylint: disable=protected-access
+  for var in op_list:
+    resource_or_ref_variable = (
+        isinstance(var, resource_variable_ops.BaseResourceVariable) or
+        isinstance(var, ref_variable.RefVariable))
+
+    if isinstance(var, saveable_object.SaveableObject):
+      names_to_saveables[var.name] = var
+    elif isinstance(var, variables.PartitionedVariable):
+      if var.name in names_to_saveables:
+        raise ValueError(
+            f"At least two variables have the same name: {var.name}")
+      names_to_saveables[var.name] = var
+    elif isinstance(var, variables.Variable) and var._save_slice_info:
+      name = var._save_slice_info.full_name
+      if name in names_to_saveables:
+        if not isinstance(names_to_saveables[name], list):
+          raise ValueError("Mixing slices and non-slices with the same name: "
+                           f"{name}")
+        names_to_saveables[name].append(var)
+      else:
+        names_to_saveables[name] = [var]
+    elif isinstance(var, trackable.Trackable) and not resource_or_ref_variable:
+      trackable_saveables = [
+          (factory() if callable(factory) else factory)
+          for factory in (
+              saveable_objects_from_trackable(var, tf1_saver=True).values())]
+      names_to_saveables.update(
+          op_list_to_dict(trackable_saveables))
+    else:
+      # Variables (reference and resource) have an _in_graph_mode property
+      # indicating whether they were created in a graph building context. We
+      # also get Tensors when graph building, which do not have this property.
+      if not getattr(var, "_in_graph_mode", True):
+        if not isinstance(var, resource_variable_ops.BaseResourceVariable):
+          raise ValueError(
+              "Can only save/restore ResourceVariables when eager execution "
+              f"is enabled. Got type: {type(var)}.")
+        set_var = names_to_saveables.setdefault(var._shared_name, var)
+        if set_var is not var:
+          raise ValueError(
+              "Two different ResourceVariable objects with the same "
+              f"shared_name '{var._shared_name}' were passed to the Saver. This"
+              " likely means that they were created in different Graphs or "
+              "isolated contexts, and may not be checkpointed together.")
+      else:
+        if convert_variable_to_tensor:
+          if isinstance(var, resource_variable_ops.BaseResourceVariable):
+            var = var._graph_element  # pylint: disable=protected-access
+          else:
+            var = ops.convert_to_tensor(var, as_ref=True)
+          if not _tensor_comes_from_variable(var):
+            raise TypeError(f"Variable to save is not a Variable: {var}")
+        if var.op.type == "ReadVariableOp":
+          name = var.op.inputs[0].op.name
+        else:
+          name = var.op.name
+        if name in names_to_saveables:
+          raise ValueError(f"At least two variables have the same name: {name}")
+        names_to_saveables[name] = var
+
+    # pylint: enable=protected-access
+  return names_to_saveables
+
+
+def _add_saveable(saveables, seen_ops, saveable):
+  """Adds the saveable to the saveables list.
+
+  Args:
+    saveables: List to append the SaveableObject to.
+    seen_ops: Set of the ops of the saveables already processed.  Used to
+      check that each saveable is only saved once.
+    saveable: The saveable.
+
+  Raises:
+    ValueError: If the saveable has already been processed.
+  """
+  if saveable.op is not None and saveable.op in seen_ops:
+    raise ValueError("The same saveable will be restored with two names: "
+                     f"{saveable.name}")
+  saveables.append(saveable)
+  seen_ops.add(saveable.op)
+
+
+def validate_and_slice_inputs(names_to_saveables):
+  """Returns the variables and names that will be used for a Saver.
+
+  Args:
+    names_to_saveables: A dict (k, v) where k is the name of an operation and
+       v is an operation to save or a BaseSaverBuilder.Saver.
+
+  Returns:
+    A list of SaveableObjects.
+
+  Raises:
+    TypeError: If any of the keys are not strings or any of the
+      values are not one of Tensor or Variable or a trackable operation.
+    ValueError: If the same operation is given in more than one value
+      (this also applies to slices of SlicedVariables).
+  """
+  saveables = []
+  seen_ops = object_identity.ObjectIdentitySet()
+  for name, op in sorted(names_to_saveables.items(),
+                         # Avoid comparing ops, sort only by name.
+                         key=lambda x: x[0]):
+    for converted_saveable_object in saveable_objects_for_op(op, name):
+      _add_saveable(saveables, seen_ops, converted_saveable_object)
+  return saveables
+
+
+def validate_saveables_for_saved_model(saveables, obj):
+  """Makes sure SaveableObjects are compatible with SavedModel."""
+  if isinstance(obj, python_state.PythonState):
+    logging.warn(
+        f"Note that object {obj} stores python values into the checkpoint. "
+        "These values will not be restored when loading the SavedModel "
+        "into python.")
+    return []
+  if any(isinstance(saveable, trackable.NoRestoreSaveable)
+         for saveable in saveables):
+    return []
+  return saveables
+
+
+class RestoredSaveableObject(saveable_object.SaveableObject):
+  """SaveableObject restored from SavedModel using the traced save/restore."""
+
+  def __init__(self, names_and_slices, save_function, restore_function, name):
+    self.save_function = save_function
+    self.restore_function = restore_function
+
+    if tensor_util.is_tf_type(name):
+      name_tensor = name
+    else:
+      with ops.init_scope():
+        name_tensor = constant_op.constant(name)
+    tensors = save_function(name_tensor)
+    specs = []
+    for (str_name, str_slice), tensor_info in zip(names_and_slices, tensors):
+      specs.append(saveable_object.SaveSpec(tensor_info["tensor"], str_slice,
+                                            name + str_name))
+    super(RestoredSaveableObject, self).__init__(None, specs, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    del restored_shapes  # unused
+    return self.restore_function(
+        *[restored_tensors[i] for i in range(len(self.specs))])
+
+
+def recreate_saveable_objects(saveable_fn_by_name, temp_session):
+  """Returns a dict of SaveableObject factories generated from loaded fns."""
+
+  names_and_slices = []
+
+  with ops.init_scope():
+
+    for save_fn, _ in saveable_fn_by_name.values():
+      for tensor_info in save_fn(""):
+        name = tensor_info["name"]
+        slice_spec = tensor_info["slice_spec"]
+        if not context.executing_eagerly():
+          sess = ops.get_default_session()
+          if sess is None:
+            if temp_session[0] is not None:
+              sess = temp_session[0]
+            else:
+              sess = temp_session[0] = session.Session()
+          name, slice_spec = sess.run([name, slice_spec])
+        names_and_slices.append((
+            _convert_to_string(name),
+            _convert_to_string(slice_spec)))
+
+  saveable_factories = {}
+  for name, (save_fn, restore_fn) in saveable_fn_by_name.items():
+    saveable_factories[name] = functools.partial(
+        RestoredSaveableObject,
+        names_and_slices=names_and_slices,
+        save_function=save_fn,
+        restore_function=restore_fn)
+  return saveable_factories
+
+
+def create_saveable_object(name, key, factory, call_with_mapped_captures):
+  """Creates a SaveableObject while potentially in a different graph.
+
+  When creating the frozen saver for SavedModel, the save and restore ops are
+  placed in a separate graph. Since RestoredSaveableObject uses tf.functions to
+  save and restore, the function captures must be mapped to the new graph.
+
+  Args:
+    name: Name of SaveableObject factory.
+    key: Checkpoint key of this SaveableObject.
+    factory: Factory method for creating the SaveableObject.
+    call_with_mapped_captures: Helper that calls a tf.function while remapping
+      the captures.
+
+  Returns:
+    a SaveableObject.
+  """
+  if call_with_mapped_captures is None:
+    return factory(name=key)
+  if name == trackable_utils.SERIALIZE_TO_TENSORS_NAME:
+    return factory(name=key,
+                   call_with_mapped_captures=call_with_mapped_captures)
+  elif is_factory_for_restored_saveable_object(factory):
+    concrete_save_fn = factory.keywords["save_function"]
+
+    def save_fn(name):
+      return call_with_mapped_captures(concrete_save_fn, [name])
+
+    concrete_restore_fn = factory.keywords["restore_function"]
+
+    def restore_fn(*restored_tensors):
+      return call_with_mapped_captures(concrete_restore_fn, restored_tensors)
+
+    return factory(save_function=save_fn, restore_function=restore_fn,
+                   name=key)
+  else:
+    return factory(name=key)
+
+
+def is_factory_for_restored_saveable_object(factory):
+  return (isinstance(factory, functools.partial) and
+          factory.func is RestoredSaveableObject)
+
+
+@tf_export("__internal__.tracking.saveable_objects_from_trackable", v1=[])
+def saveable_objects_from_trackable(obj, tf1_saver=False):
+  """Returns SaveableObject factory dict from a Trackable.
+
+  Args:
+    obj: A `Trackable`
+    tf1_saver: Boolean, whether this is being called from a TF1 Saver (
+        `tf.compat.v1.train.Saver`). When this is True, the SaveableObject will
+        be generated from `obj`'s legacy `_gather_saveables_for_checkpoint` fn.
+        When saving with TF2, `Trackable._serialize_from_tensors` is preferred.
+
+  Returns:
+    A dict mapping attribute names to SaveableObject factories (callables that
+    produce a SaveableObject).
+  """
+  if isinstance(obj, python_state.PythonState):
+    return {
+        python_state.PYTHON_STATE:
+            functools.partial(
+                _PythonStringStateSaveable,
+                state_callback=obj.serialize,
+                restore_callback=obj.deserialize)
+    }
+
+  if tf1_saver:
+    saveable_factories = obj._gather_saveables_for_checkpoint()  # pylint: disable=protected-access
+    if saveable_factories:
+      return saveable_factories
+
+  if trackable_has_serialize_to_tensor(obj):
+
+    def create_saveable(name="", call_with_mapped_captures=None):
+      save_fn = obj._serialize_to_tensors  # pylint: disable=protected-access
+      if (call_with_mapped_captures and
+          isinstance(save_fn, core.ConcreteFunction)):
+        tensor_dict = call_with_mapped_captures(save_fn, [])
+      else:
+        tensor_dict = save_fn()
+
+      specs = []
+      local_names = []
+      for tensor_name, maybe_tensor in tensor_dict.items():
+        local_names.append(tensor_name)
+
+        if not isinstance(maybe_tensor, dict):
+          maybe_tensor = {"": maybe_tensor}
+
+        spec_name = name + trackable_utils.escape_local_name(tensor_name)
+        # Create separate specs for each slice spec.
+        for slice_spec, tensor in maybe_tensor.items():
+          if isinstance(tensor, saveable_object.SaveSpec):
+            spec = tensor
+            spec.name = spec_name
+            spec.slice_spec = slice_spec
+          else:
+            spec = saveable_object.SaveSpec(tensor, slice_spec, spec_name)
+          specs.append(spec)
+
+      return TrackableSaveable(
+          obj=obj,
+          specs=specs,
+          name=name,
+          local_names=local_names,
+          prefix=saveable_compat.get_saveable_name(obj) or "",
+          call_with_mapped_captures=call_with_mapped_captures)
+
+    return {trackable_utils.SERIALIZE_TO_TENSORS_NAME: create_saveable}
+  else:
+    return obj._gather_saveables_for_checkpoint()  # pylint: disable=protected-access
+
+
+class TrackableSaveable(saveable_object.SaveableObject):
+  """A SaveableObject that defines `Trackable` checkpointing steps."""
+
+  def __init__(self, obj, specs, name, local_names, prefix,
+               call_with_mapped_captures=None):
+    self._prefix = prefix
+    self._local_names = local_names
+    self._trackable = obj
+    self._call_with_mapped_captures = call_with_mapped_captures
+    super(TrackableSaveable, self).__init__(obj, specs, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    del restored_shapes  # Unused.
+    restored_tensor_dict = {}
+    for n, local_name in enumerate(self._local_names):
+      restored_tensor_dict[local_name] = restored_tensors[n]
+
+    restore_fn = self._trackable._restore_from_tensors  # pylint: disable=protected-access
+
+    # When restoring a RefVariable, call the restore function directly.
+    # pylint: disable=protected-access
+    if not ops.executing_eagerly_outside_functions() and any([
+        spec._tensor.op.type in _REF_VARIABLE_OPS
+        for spec in self.specs
+        if isinstance(spec._tensor, tensor_lib.Tensor)]):
+      return restore_fn(restored_tensor_dict)
+    # pylint: enable=protected-access
+
+    if (self._call_with_mapped_captures and
+        isinstance(restore_fn, core.ConcreteFunction)):
+      ret = self._call_with_mapped_captures(restore_fn, [restored_tensor_dict])
+    else:
+      ret = restore_fn(restored_tensor_dict)
+    if ret is not None:
+      return ret
+    return gen_control_flow_ops.no_op()
+
+  def get_proto_names_and_checkpoint_keys(self):
+    return [(self._prefix + local_name, spec.name)
+            for local_name, spec in zip(self._local_names, self.specs)]
+
+
+class _PythonStringStateSaveable(saveable_object.SaveableObject):
+  """Saves Python state in a checkpoint."""
+
+  def __init__(self, name, state_callback, restore_callback):
+    """Configure saving.
+
+    Args:
+      name: The checkpoint key to write to.
+      state_callback: A function taking no arguments which returns a string.
+        This function is run every time a checkpoint is written.
+      restore_callback: A function taking a Python string, used to restore
+        state.
+    """
+
+    def _state_callback_wrapper():
+      with ops.init_scope():
+        return state_callback()
+
+    self._state_callback = _state_callback_wrapper
+    self._restore_callback = restore_callback
+    with ops.device("/cpu:0"):
+      self._save_string = constant_op.constant("", dtype=dtypes.string)
+    spec = saveable_object.SaveSpec(
+        self._save_string, "", name, dtype=dtypes.string)
+    super(_PythonStringStateSaveable, self).__init__(self._save_string, [spec],
+                                                     name)
+
+  def feed_dict_additions(self):
+    """When running a graph, indicates fresh state to feed."""
+    return {self._save_string: self._state_callback()}
+
+  def freeze(self):
+    """Create a frozen `SaveableObject` which saves the current state."""
+
+    def _constant_state():
+      return constant_op.constant(self._state_callback(), dtype=dtypes.string)
+
+    return trackable.NoRestoreSaveable(
+        tensor=_constant_state,
+        dtype=dtypes.string,
+        name=self.name,
+        device="cpu:0")
+
+
+def trackable_has_serialize_to_tensor(obj):
+  """Returns whether obj's class has `_serialize_to_tensors` defined."""
+  if obj is base_delegate.DelegatingTrackableMixin:
+    # DelegatingTrackableMixin always delegates "_serialize_to_tensors"
+    # to its inner `trackable`, so we check whether the inner trackable
+    # has `_serialize_to_tensor`.
+    return trackable_has_serialize_to_tensor(obj._trackable)  # pylint: disable=protected-access
+
+  try:
+    if "_serialize_to_tensors" in obj.__dict__:
+      # In some cases (e.g. restored objects), the object may have
+      # `_serialize_to_tensors` even if the class does not.
+      return True
+  except (AttributeError, TypeError):
+    # Data structure proxy wrappers don't have __dict__.
+    pass
+
+  # Use MRO so that if a parent class has `_serialize_to_tensors`, but the
+  # object class has not yet been migrated, we'll continue to use the obj
+  # class's `_gather_saveables_for_checkpoint` method.
+  for t in type(obj).mro():
+    if t is base_delegate.DelegatingTrackableMixin:
+      # DelegatingTrackableMixin always delegates "_serialize_to_tensors"
+      # to its inner `trackable`, so we check whether the inner trackable
+      # has `_serialize_to_tensor`.
+      return trackable_has_serialize_to_tensor(obj._trackable)  # pylint: disable=protected-access
+    if t is trackable.Trackable:
+      # Base case. Return False since _serialize_to_tensors will raise a
+      # NotImplemented Error.
+      return False
+    elif "_serialize_to_tensors" in t.__dict__:
+      return True
+    elif "_gather_saveables_for_checkpoint" in t.__dict__:
+      return False
+  return False
+
+
+def _convert_to_string(x):
+  return compat.as_str(tensor_util.constant_value(x))
+
+
+class SaveableCompatibilityConverter(trackable.Trackable):
+  """Converts object's `SaveableObjects` to functions used in TF2 checkpointing.
+
+  A class that converts a Trackable object's `SaveableObjects` to save and
+  restore functions with the same signatures as
+  `Trackable._serialize_to_tensors` and `Trackable._restore_from_tensors`.
+  This class also produces a method for filling the object proto.
+  """
+
+  __slots__ = ("_obj", "_saveables")
+
+  def __init__(self, obj, saveables):
+    """Constructor.
+
+    Args:
+      obj: A Trackable object.
+      saveables: A list of saveables for `obj`.
+    """
+    self._obj = obj
+    self._saveables = saveables
+
+  @property
+  def obj(self):
+    return self._obj
+
+  @property
+  def saveables(self):
+    """Returns a list of SaveableObjects generated from the Trackable object."""
+    return self._saveables
+
+  def _serialize_to_tensors(self):
+    """Returns a dict of tensors to serialize."""
+    return saveable_object_to_tensor_dict(self.saveables)
+
+  def _restore_from_tensors(self, restored_tensors):
+    """Returns the restore ops defined in the Saveables."""
+    # Map restored tensors to the corresponding SaveableObjects, then call
+    # restore. There must be an exact match between restored tensors and the
+    # expected attributes.
+    expected_keys = []
+    for saveable in self.saveables:
+      expected_keys.extend(
+          trackable_utils.extract_local_name(_convert_to_string(spec.name))
+          for spec in saveable.specs)
+    if set(expected_keys) != restored_tensors.keys():
+      raise ValueError(f"Could not restore object {self._obj} because not all "
+                       "expected tensors were in the checkpoint."
+                       f"\n\tExpected: {expected_keys}"
+                       f"\n\tGot: {list(restored_tensors.keys())}")
+
+    return saveable_object_to_restore_fn(self.saveables)(restored_tensors)
+
+
+def saveable_object_to_tensor_dict(saveables):
+  """Converts a list of SaveableObjects to a tensor dictionary."""
+  tensor_dict = {}
+  for saveable in saveables:
+    for spec in saveable.specs:
+      name = _convert_to_string(spec.name)
+      slice_spec = _convert_to_string(spec.slice_spec)
+      # Currently, tensor dict cannot handle callable tensor values (which
+      # are needed for uninitialized variables), so keep using SaveSpec.
+      tensor = spec if callable(spec._tensor) else spec._tensor  # pylint: disable=protected-access
+      if slice_spec:
+        tensor_dict.setdefault(name, {})[slice_spec] = tensor
+      else:
+        tensor_dict[name] = tensor
+  return tensor_dict
+
+
+def saveable_object_to_restore_fn(saveables):
+  """Generates `Trackable._restore_from_tensors` from SaveableObjects."""
+
+  def _restore_from_tensors(restored_tensors):
+    restore_ops = {}
+
+    for saveable in saveables:
+      saveable_restored_tensors = []
+      for spec in saveable.specs:
+        name = trackable_utils.extract_local_name(_convert_to_string(spec.name))
+        slice_spec = _convert_to_string(spec.slice_spec)
+
+        maybe_tensor = restored_tensors[name]
+        if not isinstance(maybe_tensor, dict):
+          maybe_tensor = {"": maybe_tensor}
+
+        saveable_restored_tensors.append(maybe_tensor[slice_spec])
+      restore_ops[saveable.name] = saveable.restore(
+          saveable_restored_tensors, restored_shapes=None)
+    return restore_ops
+
+  return _restore_from_tensors
+
+
+def serialized_tensors_to_saveable_cache(serialized_tensors):
+  """Converts a tensor dict to a SaveableObject cache.
+
+  Args:
+    serialized_tensors: Map from Trackable to a tensor dict. The tensor dict
+      maps checkpoint key (-> slice_spec) -> Tensor
+
+  Returns:
+    A dict mapping Trackable objects to a map from local savable name to
+    SaveableObject.
+  """
+  saveables_cache = object_identity.ObjectIdentityWeakKeyDictionary()
+
+  for obj, tensor_dict in serialized_tensors.items():
+    if not tensor_dict: continue
+    if isinstance(obj, SaveableCompatibilityConverter):
+      trackable_obj = obj.obj
+      saveables_cache[trackable_obj] = {}
+      for saveable in obj.saveables:
+        local_name = trackable_utils.extract_local_name(saveable.name)
+        saveables_cache[trackable_obj][local_name] = [saveable]
+      continue
+
+    specs = []
+    # The local names and prefixes are computed to ensure that the generated
+    # SaveableObject can call `Trackable._restore_from_tensors()`
+    local_names = []
+    prefix = saveable_compat.get_saveable_name(obj) or ""
+    for checkpoint_key, maybe_tensor in tensor_dict.items():
+      # Make sure that `maybe_tensor` is a dict from `slice_spec` to `tensor`.
+      if not isinstance(maybe_tensor, dict):
+        maybe_tensor = {"": maybe_tensor}
+
+      for slice_spec, tensor in maybe_tensor.items():
+        if isinstance(tensor, saveable_object.SaveSpec):
+          specs.append(tensor)
+        else:
+          specs.append(saveable_object.SaveSpec(tensor,
+                                                slice_spec,
+                                                checkpoint_key))
+      local_names.append(trackable_utils.extract_local_name(checkpoint_key,
+                                                            prefix))
+
+    object_name = trackable_utils.extract_object_name(
+        next(iter(tensor_dict.keys())))
+    saveables_cache[obj] = {
+        trackable_utils.SERIALIZE_TO_TENSORS_NAME: [TrackableSaveable(
+            obj, specs, object_name, local_names=local_names, prefix=prefix)]}
+  return saveables_cache

tensorflow/python/training/saving/saveable_object_util_test.py

@@ -0,0 +1,251 @@
+# 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.
+# ==============================================================================
+"""Tests for saveable_object_util."""
+
+import os
+
+from tensorflow.python.checkpoint import checkpoint
+from tensorflow.python.eager import context
+from tensorflow.python.eager import test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.trackable import base
+from tensorflow.python.trackable import resource
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.training.saving import saveable_object_util
+
+
+class _VarSaveable(saveable_object.SaveableObject):
+
+  def __init__(self, var, slice_spec, name):
+    specs = [saveable_object.SaveSpec(var.read_value(), slice_spec, name)]
+    super().__init__(var, specs, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    return self.op.assign(restored_tensors[0])
+
+
+def _create_converted_trackable(obj):
+  saveable_factories = saveable_object_util.saveable_objects_from_trackable(obj)
+  saveables = [factory(name) for name, factory in saveable_factories.items()]
+  return saveable_object_util.SaveableCompatibilityConverter(obj, saveables)
+
+
+class SaveableCompatibilityConverterTest(test.TestCase):
+
+  def test_convert_no_saveable(self):
+    t = base.Trackable()
+    converter = _create_converted_trackable(t)
+    self.assertEmpty(converter._serialize_to_tensors())
+    converter._restore_from_tensors({})
+
+    with self.assertRaisesRegex(ValueError, "Could not restore object"):
+      converter._restore_from_tensors({"": 0})
+
+  def test_convert_single_saveable(self):
+
+    class MyTrackable(base.Trackable):
+
+      def __init__(self):
+        self.a = variables.Variable(5.0)
+
+      def _gather_saveables_for_checkpoint(self):
+        return {"a": lambda name: _VarSaveable(self.a, "", name)}
+
+    t = MyTrackable()
+    converter = _create_converted_trackable(t)
+
+    serialized_tensors = converter._serialize_to_tensors()
+    self.assertLen(serialized_tensors, 1)
+    self.assertIn("a", serialized_tensors)
+    self.assertEqual(5, self.evaluate(serialized_tensors["a"]))
+
+    with self.assertRaisesRegex(ValueError, "Could not restore object"):
+      converter._restore_from_tensors({})
+    with self.assertRaisesRegex(ValueError, "Could not restore object"):
+      converter._restore_from_tensors({"not_a": 1.})
+
+    self.assertEqual(5, self.evaluate(t.a))
+    converter._restore_from_tensors({"a": 123.})
+    self.assertEqual(123, self.evaluate(t.a))
+
+  def test_convert_single_saveable_renamed(self):
+
+    class MyTrackable(base.Trackable):
+
+      def __init__(self):
+        self.a = variables.Variable(15.0)
+
+      def _gather_saveables_for_checkpoint(self):
+        return {"a": lambda name: _VarSaveable(self.a, "", name + "-value")}
+
+    t = MyTrackable()
+    converter = _create_converted_trackable(t)
+
+    serialized_tensors = converter._serialize_to_tensors()
+
+    self.assertLen(serialized_tensors, 1)
+    self.assertEqual(15, self.evaluate(serialized_tensors["a-value"]))
+
+    with self.assertRaisesRegex(ValueError, "Could not restore object"):
+      converter._restore_from_tensors({"a": 1.})
+
+    self.assertEqual(15, self.evaluate(t.a))
+    converter._restore_from_tensors({"a-value": 456.})
+    self.assertEqual(456, self.evaluate(t.a))
+
+
+class _MultiSpecSaveable(saveable_object.SaveableObject):
+
+  def __init__(self, obj, name):
+    self.obj = obj
+    specs = [
+        saveable_object.SaveSpec(obj.a, "", name + "-a"),
+        saveable_object.SaveSpec(obj.b, "", name + "-b")]
+    super(_MultiSpecSaveable, self).__init__(None, specs, name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    del restored_shapes  # Unused.
+    self.obj.a.assign(restored_tensors[0])
+    self.obj.b.assign(restored_tensors[1])
+
+
+class State(resource.TrackableResource):
+
+  def __init__(self, initial_value):
+    super().__init__()
+    self._initial_value = initial_value
+    self._initialize()
+
+  def _create_resource(self):
+    return gen_resource_variable_ops.var_handle_op(
+        shape=[],
+        dtype=dtypes.float32,
+        shared_name=context.anonymous_name(),
+        name="StateVar",
+        container="")
+
+  def _initialize(self):
+    gen_resource_variable_ops.assign_variable_op(self.resource_handle,
+                                                 self._initial_value)
+
+  def _destroy_resource(self):
+    gen_resource_variable_ops.destroy_resource_op(self.resource_handle,
+                                                  ignore_lookup_error=True)
+
+  def read(self):
+    return gen_resource_variable_ops.read_variable_op(self.resource_handle,
+                                                      dtypes.float32)
+
+  def assign(self, value):
+    gen_resource_variable_ops.assign_variable_op(self.resource_handle, value)
+
+
+class _StateSaveable(saveable_object.SaveableObject):
+
+  def __init__(self, obj, name):
+    spec = saveable_object.SaveSpec(obj.read(), "", name)
+    self.obj = obj
+    super(_StateSaveable, self).__init__(obj, [spec], name)
+
+  def restore(self, restored_tensors, restored_shapes):
+    del restored_shapes  # Unused.
+    self.obj.assign(restored_tensors[0])
+
+
+class SaveableState(State):
+
+  def _gather_saveables_for_checkpoint(self):
+    return {
+        "value": lambda name: _StateSaveable(self, name)
+    }
+
+
+class TrackableState(State):
+
+  def _serialize_to_tensors(self):
+    return {
+        "value": self.read()
+    }
+
+  def _restore_from_tensors(self, restored_tensors):
+    return self.assign(restored_tensors["value"])
+
+
+class SaveableCompatibilityEndToEndTest(test.TestCase):
+
+  def test_checkpoint_comparison(self):
+    saveable_state = SaveableState(5.)
+    trackable_state = TrackableState(10.)
+
+    # First test that SaveableState and TrackableState are equivalent by
+    # saving a checkpoint with both objects and swapping values.
+
+    self.assertEqual(5, self.evaluate(saveable_state.read()))
+    self.assertEqual(10, self.evaluate(trackable_state.read()))
+
+    ckpt_path = os.path.join(self.get_temp_dir(), "ckpt")
+    checkpoint.Checkpoint(a=saveable_state, b=trackable_state).write(ckpt_path)
+
+    status = checkpoint.Checkpoint(b=saveable_state,
+                                   a=trackable_state).read(ckpt_path)
+    status.assert_consumed()
+
+    self.assertEqual(10, self.evaluate(saveable_state.read()))
+    self.assertEqual(5, self.evaluate(trackable_state.read()))
+
+    # Test that the converted SaveableState is compatible with the checkpoint
+    # saved above.
+    to_convert = SaveableState(0.0)
+
+    converted_saveable_state = _create_converted_trackable(to_convert)
+
+    checkpoint.Checkpoint(a=converted_saveable_state).read(
+        ckpt_path).assert_existing_objects_matched().expect_partial()
+    self.assertEqual(5, self.evaluate(to_convert.read()))
+
+    checkpoint.Checkpoint(b=converted_saveable_state).read(
+        ckpt_path).assert_existing_objects_matched().expect_partial()
+    self.assertEqual(10, self.evaluate(to_convert.read()))
+
+
+class HasSerializeToTensorTest(test.TestCase):
+  def test_has_serialize_to_tensor(self):
+    class ReturnsTrue(base.Trackable):
+      def _serialize_to_tensors(self):
+        return {}
+
+    class ReturnsFalse(base.Trackable):
+      pass
+
+    class SubclassReturnsFalse(ReturnsTrue):
+      def _gather_saveables_for_checkpoint(self):
+        return {}
+
+    self.assertTrue(saveable_object_util.trackable_has_serialize_to_tensor(
+        ReturnsTrue()))
+    self.assertFalse(saveable_object_util.trackable_has_serialize_to_tensor(
+        ReturnsFalse()))
+
+    # This should return False, because even though its parent class has
+    # `_serialize_to_tensors`, the class itself defines
+    # `_gather_saveables_for_checkpoint`.
+    self.assertFalse(saveable_object_util.trackable_has_serialize_to_tensor(
+        SubclassReturnsFalse()))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/training/saving/trace_saveable_util.py

@@ -0,0 +1,116 @@
+# Copyright 2015 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 for tracing save and restore functions for SaveableObjects."""
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import type_spec
+
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.training.saving import saveable_object_util
+from tensorflow.python.util import nest
+
+
+def trace_save_restore_function_map(obj, factory_data_list):
+  """Traces all save and restore functions in the provided factory list.
+
+  Args:
+    obj: `Trackable` object.
+    factory_data_list: List of `_CheckpointFactoryData`.
+
+  Returns:
+    Dict mapping atttribute names to tuples of concrete save/restore functions.
+  """
+  saveable_fns = {}
+
+  for factory_data in factory_data_list:
+    saveable_factory = factory_data.factory
+    attribute_name = factory_data.name
+
+    # If object revives as a resource (or TPU/Mirrored) variable,
+    # there is no need to trace the save and restore functions.
+    if (resource_variable_ops.is_resource_variable(obj) or
+        resource_variable_ops.is_resource_variable(saveable_factory) or
+        not callable(saveable_factory)):
+      continue
+
+    concrete_save, concrete_restore = (
+        _trace_save_restore_functions(saveable_factory, obj))
+    if not concrete_save:
+      continue
+    saveable_fns[attribute_name] = (concrete_save, concrete_restore)
+  return saveable_fns
+
+
+def _trace_save_restore_functions(saveable_factory, obj):
+  """Traces save and restore functions."""
+  if saveable_object_util.is_factory_for_restored_saveable_object(
+      saveable_factory):
+    return (
+        saveable_factory.keywords["save_function"],
+        saveable_factory.keywords["restore_function"],
+    )
+
+  saveables = []  # Store the saveables in a data structure accessible to both
+  # the save and restore functions.
+
+  @def_function.function(
+      input_signature=[tensor_spec.TensorSpec([], dtypes.string)]
+  )
+  def save_fn(checkpoint_key):
+    maybe_saveable = saveable_factory(name=checkpoint_key)
+    if isinstance(maybe_saveable, saveable_object.SaveableObject):
+      maybe_saveable = [maybe_saveable]
+    saveables[:] = maybe_saveable
+
+    # Return list of all SaveSpecs created by the factory.
+    ret = []
+    for saveable in saveables:
+      for spec in saveable.specs:
+        ret.append({"name": spec.name, "tensor": spec.tensor,
+                    "slice_spec": spec.slice_spec})
+    return ret
+
+  concrete_save = save_fn.get_concrete_function()
+
+  # The SaveableObjects are produced when `save_fn` is traced.
+  saveables = saveable_object_util.validate_saveables_for_saved_model(
+      saveables, obj)
+  if not saveables:
+    return None, None
+
+  # Use the SaveSpecs to define the input signature of the restore function.
+  restored_type_specs = []
+  tensor_structure = []
+  for saveable in saveables:
+    saveable_tensor_structure = []
+    tensor_structure.append(saveable_tensor_structure)
+    for spec in saveable.specs:
+      restored_type_specs.append(type_spec.type_spec_from_value(spec.tensor))
+      saveable_tensor_structure.append(spec.name)
+
+  @def_function.function(input_signature=restored_type_specs)
+  def restore_fn(*restored_tensors):
+    structured_restored_tensors = nest.pack_sequence_as(
+        tensor_structure, restored_tensors)
+    for saveable, restored_tensors in zip(saveables,
+                                          structured_restored_tensors):
+      saveable.restore(restored_tensors, restored_shapes=None)
+    return 1  # Return dummy tensor
+
+  concrete_restore = restore_fn.get_concrete_function()
+  return concrete_save, concrete_restore