chore: import upstream snapshot with attribution

tensorflow/python/saved_model/README.md

@@ -0,0 +1,241 @@
+# TensorFlow SavedModel
+
+[TOC]
+
+## Overview
+
+SavedModel is the universal serialization format for
+[TensorFlow](https://www.tensorflow.org/) models.
+
+SavedModel provides a language-neutral format to save machine-learning models
+that is recoverable and hermetic. It enables higher-level systems and tools to
+produce, consume and transform TensorFlow models.
+
+## Guides
+* [Using the SavedModel Format](https://www.tensorflow.org/guide/saved_model)
+* [Save and load Keras models](https://www.tensorflow.org/guide/keras/save_and_serialize)
+* [Save and load with checkpointing in Keras](https://www.tensorflow.org/tutorials/keras/save_and_load)
+* [Training checkpoints](https://www.tensorflow.org/guide/checkpoint)
+* [Save and load a model using a distribution strategy](https://www.tensorflow.org/tutorials/distribute/save_and_load)
+
+
+## [Public API](https://www.tensorflow.org/api_docs/python/tf/saved_model)
+* [`tf.saved_model.save`](https://www.tensorflow.org/api_docs/python/tf/saved_model/save)
+* [`tf.saved_model.load`](https://www.tensorflow.org/api_docs/python/tf/saved_model/load)
+* [`tf.saved_model.SaveOptions`](https://www.tensorflow.org/api_docs/python/tf/saved_model/SaveOptions)
+* [`tf.saved_model.LoadOptions`](https://www.tensorflow.org/api_docs/python/tf/saved_model/LoadOptions)
+* [`tf.saved_model.Asset`](https://www.tensorflow.org/api_docs/python/tf/saved_model/Asset)
+* [`tf.saved_model.contains_saved_model`](https://www.tensorflow.org/api_docs/python/tf/saved_model/contains_saved_model)
+
+### Related Modules and Functions
+* [`tf.keras.models.save_model`](https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model)
+* [`tf.keras.models.load_model`](https://www.tensorflow.org/api_docs/python/tf/keras/models/load_model)
+* [`tf.train.Checkpoint`](https://www.tensorflow.org/api_docs/python/tf/train/Checkpoint)
+
+
+## The SavedModel Format
+A SavedModel directory has the following structure:
+
+```
+assets/
+assets.extra/
+variables/
+    variables.data-?????-of-?????
+    variables.index
+saved_model.pb
+```
+
+*   SavedModel protocol buffer
+    *   [`saved_model.pb`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/protobuf/saved_model.proto)
+        or `saved_model.pbtxt`
+    *   Includes the graph definitions as `MetaGraphDef` protocol buffers.
+*   Assets
+    *   Subfolder called `assets`.
+    *   Contains auxiliary files such as vocabularies, etc.
+*   Extra assets
+    *   Subfolder where higher-level libraries and users can add their own
+        assets that co-exist with the model, but are not loaded by the graph.
+    *   This subfolder is not managed by the SavedModel libraries.
+*   Variables
+    *   Subfolder called `variables`.
+        *   `variables.data-?????-of-?????`
+        *   `variables.index`
+
+---
+
+## SavedModel in TensorFlow 1.x
+
+SavedModel had slightly different semantics in TF 1.x. Conventions that are
+generally only supported in TF 1.x are noted as such.
+
+### Features
+
+The following is a summary of the features in SavedModel:
+
+* (TF1-only) Multiple graphs sharing a single set of variables and assets can be added to a
+  single SavedModel. Each graph is associated with a specific set of tags to
+  allow identification during a load or restore operation.
+* (TF1-only) Support for `SignatureDefs`
+    * Graphs that are used for inference tasks typically have a set of inputs
+      and outputs. This is called a `Signature`.
+    * SavedModel uses [SignatureDefs](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/protobuf/meta_graph.proto)
+      to allow generic support for signatures that may need to be saved with the graphs.
+    * For commonly used SignatureDefs in the context of TensorFlow Serving,
+      please see documentation [here](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/g3doc/signature_defs.md).
+* Support for `Assets`.
+    * For cases where ops depend on external files for initialization, such as
+      vocabularies, SavedModel supports this via `assets`.
+    * Assets are copied to the SavedModel location and can be read when loading
+      a specific meta graph def.
+* Support to clear devices before generating the SavedModel.
+
+The following is a summary of features that are NOT supported in SavedModel.
+Higher-level frameworks and tools that use SavedModel may provide these.
+
+* Implicit versioning.
+* Garbage collection.
+* Atomic writes to the SavedModel location.
+
+### TF1 SavedModel Background
+SavedModel manages and builds upon existing TensorFlow primitives such as
+`TensorFlow Saver` and `MetaGraphDef`. Specifically, SavedModel wraps a [TensorFlow Saver](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/training/saver.py).
+The Saver is primarily used to generate the variable checkpoints. SavedModel
+will replace the existing [TensorFlow Inference Model Format](https://github.com/tensorflow/tensorflow/tree/r1.15/tensorflow/contrib/session_bundle#tensorflow-inference-model-format)
+as the canonical way to export TensorFlow graphs for serving.
+
+
+### APIs
+The APIs for building and loading a SavedModel are described in this section.
+
+#### (TF1-only) Builder
+The SavedModel [builder](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/builder.py)
+is implemented in Python.
+
+The `SavedModelBuilder` class provides functionality to save multiple meta graph
+defs, associated variables and assets.
+
+To build a SavedModel, the first meta graph must be saved with variables.
+Subsequent meta graphs will simply be saved with their graph definitions. If
+assets need to be saved and written or copied to disk, they can be provided
+when the meta graph def is added. If multiple meta graph defs are associated
+with an asset of the same name, only the first version is retained.
+
+#### (TF1-only) Tags
+Each meta graph added to the SavedModel must be annotated with user specified
+tags, which reflect the meta graph capabilities or use-cases.
+More specifically, these tags typically annotate a meta graph with its
+functionality (e.g. serving or training), and possibly hardware specific aspects
+such as GPU.
+In the SavedModel, the meta graph def whose tag-set exactly matches those
+specified in the loader API, will be the one loaded by the loader.
+If no meta graph def is found matching the specified tags, an error is returned.
+For example, a loader with a requirement to serve on GPU hardware would be able
+to load only meta graph annotated with tags='serve,gpu' by specifying this set
+of tags in tensorflow::LoadSavedModel(...).
+
+
+#### Usage
+The typical usage of `builder` is as follows:
+
+~~~python
+export_dir = ...
+...
+builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
+with tf.Session(graph=tf.Graph()) as sess:
+  ...
+  builder.add_meta_graph_and_variables(sess,
+                                       [tf.saved_model.tag_constants.TRAINING],
+                                       signature_def_map=foo_signatures,
+                                       assets_collection=foo_assets)
+...
+with tf.Session(graph=tf.Graph()) as sess:
+  ...
+  builder.add_meta_graph(["bar-tag", "baz-tag"])
+...
+builder.save()
+~~~
+
+#### (TF1-only) Stripping Default valued attributes
+The SavedModelBuilder class allows users to control whether default-valued
+attributes must be stripped from the NodeDefs while adding a meta graph to the
+SavedModel bundle. Both `SavedModelBuilder.add_meta_graph_and_variables` and
+`SavedModelBuilder.add_meta_graph` methods accept a Boolean flag
+`strip_default_attrs` that controls this behavior.
+
+If `strip_default_attrs` is `False`, the exported MetaGraphDef will have the
+default valued attributes in all it's NodeDef instances. This can break forward
+compatibility with a sequence of events such as the following:
+
+* An existing Op (`Foo`) is updated to include a new attribute (`T`) with a
+  default (`bool`) at version 101.
+* A model producer (such as a Trainer) binary picks up this change
+  (version 101) to the OpDef and re-exports an existing model that uses Op `Foo`.
+* A model consumer (such as Tensorflow Serving) running an older binary
+  (version 100) doesn't have attribute `T` for Op `Foo`, but tries to import
+  this model. The model consumer doesn't recognize attribute `T` in a NodeDef
+  that uses Op `Foo` and therefore fails to load the model.
+
+By setting `strip_default_attrs` to `True`, the model producers can strip away
+any default valued attributes in the NodeDefs. This helps ensure that newly
+added attributes with defaults don't cause older model consumers to fail loading
+models regenerated with newer training binaries.
+
+TIP: If you care about forward compatibility, then set `strip_default_attrs`
+to `True` while using `SavedModelBuilder.add_meta_graph_and_variables` and
+`SavedModelBuilder.add_meta_graph`.
+
+### Loader
+The SavedModel loader is implemented in C++ and Python.
+
+#### (TF1-only) Python
+The Python version of the SavedModel [loader](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/loader.py)
+provides load and restore capability for a SavedModel. The `load` operation
+requires the session in which to restore the graph definition and variables, the
+tags used to identify the meta graph def to load and the location of the
+SavedModel. Upon a load, the subset of variables and assets supplied as part of
+the specific meta graph def, will be restored into the supplied session.
+
+~~~python
+export_dir = ...
+...
+with tf.Session(graph=tf.Graph()) as sess:
+  tf.saved_model.loader.load(sess, [tag_constants.TRAINING], export_dir)
+  ...
+~~~
+
+#### C++
+The C++ version of the SavedModel [loader](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/saved_model/loader.h)
+provides an API to load a SavedModel from a path, while allowing
+`SessionOptions` and `RunOptions`. Similar to the Python version, the C++
+version requires the tags associated with the graph to be loaded, to be
+specified. The loaded version of SavedModel is referred to as `SavedModelBundle`
+and contains the meta graph def and the session within which it is loaded.
+
+~~~c++
+const string export_dir = ...
+SavedModelBundle bundle;
+...
+LoadSavedModel(session_options, run_options, export_dir, {kSavedModelTagTrain},
+               &bundle);
+~~~
+
+### Constants
+SavedModel offers the flexibility to build and load TensorFlow graphs for a
+variety of use-cases. For the set of most common expected use-cases,
+SavedModel's APIs provide a set of constants in Python and C++ that are easy to
+reuse and share across tools consistently.
+
+#### (TF1-specific) Tag constants
+Sets of tags can be used to uniquely identify a `MetaGraphDef` saved in a
+SavedModel. A subset of commonly used tags is specified in:
+
+* [Python](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/tag_constants.py)
+* [C++](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/saved_model/tag_constants.h).
+
+#### (TF1-specific) Signature constants
+SignatureDefs are used to define the signature of a computation supported in a
+TensorFlow graph. Commonly used input keys, output keys and method names are
+defined in:
+
+* [Python](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/signature_constants.py)
+* [C++](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/cc/saved_model/signature_constants.h).

tensorflow/python/saved_model/builder.py

@@ -0,0 +1,25 @@
+# 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.
+# ==============================================================================
+"""SavedModel builder.
+
+Builds a SavedModel that can be saved to storage, is language neutral, and
+enables systems to produce, consume, or transform TensorFlow Models.
+
+"""
+
+# pylint: disable=unused-import
+from tensorflow.python.saved_model.builder_impl import _SavedModelBuilder
+from tensorflow.python.saved_model.builder_impl import SavedModelBuilder
+# pylint: enable=unused-import

tensorflow/python/saved_model/builder_impl.py

@@ -0,0 +1,861 @@
+# 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.
+# ==============================================================================
+"""SavedModel builder implementation."""
+
+import functools
+import os
+
+from google.protobuf.any_pb2 import Any
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.core.protobuf import saved_model_pb2
+from tensorflow.core.protobuf import saver_pb2
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.saved_model import fingerprinting_utils
+from tensorflow.python.saved_model import path_helpers
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model.pywrap_saved_model import constants
+from tensorflow.python.saved_model.pywrap_saved_model import metrics
+from tensorflow.python.training import saver as tf_saver
+from tensorflow.python.util import compat
+from tensorflow.python.util.deprecation import deprecated_args
+from tensorflow.python.util.tf_export import tf_export
+# copybara:uncomment # Placeholder for protosplitter import.
+# API label for SavedModel metrics.
+_SAVE_BUILDER_LABEL = "save_v1_builder"
+
+
+# Base class for the SavedModelBuilder that is only used by Tensorflow
+# internally. Please use tf.compat.v1.saved_model.SavedModelBuilder instead.
+@tf_export("__internal__.saved_model.SavedModelBuilder", v1=[])
+class _SavedModelBuilder(object):
+  """Builds the `SavedModel` protocol buffer and saves variables and assets.
+
+  The `SavedModelBuilder` class provides the functionality to build a
+  `SavedModel` protocol buffer. Specifically, this allows multiple meta
+  graphs to be saved as part of a single language-neutral `SavedModel`,
+  while sharing variables and assets.
+
+  To build a SavedModel, the first meta graph must be saved with variables.
+  Subsequent meta graphs will simply be saved with their graph definitions. If
+  assets need to be saved and written or copied to disk, they can be provided
+  when the meta graph def is added. If multiple meta graph defs are associated
+  an asset of the same name, only the first version is retained.
+
+  Each meta graph added to the SavedModel must be annotated with tags. The tags
+  provide a means to identify the specific meta graph to load and restore, along
+  with the shared set of variables and assets.
+
+  Typical usage for the `SavedModelBuilder`:
+
+  ```python
+  ...
+  builder = tf.compat.v1.saved_model.Builder(export_dir)
+
+  with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+    ...
+    builder.add_meta_graph_and_variables(sess,
+                                    ["foo-tag"],
+                                    signature_def_map=foo_signatures,
+                                    assets_list=foo_assets)
+  ...
+
+  with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+    ...
+    builder.add_meta_graph(["bar-tag", "baz-tag"])
+  ...
+
+  builder.save()
+  ```
+
+  Note: This function will only be available through the v1 compatibility
+  library as tf.compat.v1.saved_model.builder.SavedModelBuilder or
+  tf.compat.v1.saved_model.Builder. Tensorflow 2.0 will introduce a new
+  object-based method of creating SavedModels.
+  """
+
+  def __init__(self, export_dir):
+    self._saved_model = saved_model_pb2.SavedModel()
+    self._saved_model.saved_model_schema_version = (
+        constants.SAVED_MODEL_SCHEMA_VERSION)
+
+    self._export_dir = export_dir
+    if file_io.file_exists(export_dir):
+      if file_io.list_directory(export_dir):
+        raise AssertionError(
+            f"Export directory {export_dir} already exists, and isn't empty. "
+            "Please choose a different export directory, or delete all the "
+            "contents of the specified directory.")
+    else:
+      file_io.recursive_create_dir(self._export_dir)
+
+    # Boolean to track whether variables and assets corresponding to the
+    # SavedModel have been saved. Specifically, the first meta graph to be added
+    # MUST use the add_meta_graph_and_variables() API. Subsequent add operations
+    # on the SavedModel MUST use the add_meta_graph() API which does not save
+    # weights.
+    self._has_saved_variables = False
+    self._saved_asset_files = set()
+
+  def _save_and_write_assets(self, meta_graph_def, assets_list=None):
+    """Saves asset to the meta graph and writes asset files to disk.
+
+    Args:
+      meta_graph_def: The meta graph def to which the assets will be added.
+      assets_list: The list where the asset paths are setup.
+    """
+    # Creates a function that adds assets into the meta graph def.
+    write_fn = functools.partial(_add_asset_to_metagraph, meta_graph_def)
+    asset_filename_map = _maybe_save_assets(write_fn, assets_list)
+
+    # Return if there are no assets to write.
+    if not asset_filename_map:
+      tf_logging.info("No assets to write.")
+      return
+
+    # Copy assets from source path to destination path.
+    copy_assets_to_destination_dir(asset_filename_map, self._export_dir,
+                                   self._saved_asset_files)
+
+  def _tag_and_add_meta_graph(self, meta_graph_def, tags, signature_def_map):
+    """Tags the meta graph def and adds it to the SavedModel.
+
+    Tags the meta graph def with the supplied tags, adds signature defs to it if
+    provided and appends the meta graph def to the SavedModel proto.
+
+    Args:
+      meta_graph_def: The meta graph def to add to the SavedModel.
+      tags: The set of tags to annotate the meta graph def with.
+      signature_def_map: The map of signature defs to be added to the meta graph
+        def.
+    """
+    for tag in tags:
+      meta_graph_def.meta_info_def.tags.append(tag)
+
+    if signature_def_map is not None:
+      for key in signature_def_map:
+        meta_graph_def.signature_def[key].CopyFrom(signature_def_map[key])
+
+    proto_meta_graph_def = self._saved_model.meta_graphs.add()
+    proto_meta_graph_def.CopyFrom(meta_graph_def)
+
+  def _validate_tensor_info(self, tensor_info):
+    """Validates the `TensorInfo` proto.
+
+    Checks if the `encoding` (`name` or `coo_sparse` or `type_spec`) and
+    `dtype` fields exist and are non-empty.
+
+    Args:
+      tensor_info: `TensorInfo` protocol buffer to validate.
+
+    Raises:
+      AssertionError: If the `encoding` or `dtype` fields of the supplied
+          `TensorInfo` proto are not populated.
+    """
+    if tensor_info is None:
+      raise AssertionError(
+          "All TensorInfo protos used in the SignatureDefs must have the name "
+          "and dtype fields set.")
+    if tensor_info.WhichOneof("encoding") is None:
+      # TODO(soergel) validate each of the fields of coo_sparse
+      raise AssertionError(
+          f"Invalid `tensor_info`: {tensor_info}. All TensorInfo protos used "
+          "in the SignatureDefs must have one of the 'encoding' fields (e.g., "
+          "name or coo_sparse) set.")
+    if tensor_info.WhichOneof("encoding") == "composite_tensor":
+      for component in tensor_info.composite_tensor.components:
+        self._validate_tensor_info(component)
+    elif tensor_info.dtype == types_pb2.DT_INVALID:
+      raise AssertionError(
+          f"Invalid `tensor_info`: {tensor_info}. All TensorInfo protos used in"
+          " the SignatureDefs must have the dtype field set.")
+
+  def _validate_signature_def_map(self, signature_def_map):
+    """Validates the `SignatureDef` entries in the signature def map.
+
+    Validation of entries in the signature def map includes ensuring that the
+    `name` and `dtype` fields of the TensorInfo protos of the `inputs` and
+    `outputs` of each `SignatureDef` are populated. Also ensures that reserved
+    SignatureDef keys for the initialization and train ops are not used.
+
+    Args:
+      signature_def_map: The map of signature defs to be validated.
+
+    Raises:
+      AssertionError: If a TensorInfo is not valid.
+      KeyError: If a reserved signature key is used in the map.
+    """
+    for signature_def_key in signature_def_map:
+      signature_def = signature_def_map[signature_def_key]
+      inputs = signature_def.inputs
+      outputs = signature_def.outputs
+      for inputs_key in inputs:
+        self._validate_tensor_info(inputs[inputs_key])
+      for outputs_key in outputs:
+        self._validate_tensor_info(outputs[outputs_key])
+    if constants.INIT_OP_SIGNATURE_KEY in signature_def_map:
+      raise KeyError(
+          f"SignatureDef map key \"{constants.INIT_OP_SIGNATURE_KEY}\" is "
+          "reserved for initialization. Please use a different key.")
+    if constants.TRAIN_OP_SIGNATURE_KEY in signature_def_map:
+      raise KeyError(
+          f"SignatureDef map key \"{constants.TRAIN_OP_SIGNATURE_KEY}\" is "
+          f"reserved for the train op. Please use a different key.")
+
+  def _maybe_create_saver(self, saver=None):
+    """Creates a sharded saver if one does not already exist."""
+    if not saver:
+      # Initialize a saver to generate a sharded output for all saveables in the
+      # current scope.
+      saver = tf_saver.Saver(
+          variables._all_saveable_objects(),  # pylint: disable=protected-access
+          sharded=True,
+          write_version=saver_pb2.SaverDef.V2,
+          allow_empty=True)
+    return saver
+
+  def add_meta_graph(self,
+                     tags,
+                     signature_def_map=None,
+                     assets_list=None,
+                     clear_devices=False,
+                     init_op=None,
+                     train_op=None,
+                     saver=None):
+    """Adds the current meta graph to the SavedModel.
+
+    Creates a Saver in the current scope and uses the Saver to export the meta
+    graph def. Invoking this API requires the `add_meta_graph_and_variables()`
+    API to have been invoked before.
+
+    Args:
+      tags: The set of tags to annotate the meta graph def with.
+      signature_def_map: The map of signature defs to be added to the meta graph
+        def.
+      assets_list: Assets to be saved with SavedModel. Note
+          that this list should be a subset of the assets saved as part of
+          the first meta graph in the SavedModel.
+      clear_devices: Set to true if the device info on the default graph should
+        be cleared.
+      init_op: Op or group of ops to execute when the graph is loaded. Note
+          that when the init_op is specified it is run after the restore op at
+        load-time.
+      train_op: Op or group of opts that trains the model when run. This will
+        not be run automatically when the graph is loaded, instead saved in
+        a SignatureDef accessible through the exported MetaGraph.
+      saver: An instance of tf.compat.v1.train.Saver that will be used to export
+        the metagraph. If None, a sharded Saver that restores all variables will
+        be used.
+
+    Raises:
+      AssertionError: If the variables for the SavedModel have not been saved
+          yet, or if the graph already contains one or more legacy init ops.
+    """
+    if not self._has_saved_variables:
+      raise AssertionError(
+          "Graph state including variables and assets has not been saved yet. "
+          "Please invoke `add_meta_graph_and_variables()` first.")
+
+    # Validate the signature def map to ensure all included TensorInfos are
+    # properly populated.
+    signature_def_map = signature_def_map or {}
+    self._validate_signature_def_map(signature_def_map)
+
+    # Create a SignatureDef pointing to the graph initialization op, which will
+    # be added to the MetaGraphDef.
+    _add_op_to_signature_def_map(signature_def_map, init_op,
+                                 constants.INIT_OP_SIGNATURE_KEY)
+    _add_op_to_signature_def_map(signature_def_map, train_op,
+                                 constants.TRAIN_OP_SIGNATURE_KEY)
+
+    saver = self._maybe_create_saver(saver)
+
+    # The graph almost certainly previously contained at least one Saver, and
+    # possibly several (e.g. one for loading a pretrained embedding, and another
+    # for the model weights).  Removing the preexisting ones was the
+    # motivation for the clear_extraneous_savers option, but it turns out that
+    # there are edge cases where that option breaks the graph.  Until that is
+    # resolved, we just leave the option set to False for now.
+    # TODO(soergel): Reinstate clear_extraneous_savers=True when possible.
+    meta_graph_def = saver.export_meta_graph(
+        clear_devices=clear_devices, strip_default_attrs=True)
+
+    # Save asset files and write them to disk, if any.
+    self._save_and_write_assets(meta_graph_def, assets_list)
+
+    # Tag the meta graph def and add it to the SavedModel.
+    self._tag_and_add_meta_graph(meta_graph_def, tags, signature_def_map)
+
+  def add_meta_graph_and_variables(self,
+                                   sess,
+                                   tags,
+                                   signature_def_map=None,
+                                   assets_list=None,
+                                   clear_devices=False,
+                                   init_op=None,
+                                   train_op=None,
+                                   strip_default_attrs=False,
+                                   saver=None):
+    # pylint: disable=line-too-long
+    """Adds the current meta graph to the SavedModel and saves variables.
+
+    Creates a Saver to save the variables from the provided session. Exports the
+    corresponding meta graph def. This function assumes that the variables to be
+    saved have been initialized. For a given `SavedModelBuilder`, this API must
+    be called exactly once and for the first meta graph to save. For subsequent
+    meta graph defs to be added, the `add_meta_graph()` API must be used.
+
+    Args:
+      sess: The TensorFlow session from which to save the meta graph and
+        variables.
+      tags: The set of tags with which to save the meta graph.
+      signature_def_map: The map of signature def map to add to the meta graph
+        def.
+      assets_list: Assets to be saved with SavedModel.
+      clear_devices: Set to true if the device info on the default graph should
+        be cleared.
+      init_op: Op or group of ops to execute when the graph is loaded. Note
+          that when the init_op is specified it is run after the restore op at
+        load-time.
+      train_op: Op or group of ops that trains the model when run. This will
+        not be run automatically when the graph is loaded, instead saved in
+        a SignatureDef accessible through the exported MetaGraph.
+      strip_default_attrs: Boolean. If `True`, default-valued attributes will be
+        removed from the NodeDefs. For a detailed guide, see
+        [Stripping Default-Valued Attributes](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md#stripping-default-valued-attributes).
+      saver: An instance of tf.compat.v1.train.Saver that will be used to export the
+        metagraph and save variables. If None, a sharded Saver that restores
+        all variables will be used.
+
+    """
+    # pylint: enable=line-too-long
+    if self._has_saved_variables:
+      raise AssertionError("Graph state including variables and assets has "
+                           "already been saved. Please invoke "
+                           "`add_meta_graph()` instead.")
+
+    # Validate the signature def map to ensure all included TensorInfos are
+    # properly populated.
+    signature_def_map = signature_def_map or {}
+    self._validate_signature_def_map(signature_def_map)
+
+    # Create a SignatureDef pointing to the graph initialization op, which will
+    # be added to the MetaGraphDef.
+    _add_op_to_signature_def_map(signature_def_map, init_op,
+                                 constants.INIT_OP_SIGNATURE_KEY)
+    _add_op_to_signature_def_map(signature_def_map, train_op,
+                                 constants.TRAIN_OP_SIGNATURE_KEY)
+
+    path_helpers.get_or_create_variables_dir(self._export_dir)
+    variables_path = path_helpers.get_variables_path(self._export_dir)
+
+    saver = self._maybe_create_saver(saver)
+
+    # Save the variables. Also, disable writing the checkpoint state proto. The
+    # file is not used during SavedModel loading. In addition, since a
+    # SavedModel can be copied or moved, this avoids the checkpoint state to
+    # become outdated.
+    saver.save(sess, variables_path, write_meta_graph=False, write_state=False)
+
+    # Export the meta graph def.
+
+    # The graph almost certainly previously contained at least one Saver, and
+    # possibly several (e.g. one for loading a pretrained embedding, and another
+    # for the model weights).  Removing the preexisting ones was the
+    # motivation for the clear_extraneous_savers option, but it turns out that
+    # there are edge cases where that option breaks the graph.  Until that is
+    # resolved, we just leave the option set to False for now.
+    # TODO(soergel): Reinstate clear_extraneous_savers=True when possible.
+    meta_graph_def = saver.export_meta_graph(
+        clear_devices=clear_devices, strip_default_attrs=strip_default_attrs)
+
+    # Save asset files and write them to disk, if any.
+    self._save_and_write_assets(meta_graph_def, assets_list)
+
+    # Tag the meta graph def and add it to the SavedModel.
+    self._tag_and_add_meta_graph(meta_graph_def, tags, signature_def_map)
+
+    # Mark this instance of SavedModel as having saved variables, such that
+    # subsequent attempts to save variables will fail.
+    self._has_saved_variables = True
+
+  def save(self, as_text=False, experimental_image_format=False,
+           experimental_image_writer_options=None):
+    """Writes a `SavedModel` protocol buffer to disk.
+
+    The function writes the SavedModel protocol buffer to the export directory
+    in a serialized format.
+
+    Args:
+      as_text: Writes the SavedModel protocol buffer in text format to disk.
+        Protocol buffers in text format are useful for debugging, but parsing
+        fails when it encounters an unknown field and so is not forward
+        compatible. This means changes to TensorFlow may prevent deployment of
+        new text format SavedModels to existing serving binaries. Do not deploy
+        `as_text` SavedModels to production.
+      experimental_image_format: Writes the SavedModel protobuf in the
+        experimental image format. See
+      https://www.tensorflow.org/api_docs/python/tf/saved_model/SaveOptions for
+        more details. This allows `SavedModelBuilder` to save models larger than
+        2 GiB.
+      experimental_image_writer_options: Optional options for the experimental
+        image writer. See
+        https://github.com/google/riegeli/blob/master/doc/record_writer_options.md
+          for available options.
+
+    Raises:
+       RuntimeError: When trying to use `proto_splitter` but `proto_splitter` is
+         not imported. This check is here because `proto_splitter` is not
+         available in OSS at the moment.
+
+    Returns:
+      The path to which the SavedModel protocol buffer was written.
+    """
+    metrics.IncrementWriteApi(_SAVE_BUILDER_LABEL)
+    if not file_io.file_exists(self._export_dir):
+      file_io.recursive_create_dir(self._export_dir)
+
+    if as_text:
+      path = file_io.join(
+          compat.as_bytes(self._export_dir),
+          compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
+      file_io.write_string_to_file(path, str(self._saved_model))
+    else:
+      if experimental_image_format:
+        path = file_io.join(
+            self._export_dir,
+            constants.SAVED_MODEL_FILENAME_PREFIX,
+        )
+        if (
+            locals().get("proto_splitter", globals().get("proto_splitter"))
+            is None
+        ):
+          raise RuntimeError(
+              "No proto_splitter is provided, cannot use"
+              " experimental_image_format."
+          )
+        # Overwrites path to record whether the saved_model is split, i.e.,
+        # whether the suffix is `.pb` or `.cpb`.
+        path = proto_splitter.SavedModelSplitter(self._saved_model).write(
+            path, experimental_image_writer_options
+        )
+      else:
+        path = file_io.join(
+            compat.as_bytes(self._export_dir),
+            compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB),
+        )
+        file_io.write_string_to_file(
+            path, self._saved_model.SerializeToString(deterministic=True)
+        )
+    # Placeholder for internal TF1 model fingerprint write
+    tf_logging.info("SavedModel written to: %s", compat.as_text(path))
+    metrics.IncrementWrite(write_version="1")
+
+    return path
+
+
+@tf_export(v1=["saved_model.Builder", "saved_model.builder.SavedModelBuilder"])  # pylint: disable=missing-docstring
+class SavedModelBuilder(_SavedModelBuilder):
+  __doc__ = (_SavedModelBuilder.__doc__ or "").replace(
+      "assets_list", "assets_collection"
+  )
+
+  def __init__(self, export_dir):
+    super(SavedModelBuilder, self).__init__(export_dir=export_dir)
+
+  def _add_collections(self, assets_collection, main_op, train_op):
+    """Add asset and op collections to be saved."""
+    # Save asset files and write them to disk, if any.
+    self._save_and_write_assets(assets_collection)
+
+    self._maybe_add_main_op(main_op)
+
+    self._add_train_op(train_op)
+
+  def _save_and_write_assets(self, assets_collection_to_add=None):
+    """Saves asset to the meta graph and writes asset files to disk.
+
+    Args:
+      assets_collection_to_add: The collection where the asset paths are setup.
+    """
+    # Add assets to the collection with key `saved_model.ASSETS_KEY`, in the
+    # graph.
+    asset_filename_map = _maybe_save_assets(_add_asset_to_collection,
+                                            assets_collection_to_add)
+
+    # Return if there are no assets to write.
+    if not asset_filename_map:
+      tf_logging.info("No assets to write.")
+      return
+
+    # Copy assets from source path to destination path.
+    copy_assets_to_destination_dir(asset_filename_map, self._export_dir,
+                                   self._saved_asset_files)
+
+  def _maybe_add_main_op(self, main_op):
+    """Adds main op to the SavedModel.
+
+    Args:
+      main_op: Main op to run as part of graph initialization. If None, no main
+        op will be added to the graph.
+
+    Raises:
+      TypeError: If the main op is provided but is not of type `Operation`.
+      ValueError: if the Graph already contains an init op.
+    """
+    if main_op is None:
+      return
+
+    if not isinstance(main_op, ops.Operation):
+      raise TypeError(f"Expected {main_op} to be an Operation but got type "
+                      f"{type(main_op)} instead.")
+
+    # Validate that no other init ops have been added to this graph already.
+    # We check main_op and legacy_init_op for thoroughness and explicitness.
+    for init_op_key in (constants.MAIN_OP_KEY, constants.LEGACY_INIT_OP_KEY):
+      if ops.get_collection(init_op_key):
+        raise ValueError(
+            "Graph already contains one or more main ops under the "
+            f"collection {init_op_key}.")
+
+    ops.add_to_collection(constants.MAIN_OP_KEY, main_op)
+
+  def _add_train_op(self, train_op):
+    """Add train op to the SavedModel.
+
+    Note that this functionality is in development, and liable to be
+    moved elsewhere.
+
+    Args:
+      train_op: Op or group of ops that are used for training. These are stored
+        as a collection with key TRAIN_OP_KEY, but not executed.
+
+    Raises:
+      TypeError if Train op is not of type `Operation`.
+    """
+    if train_op is not None:
+      if (not isinstance(train_op, tensor.Tensor) and
+          not isinstance(train_op, ops.Operation)):
+        raise TypeError(f"`train_op` {train_op} needs to be a Tensor or Op.")
+      ops.add_to_collection(constants.TRAIN_OP_KEY, train_op)
+
+  @deprecated_args(None,
+                   "Pass your op to the equivalent parameter main_op instead.",
+                   "legacy_init_op")
+  def add_meta_graph(self,
+                     tags,
+                     signature_def_map=None,
+                     assets_collection=None,
+                     legacy_init_op=None,
+                     clear_devices=False,
+                     main_op=None,
+                     strip_default_attrs=False,
+                     saver=None):
+    if not self._has_saved_variables:
+      raise AssertionError(
+          "Graph state including variables and assets has not been saved yet. "
+          "Please invoke `add_meta_graph_and_variables()` first.")
+
+    # Validate the signature def map to ensure all included TensorInfos are
+    # properly populated.
+    signature_def_map = signature_def_map or {}
+    self._validate_signature_def_map(signature_def_map)
+
+    # legacy_init_op is deprecated, and going away in TF 2.0.
+    # Re-mapping to main_op, as treatment is identical regardless.
+    main_op = main_op if main_op is not None else legacy_init_op
+
+    # Add assets and ops
+    self._add_collections(assets_collection, main_op, None)
+
+    saver = self._maybe_create_saver(saver)
+
+    # The graph almost certainly previously contained at least one Saver, and
+    # possibly several (e.g. one for loading a pretrained embedding, and another
+    # for the model weights).  Removing the preexisting ones was the
+    # motivation for the clear_extraneous_savers option, but it turns out that
+    # there are edge cases where that option breaks the graph.  Until that is
+    # resolved, we just leave the option set to False for now.
+    # TODO(soergel): Reinstate clear_extraneous_savers=True when possible.
+    meta_graph_def = saver.export_meta_graph(
+        clear_devices=clear_devices, strip_default_attrs=strip_default_attrs)
+
+    # Tag the meta graph def and add it to the SavedModel.
+    self._tag_and_add_meta_graph(meta_graph_def, tags, signature_def_map)
+
+  @deprecated_args(None,
+                   "Pass your op to the equivalent parameter main_op instead.",
+                   "legacy_init_op")
+  def add_meta_graph_and_variables(self,
+                                   sess,
+                                   tags,
+                                   signature_def_map=None,
+                                   assets_collection=None,
+                                   legacy_init_op=None,
+                                   clear_devices=False,
+                                   main_op=None,
+                                   strip_default_attrs=False,
+                                   saver=None):
+    if self._has_saved_variables:
+      raise AssertionError("Graph state including variables and assets has "
+                           "already been saved. Please invoke "
+                           "`add_meta_graph()` instead.")
+
+    # Validate the signature def map to ensure all included TensorInfos are
+    # properly populated.
+    signature_def_map = signature_def_map or {}
+    self._validate_signature_def_map(signature_def_map)
+
+    # legacy_init_op is deprecated, and going away in TF 2.0.
+    # Re-mapping to main_op, as treatment is identical regardless.
+    main_op = main_op or legacy_init_op
+
+    # Add assets and ops
+    self._add_collections(assets_collection, main_op, None)
+
+    path_helpers.get_or_create_variables_dir(self._export_dir)
+    variables_path = path_helpers.get_variables_path(self._export_dir)
+
+    saver = self._maybe_create_saver(saver)
+
+    # Save the variables. Also, disable writing the checkpoint state proto. The
+    # file is not used during SavedModel loading. In addition, since a
+    # SavedModel can be copied or moved, this avoids the checkpoint state to
+    # become outdated.
+    saver.save(sess, variables_path, write_meta_graph=False, write_state=False)
+
+    # Export the meta graph def.
+
+    # The graph almost certainly previously contained at least one Saver, and
+    # possibly several (e.g. one for loading a pretrained embedding, and another
+    # for the model weights).  Removing the preexisting ones was the
+    # motivation for the clear_extraneous_savers option, but it turns out that
+    # there are edge cases where that option breaks the graph.  Until that is
+    # resolved, we just leave the option set to False for now.
+    # TODO(soergel): Reinstate clear_extraneous_savers=True when possible.
+    meta_graph_def = saver.export_meta_graph(
+        clear_devices=clear_devices, strip_default_attrs=strip_default_attrs)
+
+    # Tag the meta graph def and add it to the SavedModel.
+    self._tag_and_add_meta_graph(meta_graph_def, tags, signature_def_map)
+
+    # Mark this instance of SavedModel as having saved variables, such that
+    # subsequent attempts to save variables will fail.
+    self._has_saved_variables = True
+
+  add_meta_graph.__doc__ = (
+      _SavedModelBuilder.add_meta_graph.__doc__ or ""
+  ).replace("assets_list", "assets_collection")
+  add_meta_graph_and_variables.__doc__ = (
+      _SavedModelBuilder.add_meta_graph_and_variables.__doc__ or ""
+  ).replace("assets_list", "assets_collection")
+
+
+def _maybe_save_assets(write_fn, assets_to_add=None):
+  """Saves assets to the meta graph.
+
+  Args:
+    write_fn: A function callback that writes assets into meta graph.
+    assets_to_add: The list where the asset paths are setup.
+
+  Returns:
+    A dict of asset basenames for saving to the original full path to the asset.
+
+  Raises:
+    ValueError: Indicating an invalid filepath tensor.
+  """
+  # Map of target file names to original filenames
+  asset_filename_map = {}
+
+  if assets_to_add is None:
+    tf_logging.info("No assets to save.")
+    return asset_filename_map
+
+  # Iterate over the supplied assets, build the `AssetFile` proto and add them
+  # to the meta graph.
+  for asset_tensor in assets_to_add:
+    asset_source_filepath = _asset_path_from_tensor(asset_tensor)
+    if not asset_source_filepath:
+      raise ValueError(f"Asset filepath tensor {asset_tensor} in is invalid.")
+
+    asset_filename = get_asset_filename_to_add(
+        asset_source_filepath, asset_filename_map)
+
+    # Call the passed-in function that builds AssetFileDef proto and adds it
+    # to either the collection or asset_file_def field of the meta graph.
+    # Note that this should be done even when the file is a duplicate of an
+    # already-added file, as the tensor reference should still exist.
+    write_fn(asset_filename, asset_tensor)
+
+    # In the cases where we are adding a duplicate, this will result in the
+    # last of the filepaths being the one used for copying the file to the
+    # SavedModel. Since the files in question are the same, it doesn't matter
+    # either way.
+    asset_filename_map[asset_filename] = asset_source_filepath
+
+  tf_logging.info("Assets added to graph.")
+  return asset_filename_map
+
+
+def get_asset_filename_to_add(asset_filepath, asset_filename_map):
+  """Get a unique basename to add to the SavedModel if this file is unseen.
+
+  Assets come from users as full paths, and we save them out to the
+  SavedModel as basenames. In some cases, the basenames collide. Here,
+  we dedupe asset basenames by first checking if the file is the same,
+  and, if different, generate and return an index-suffixed basename
+  that can be used to add the asset to the SavedModel.
+
+  Args:
+    asset_filepath: the full path to the asset that is being saved
+    asset_filename_map: a dict of filenames used for saving the asset in
+      the SavedModel to full paths from which the filenames were derived.
+
+  Returns:
+    Uniquified filename string if the file is not a duplicate, or the original
+    filename if the file has already been seen and saved.
+  """
+  asset_filename = os.path.basename(asset_filepath)
+
+  if asset_filename not in asset_filename_map:
+    # This is an unseen asset. Safe to add.
+    return asset_filename
+
+  other_asset_filepath = asset_filename_map[asset_filename]
+  if other_asset_filepath == asset_filepath:
+    # This is the same file, stored twice in the list. No need
+    # to make unique.
+    return asset_filename
+
+  # Else, asset_filename is in the map, and the filepath is different. Dedupe.
+  if not file_io.filecmp(asset_filepath, other_asset_filepath):
+    # Files are different; dedupe filenames.
+    return _get_unique_asset_filename(asset_filename, asset_filename_map)
+
+  # Files are the same; don't make unique.
+  return asset_filename
+
+
+def _get_unique_asset_filename(asset_filename, asset_filename_map):
+  i = 1
+  unique_filename = asset_filename
+  while unique_filename in asset_filename_map:
+    unique_filename = compat.as_bytes("_").join(
+        [compat.as_bytes(asset_filename), compat.as_bytes(str(i))])
+    i += 1
+  return unique_filename
+
+
+def _asset_path_from_tensor(path_tensor):
+  """Returns the filepath value stored in constant `path_tensor`.
+
+  Args:
+    path_tensor: Tensor of a file-path.
+
+  Returns:
+    The string value i.e. path of the tensor, if valid.
+
+  Raises:
+    TypeError if tensor does not match expected op type, dtype or value.
+  """
+  if not isinstance(path_tensor, tensor.Tensor):
+    raise TypeError(f"Asset path tensor {path_tensor} must be a Tensor.")
+  if path_tensor.op.type != "Const":
+    raise TypeError(f"Asset path tensor {path_tensor} must be of type constant."
+                    f"Has type {path_tensor.op.type} instead.")
+  if path_tensor.dtype != dtypes.string:
+    raise TypeError(f"Asset path tensor {path_tensor}` must be of dtype string."
+                    f"Has type {path_tensor.dtype} instead.")
+  str_values = path_tensor.op.get_attr("value").string_val
+  if len(str_values) != 1:
+    raise TypeError(f"Asset path tensor {path_tensor} must be a scalar.")
+  return str_values[0]
+
+
+def _add_asset_to_metagraph(meta_graph_def, asset_filename, asset_tensor):
+  """Builds an asset proto and adds it to the meta graph def.
+
+  Args:
+    meta_graph_def: The meta graph def to which the asset will be added.
+    asset_filename: The filename of the asset to be added.
+    asset_tensor: The asset tensor used to populate the tensor info of the asset
+      proto.
+  """
+  asset_proto = meta_graph_def.asset_file_def.add()
+  asset_proto.filename = asset_filename
+  asset_proto.tensor_info.name = asset_tensor.name
+
+
+def copy_assets_to_destination_dir(asset_filename_map, destination_dir,
+                                   saved_files=None):
+  """Copy all assets from source path to destination path.
+
+  Args:
+    asset_filename_map: a dict of filenames used for saving the asset in
+      the SavedModel to full paths from which the filenames were derived.
+    destination_dir: the destination directory that assets are stored in.
+    saved_files: a set of destination filepaths that have already been copied
+      and will be skipped
+  """
+  if saved_files is None:
+    saved_files = set()
+
+  assets_destination_dir = path_helpers.get_or_create_assets_dir(
+      destination_dir)
+
+  # Copy each asset from source path to destination path.
+  for asset_basename, asset_source_filepath in asset_filename_map.items():
+    asset_destination_filepath = file_io.join(
+        compat.as_bytes(assets_destination_dir),
+        compat.as_bytes(asset_basename))
+
+    # Copy if source file exists, src & dst are not the same, and dst is not in
+    # saved_files
+    if (file_io.file_exists(asset_source_filepath) and
+        asset_source_filepath != asset_destination_filepath and
+        asset_destination_filepath not in saved_files):
+      file_io.copy(
+          asset_source_filepath, asset_destination_filepath, overwrite=True)
+      saved_files.add(asset_destination_filepath)
+
+  tf_logging.info("Assets written to: %s",
+                  compat.as_text(assets_destination_dir))
+
+
+def _add_asset_to_collection(asset_filename, asset_tensor):
+  """Builds an asset proto and adds it to the asset collection of the graph.
+
+  Args:
+    asset_filename: The filename of the asset to be added.
+    asset_tensor: The asset tensor used to populate the tensor info of the
+        asset proto.
+  """
+  asset_proto = meta_graph_pb2.AssetFileDef()
+  asset_proto.filename = asset_filename
+  asset_proto.tensor_info.name = asset_tensor.name
+
+  asset_any_proto = Any()
+  asset_any_proto.Pack(asset_proto)
+  ops.add_to_collection(constants.ASSETS_KEY, asset_any_proto)
+
+
+def _add_op_to_signature_def_map(signature_def_map, op, key):
+  if op is not None:
+    signature_def_map[key] = signature_def_utils.op_signature_def(op, key)

tensorflow/python/saved_model/constants.py

@@ -0,0 +1,135 @@
+# Copyright 2016 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 for SavedModel save and restore operations.
+
+The source of truth for these constants is in
+tensorflow/cc/saved_model/constants.h.
+
+"""
+from tensorflow.python.saved_model.pywrap_saved_model import constants
+from tensorflow.python.util.tf_export import tf_export
+
+# Subdirectory name containing the asset files.
+ASSETS_DIRECTORY = constants.ASSETS_DIRECTORY
+tf_export(
+    "saved_model.ASSETS_DIRECTORY",
+    v1=[
+        "saved_model.ASSETS_DIRECTORY", "saved_model.constants.ASSETS_DIRECTORY"
+    ]).export_constant(__name__, "ASSETS_DIRECTORY")
+
+# Subdirectory name containing unmanaged files from higher-level APIs.
+EXTRA_ASSETS_DIRECTORY = constants.EXTRA_ASSETS_DIRECTORY
+
+# CollectionDef key containing SavedModel assets.
+ASSETS_KEY = constants.ASSETS_KEY
+tf_export(
+    "saved_model.ASSETS_KEY",
+    v1=["saved_model.ASSETS_KEY",
+        "saved_model.constants.ASSETS_KEY"]).export_constant(
+            __name__, "ASSETS_KEY")
+
+# CollectionDef key for the legacy init op.
+LEGACY_INIT_OP_KEY = constants.LEGACY_INIT_OP_KEY
+tf_export(
+    v1=[
+        "saved_model.LEGACY_INIT_OP_KEY",
+        "saved_model.constants.LEGACY_INIT_OP_KEY"
+    ]).export_constant(__name__, "LEGACY_INIT_OP_KEY")
+
+# CollectionDef key for the SavedModel main op.
+MAIN_OP_KEY = constants.MAIN_OP_KEY
+tf_export(
+    v1=["saved_model.MAIN_OP_KEY",
+        "saved_model.constants.MAIN_OP_KEY"]).export_constant(
+            __name__, "MAIN_OP_KEY")
+
+# CollectionDef key for the SavedModel train op.
+# Not exported while export_all_saved_models is experimental.
+TRAIN_OP_KEY = constants.TRAIN_OP_KEY
+
+# Schema version for SavedModel.
+SAVED_MODEL_SCHEMA_VERSION = constants.SAVED_MODEL_SCHEMA_VERSION
+tf_export(
+    "saved_model.SAVED_MODEL_SCHEMA_VERSION",
+    v1=[
+        "saved_model.SAVED_MODEL_SCHEMA_VERSION",
+        "saved_model.constants.SAVED_MODEL_SCHEMA_VERSION"
+    ]).export_constant(__name__, "SAVED_MODEL_SCHEMA_VERSION")
+
+# File name prefix for SavedModel protocol buffer.
+SAVED_MODEL_FILENAME_PREFIX = constants.SAVED_MODEL_FILENAME_PREFIX
+
+# File name for SavedModel protocol buffer.
+SAVED_MODEL_FILENAME_PB = constants.SAVED_MODEL_FILENAME_PB
+tf_export(
+    "saved_model.SAVED_MODEL_FILENAME_PB",
+    v1=[
+        "saved_model.SAVED_MODEL_FILENAME_PB",
+        "saved_model.constants.SAVED_MODEL_FILENAME_PB"
+    ]).export_constant(__name__, "SAVED_MODEL_FILENAME_PB")
+
+# File name for SavedModel chunked protocol buffer (experimental).
+SAVED_MODEL_FILENAME_CPB = constants.SAVED_MODEL_FILENAME_CPB
+
+# File name for text version of SavedModel protocol buffer.
+SAVED_MODEL_FILENAME_PBTXT = constants.SAVED_MODEL_FILENAME_PBTXT
+tf_export(
+    "saved_model.SAVED_MODEL_FILENAME_PBTXT",
+    v1=[
+        "saved_model.SAVED_MODEL_FILENAME_PBTXT",
+        "saved_model.constants.SAVED_MODEL_FILENAME_PBTXT"
+    ]).export_constant(__name__, "SAVED_MODEL_FILENAME_PBTXT")
+
+# Subdirectory where debugging related files are written.
+DEBUG_DIRECTORY = constants.DEBUG_DIRECTORY
+tf_export(
+    "saved_model.DEBUG_DIRECTORY",
+    v1=[
+        "saved_model.DEBUG_DIRECTORY",
+        "saved_model.constants.DEBUG_DIRECTORY",
+    ]).export_constant(__name__, "DEBUG_DIRECTORY")
+
+# File name for GraphDebugInfo protocol buffer which corresponds to the
+# SavedModel.
+DEBUG_INFO_FILENAME_PB = constants.DEBUG_INFO_FILENAME_PB
+tf_export(
+    "saved_model.DEBUG_INFO_FILENAME_PB",
+    v1=[
+        "saved_model.DEBUG_INFO_FILENAME_PB",
+        "saved_model.constants.DEBUG_INFO_FILENAME_PB"
+    ]).export_constant(__name__, "DEBUG_INFO_FILENAME_PB")
+
+# Subdirectory name containing the variables/checkpoint files.
+VARIABLES_DIRECTORY = constants.VARIABLES_DIRECTORY
+tf_export(
+    "saved_model.VARIABLES_DIRECTORY",
+    v1=[
+        "saved_model.VARIABLES_DIRECTORY",
+        "saved_model.constants.VARIABLES_DIRECTORY"
+    ]).export_constant(__name__, "VARIABLES_DIRECTORY")
+
+# File name used for variables.
+VARIABLES_FILENAME = constants.VARIABLES_FILENAME
+tf_export(
+    "saved_model.VARIABLES_FILENAME",
+    v1=[
+        "saved_model.VARIABLES_FILENAME",
+        "saved_model.constants.VARIABLES_FILENAME"
+    ]).export_constant(__name__, "VARIABLES_FILENAME")
+
+# The initialization and train ops for a MetaGraph are stored in the
+# signature def map. The ops are added to the map with the following keys.
+INIT_OP_SIGNATURE_KEY = constants.INIT_OP_SIGNATURE_KEY
+TRAIN_OP_SIGNATURE_KEY = constants.TRAIN_OP_SIGNATURE_KEY

tensorflow/python/saved_model/fingerprinting.md

@@ -0,0 +1,26 @@
+## SavedModel Fingerprinting
+
+This document describes the implementation details of SavedModel fingerprinting. 
+The design document (RFC) can be found [here](https://github.com/tensorflow/community/pull/415).
+
+### Implementation
+
+The code that implements SavedModel fingerprinting can be found in :
+- `tensorflow/python/saved_model/fingerprinting.py`: Public python methods for accessing the fingerprint.
+- `tensorflow/python/saved_model/pywrap_saved_model_fingerprinting.*`: Python wrappers for C++ fingerprint methods. For internal use only.
+- `tensorflow/cc/saved_model/fingerprint.*`: C++ methods for creating and reading the fingerprint.
+- `tensorflow/core/graph/regularization/`: Code that "regularizes" the GraphDef. See the README
+
+Generally speaking, most of the implementation for SavedModel fingerprinting is in C++. The code in this directory is meant to make these methods accessible in Python for the purposes of creating a public API
+as well as instrumenting the Python side of the code base.
+
+### Instrumentation
+
+The current SavedModel reading and loading APIs are instrumented such that they log
+the fingerprint every time they are called. The APIs that are instrumented are:
+- `tf.saved_model.save`
+- `tf.saved_model.load`
+- `tensorflow::LoadSavedModel`
+- `tensorflow::SavedModelV2Bundle::Load`
+
+

tensorflow/python/saved_model/fingerprinting.py

@@ -0,0 +1,178 @@
+# 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.
+# ==============================================================================
+"""Methods for SavedModel fingerprinting.
+
+This module contains classes and functions for reading the SavedModel
+fingerprint.
+"""
+
+from typing import Any
+
+from tensorflow.core.protobuf import fingerprint_pb2
+from tensorflow.python.saved_model.pywrap_saved_model import fingerprinting as fingerprinting_pywrap
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("saved_model.experimental.Fingerprint", v1=[])
+class Fingerprint:
+  """The SavedModel fingerprint.
+
+  Each attribute of this class is named after a field name in the
+  FingerprintDef proto and contains the value of the respective field in the
+  protobuf.
+
+  Attributes:
+    saved_model_checksum: A uint64 containing the `saved_model_checksum`.
+    graph_def_program_hash: A uint64 containing `graph_def_program_hash`.
+    signature_def_hash: A uint64 containing the `signature_def_hash`.
+    saved_object_graph_hash: A uint64 containing the `saved_object_graph_hash`.
+    checkpoint_hash: A uint64 containing the`checkpoint_hash`.
+    version: An int32 containing the producer field of the VersionDef.
+  """
+
+  def __init__(
+      self,
+      saved_model_checksum: int = None,
+      graph_def_program_hash: int = None,
+      signature_def_hash: int = None,
+      saved_object_graph_hash: int = None,
+      checkpoint_hash: int = None,
+      version: int = None,
+  ):
+    """Initializes the instance based on values in the SavedModel fingerprint.
+
+    Args:
+      saved_model_checksum: Value of the`saved_model_checksum`.
+      graph_def_program_hash: Value of the `graph_def_program_hash`.
+      signature_def_hash: Value of the `signature_def_hash`.
+      saved_object_graph_hash: Value of the `saved_object_graph_hash`.
+      checkpoint_hash: Value of the `checkpoint_hash`.
+      version: Value of the producer field of the VersionDef.
+    """
+    self.saved_model_checksum = saved_model_checksum
+    self.graph_def_program_hash = graph_def_program_hash
+    self.signature_def_hash = signature_def_hash
+    self.saved_object_graph_hash = saved_object_graph_hash
+    self.checkpoint_hash = checkpoint_hash
+    self.version = version
+
+  @classmethod
+  def from_proto(cls, proto: fingerprint_pb2.FingerprintDef) -> "Fingerprint":
+    """Constructs Fingerprint object from protocol buffer message."""
+    if isinstance(proto, bytes):
+      proto = fingerprint_pb2.FingerprintDef.FromString(proto)
+    try:
+      return Fingerprint(
+          proto.saved_model_checksum,
+          proto.graph_def_program_hash,
+          proto.signature_def_hash,
+          proto.saved_object_graph_hash,
+          proto.checkpoint_hash,
+          proto.version)
+    except AttributeError as e:
+      raise ValueError(
+          f"Given proto could not be deserialized as fingerprint."
+          f"{e}") from None
+
+  def __eq__(self, other: Any) -> bool:
+    if (isinstance(other, Fingerprint) or
+        isinstance(other, fingerprint_pb2.FingerprintDef)):
+      try:
+        return (
+            self.saved_model_checksum == other.saved_model_checksum and
+            self.graph_def_program_hash == other.graph_def_program_hash and
+            self.signature_def_hash == other.signature_def_hash and
+            self.saved_object_graph_hash == other.saved_object_graph_hash and
+            self.checkpoint_hash == other.checkpoint_hash)
+      except AttributeError:
+        pass
+    return False
+
+  def __str__(self) -> str:
+    return "\n".join([
+        "SavedModel Fingerprint",
+        f"  saved_model_checksum: {self.saved_model_checksum}",
+        f"  graph_def_program_hash: {self.graph_def_program_hash}",
+        f"  signature_def_hash: {self.signature_def_hash}",
+        f"  saved_object_graph_hash: {self.saved_object_graph_hash}",
+        f"  checkpoint_hash: {self.checkpoint_hash}"
+    ])
+
+  def __repr__(self) -> str:
+    return (f"Fingerprint({self.saved_model_checksum}, "
+            f"{self.graph_def_program_hash}, "
+            f"{self.signature_def_hash}, "
+            f"{self.saved_object_graph_hash}, "
+            f"{self.checkpoint_hash})")
+
+  def singleprint(self) -> fingerprinting_pywrap.Singleprint:
+    """Canonical fingerprinting ID for a SavedModel.
+
+    Uniquely identifies a SavedModel based on the regularized fingerprint
+    attributes. (saved_model_checksum is sensitive to immaterial changes and
+    thus non-deterministic.)
+
+    Returns:
+      The string concatenation of `graph_def_program_hash`,
+      `signature_def_hash`, `saved_object_graph_hash`, and `checkpoint_hash`
+      fingerprint attributes (separated by '/').
+
+    Raises:
+      ValueError: If the fingerprint fields cannot be used to construct the
+      singleprint.
+    """
+    try:
+      return fingerprinting_pywrap.Singleprint(self.graph_def_program_hash,
+                                               self.signature_def_hash,
+                                               self.saved_object_graph_hash,
+                                               self.checkpoint_hash)
+    except (TypeError, fingerprinting_pywrap.FingerprintException) as e:
+      raise ValueError(
+          f"Encounted invalid fingerprint values when constructing singleprint."
+          f"graph_def_program_hash: {self.graph_def_program_hash}"
+          f"signature_def_hash: {self.signature_def_hash}"
+          f"saved_object_graph_hash: {self.saved_object_graph_hash}"
+          f"checkpoint_hash: {self.checkpoint_hash}"
+          f"{e}") from None
+
+
+@tf_export("saved_model.experimental.read_fingerprint", v1=[])
+def read_fingerprint(export_dir: str) -> Fingerprint:
+  """Reads the fingerprint of a SavedModel in `export_dir`.
+
+  Returns a `tf.saved_model.experimental.Fingerprint` object that contains
+  the values of the SavedModel fingerprint, which is persisted on disk in the
+  `fingerprint.pb` file in the `export_dir`.
+
+  Read more about fingerprints in the SavedModel guide at
+  https://www.tensorflow.org/guide/saved_model.
+
+  Args:
+    export_dir: The directory that contains the SavedModel.
+
+  Returns:
+    A `tf.saved_model.experimental.Fingerprint`.
+
+  Raises:
+    FileNotFoundError: If no or an invalid fingerprint is found.
+  """
+  try:
+    fingerprint = fingerprinting_pywrap.ReadSavedModelFingerprint(export_dir)
+  except fingerprinting_pywrap.FileNotFoundException as e:
+    raise FileNotFoundError(f"SavedModel Fingerprint Error: {e}") from None  # pylint: disable=raise-missing-from
+  except fingerprinting_pywrap.FingerprintException as e:
+    raise RuntimeError(f"SavedModel Fingerprint Error: {e}") from None  # pylint: disable=raise-missing-from
+  return Fingerprint.from_proto(
+      fingerprint_pb2.FingerprintDef().FromString(fingerprint))

tensorflow/python/saved_model/fingerprinting_test.py

@@ -0,0 +1,207 @@
+# 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 SavedModel fingerprinting.
+
+These tests verify that fingerprint is written correctly and that APIs for
+reading it are correct.
+"""
+import os
+import shutil
+
+from tensorflow.core.config import flags
+from tensorflow.core.protobuf import fingerprint_pb2
+from tensorflow.core.protobuf import saved_model_pb2
+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 dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.saved_model import fingerprinting
+from tensorflow.python.saved_model import fingerprinting_utils
+from tensorflow.python.saved_model import save
+from tensorflow.python.saved_model.pywrap_saved_model import constants
+from tensorflow.python.saved_model.pywrap_saved_model import fingerprinting as fingerprinting_pywrap
+from tensorflow.python.trackable import autotrackable
+
+
+class FingerprintingTest(test.TestCase):
+
+  def _create_saved_model(self):
+    root = autotrackable.AutoTrackable()
+    save_dir = os.path.join(self.get_temp_dir(), "saved_model")
+    save.save(root, save_dir)
+    self.addCleanup(shutil.rmtree, save_dir)
+    return save_dir
+
+  def _create_model_with_function(self):
+    root = autotrackable.AutoTrackable()
+    root.f = def_function.function(lambda x: 2. * x)
+    return root
+
+  def _create_model_with_input_signature(self):
+    root = autotrackable.AutoTrackable()
+    root.f = def_function.function(
+        lambda x: 2. * x,
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+    return root
+
+  def _create_model_with_data(self):
+    root = autotrackable.AutoTrackable()
+    root.x = constant_op.constant(1.0, dtype=dtypes.float32)
+    root.f = def_function.function(
+        lambda x: root.x * x,
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+    return root
+
+  def _read_fingerprint(self, filename):
+    fingerprint_def = fingerprint_pb2.FingerprintDef()
+    with file_io.FileIO(filename, "rb") as f:
+      fingerprint_def.ParseFromString(f.read())
+    return fingerprint_def
+
+  def setUp(self):
+    super().setUp()
+    flags.config().saved_model_fingerprinting.reset(True)
+
+  def test_basic_module(self):
+    save_dir = self._create_saved_model()
+    files = file_io.list_directory_v2(save_dir)
+
+    self.assertLen(files, 4)
+    self.assertIn(constants.FINGERPRINT_FILENAME, files)
+
+    fingerprint_def = self._read_fingerprint(
+        file_io.join(save_dir, constants.FINGERPRINT_FILENAME))
+
+    # We cannot check this value due to non-determinism in saving.
+    self.assertGreater(fingerprint_def.saved_model_checksum, 0)
+    self.assertEqual(fingerprint_def.graph_def_program_hash,
+                     14830488309055091319)
+    self.assertEqual(fingerprint_def.signature_def_hash, 12089566276354592893)
+    self.assertEqual(fingerprint_def.saved_object_graph_hash, 0)
+    # TODO(b/242348400): The checkpoint hash is non-deterministic, so we cannot
+    # check its value here.
+    self.assertGreater(fingerprint_def.checkpoint_hash, 0)
+
+  def test_model_saved_with_different_signature_options(self):
+    model = self._create_model_with_function()
+    # Save the model with signatures specified in SaveOptions.
+    sig_dir = os.path.join(self.get_temp_dir(), "saved_model")
+    save.save(
+        model,
+        sig_dir,
+        signatures=model.f.get_concrete_function(
+            tensor_spec.TensorSpec(None, dtypes.float32)))
+    # Save the model without signatures.
+    no_sig_dir = os.path.join(self.get_temp_dir(), "saved_model2")
+    save.save(model, no_sig_dir)
+    # Save the model with an input signature specified.
+    input_sig_dir = os.path.join(self.get_temp_dir(), "saved_model3")
+    save.save(self._create_model_with_input_signature(), input_sig_dir)
+
+    fingerprint_sig = self._read_fingerprint(
+        file_io.join(sig_dir, constants.FINGERPRINT_FILENAME))
+    fingerprint_no_sig = self._read_fingerprint(
+        file_io.join(no_sig_dir, constants.FINGERPRINT_FILENAME))
+    fingerprint_input_sig = self._read_fingerprint(
+        file_io.join(input_sig_dir, constants.FINGERPRINT_FILENAME))
+
+    # Check that the model saved with different options has different
+    # SignatureDef hashes.
+    self.assertNotEqual(fingerprint_sig.signature_def_hash,
+                        fingerprint_no_sig.signature_def_hash)
+    # Check that the model saved with the same concrete function has the same
+    # regularized hashes.
+    self.assertEqual(fingerprint_sig.graph_def_program_hash,
+                     fingerprint_input_sig.graph_def_program_hash)
+    self.assertEqual(fingerprint_sig.signature_def_hash,
+                     fingerprint_input_sig.signature_def_hash)
+
+  def test_read_fingerprint_api(self):
+    save_dir = self._create_saved_model()
+    fingerprint = fingerprinting.read_fingerprint(save_dir)
+
+    fingerprint_def = self._read_fingerprint(
+        file_io.join(save_dir, constants.FINGERPRINT_FILENAME))
+
+    self.assertEqual(fingerprint, fingerprint_def)
+
+  def test_read_fingerprint_file_not_found(self):
+    with self.assertRaisesRegex(FileNotFoundError,
+                                "SavedModel Fingerprint Error"):
+      fingerprinting.read_fingerprint("foo")
+
+  def test_write_fingerprint(self):
+    save_dir = os.path.join(self.get_temp_dir(), "model_and_fingerprint")
+    save.save_and_return_nodes(
+        self._create_model_with_data(), save_dir,
+        experimental_skip_checkpoint=True)  # checkpoint data won't be loaded
+
+    fingerprint_def = fingerprinting.read_fingerprint(save_dir)
+
+    # We cannot check this value due to non-determinism in serialization.
+    self.assertGreater(fingerprint_def.saved_model_checksum, 0)
+    self.assertEqual(fingerprint_def.graph_def_program_hash,
+                     8947653168630125217)
+    self.assertEqual(fingerprint_def.signature_def_hash, 15354238402988963670)
+    self.assertEqual(fingerprint_def.checkpoint_hash, 0)
+
+  def test_valid_singleprint(self):
+    save_dir = os.path.join(self.get_temp_dir(), "singleprint_model")
+    save.save(self._create_model_with_data(), save_dir)
+    fingerprint = fingerprinting.read_fingerprint(save_dir)
+    singleprint = fingerprint.singleprint()
+
+    # checkpoint_hash is non-deterministic and not included
+    self.assertRegex(singleprint,
+                     "/".join(["8947653168630125217",  # graph_def_program_hash
+                               "15354238402988963670",  # signature_def_hash
+                               "1613952301283913051"  # saved_object_graph_hash
+                               ]))
+
+  def test_invalid_singleprint(self):
+    fingerprint = fingerprinting.Fingerprint()
+    with self.assertRaisesRegex(ValueError,
+                                "Encounted invalid fingerprint values"):
+      fingerprint.singleprint()
+
+  def test_valid_from_proto(self):
+    save_dir = os.path.join(self.get_temp_dir(), "from_proto_model")
+    save.save(self._create_model_with_data(), save_dir)
+    fingerprint_def = fingerprint_pb2.FingerprintDef().FromString(
+        fingerprinting_pywrap.ReadSavedModelFingerprint(save_dir))
+    fingerprint = fingerprinting.Fingerprint.from_proto(fingerprint_def)
+    self.assertEqual(fingerprint, fingerprint_def)
+
+  def test_invalid_from_proto(self):
+    save_dir = os.path.join(self.get_temp_dir(), "from_proto_model")
+    save.save(self._create_model_with_data(), save_dir)
+    wrong_def = saved_model_pb2.SavedModel(
+        saved_model_schema_version=1)
+    with self.assertRaisesRegex(ValueError,
+                                "Given proto could not be deserialized as"):
+      fingerprinting.Fingerprint.from_proto(wrong_def)
+
+  def test_fingerprint_to_proto(self):
+    save_dir = os.path.join(self.get_temp_dir(), "from_proto_model")
+    save.save(self._create_model_with_data(), save_dir)
+    fingerprint = fingerprinting.read_fingerprint(save_dir)
+    fingerprint_def = fingerprinting_utils.to_proto(fingerprint)
+    self.assertEqual(fingerprint, fingerprint_def)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/fingerprinting_utils.py

@@ -0,0 +1,151 @@
+# 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.
+# ==============================================================================
+"""Utilities for SavedModel fingerprinting.
+
+This module contains utility classes and functions for working with the
+SavedModel fingerprint.
+"""
+
+from absl import logging
+
+from tensorflow.core.config import flags
+from tensorflow.core.protobuf import fingerprint_pb2
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.saved_model import fingerprinting
+from tensorflow.python.saved_model.pywrap_saved_model import constants
+from tensorflow.python.saved_model.pywrap_saved_model import fingerprinting as fingerprinting_pywrap
+from tensorflow.python.saved_model.pywrap_saved_model import metrics
+from tensorflow.python.util import compat
+
+FingerprintException = fingerprinting_pywrap.FingerprintException
+
+
+def write_fingerprint(export_dir: str) -> None:
+  """Write fingerprint protobuf, if requested.
+
+  Writes a `tf.saved_model.experimental.Fingerprint` object to a
+  `fingerprint.pb` file in the `export_dir` using the `saved_model.pb` file
+  contained in `export_dir`.
+
+  Args:
+    export_dir: The directory in which to write the fingerprint.
+  """
+  if flags.config().saved_model_fingerprinting.value():
+    fingerprint_path = file_io.join(
+        compat.as_str(export_dir),
+        compat.as_str(constants.FINGERPRINT_FILENAME))
+    logging.info("Writing fingerprint to %s", fingerprint_path)
+
+    try:
+      fingerprint_serialized = fingerprinting_pywrap.CreateFingerprintDef(
+          export_dir)
+    except FingerprintException as e:
+      raise ValueError(e) from None
+    file_io.atomic_write_string_to_file(fingerprint_path,
+                                        fingerprint_serialized)
+
+    metrics.SetWriteFingerprint(fingerprint=fingerprint_serialized)
+    try:
+      metrics.SetWritePathAndSingleprint(
+          path=export_dir,
+          singleprint=singleprint_from_fingerprint_proto(export_dir))
+    except metrics.MetricException:
+      logging.info("path_and_singleprint metric could not be set. "
+                   "Model saving will continue.")
+
+
+def singleprint_from_saved_model_proto(export_dir: str) -> str:
+  """Returns the singleprint of `saved_model.pb` in `export_dir`.
+
+  Args:
+    export_dir: The directory that contains `saved_model.pb`.
+
+  Returns:
+    A string containing the singleprint of `saved_model.pb` in `export_dir`.
+
+  Raises:
+    ValueError: If a valid singleprint cannot be constructed from 
+    `saved_model.pb`.
+  """
+  try:
+    return fingerprinting_pywrap.SingleprintFromSM(export_dir)
+  except FingerprintException as e:
+    raise ValueError(e) from None
+
+
+def singleprint_from_fingerprint_proto(export_dir: str) -> str:
+  """Returns the singleprint of `fingerprint.pb` in `export_dir`.
+
+  Args:
+    export_dir: The directory that contains `fingerprint.pb`.
+
+  Returns:
+    A string containing the singleprint of `fingerprint.pb` in `export_dir`.
+
+  Raises:
+    ValueError: If a valid singleprint cannot be constructed from
+    `fingerprint.pb`.
+  """
+  try:
+    return fingerprinting_pywrap.SingleprintFromFP(export_dir)
+  except FingerprintException as e:
+    raise ValueError(e) from None
+
+
+def singleprint_from_saved_model(export_dir: str) -> str:
+  """Returns the singleprint of the SavedModel in `export_dir`.
+
+  First tries to construct the singleprint from `fingerprint.pb`, then from
+  `saved_model.pb`. Attempts to write the `fingerprint.pb` if not found, but
+  doesn't return an error if it isn't writeable.
+
+  Args:
+    export_dir: The directory that contains the SavedModel.
+
+  Returns:
+    A string containing the singleprint of the SavedModel in `export_dir`.
+
+  Raises:
+    ValueError: If a valid singleprint cannot be constructed from the
+    SavedModel.
+  """
+  # try generating the singleprint from `fingerprint.pb`
+  try:
+    return singleprint_from_fingerprint_proto(export_dir)
+  except ValueError:
+    pass
+
+  # try creating `fingerprint.pb`, then generating the singleprint
+  try:
+    write_fingerprint(export_dir)
+    return singleprint_from_fingerprint_proto(export_dir)
+  except ValueError:
+    pass
+
+  # try generating the singleprint from `saved_model.pb`
+  try:
+    return singleprint_from_saved_model_proto(export_dir)
+  except ValueError as e:
+    raise ValueError(e) from None
+
+
+def to_proto(
+    fingerprint: fingerprinting.Fingerprint) -> fingerprint_pb2.FingerprintDef:
+  return fingerprint_pb2.FingerprintDef(
+      saved_model_checksum=fingerprint.saved_model_checksum,
+      graph_def_program_hash=fingerprint.graph_def_program_hash,
+      signature_def_hash=fingerprint.signature_def_hash,
+      saved_object_graph_hash=fingerprint.saved_object_graph_hash,
+      checkpoint_hash=fingerprint.checkpoint_hash)

tensorflow/python/saved_model/function_deserialization.py

@@ -0,0 +1,715 @@
+# 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.
+# ==============================================================================
+"""Tools for deserializing `Function`s."""
+
+import collections
+import pprint
+import re
+
+from absl import logging
+
+from tensorflow.core.function import trace_type
+from tensorflow.core.function.polymorphism import function_type as function_type_lib
+from tensorflow.core.protobuf import saved_object_graph_pb2
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as function_lib
+from tensorflow.python.eager.polymorphic_function import function_type_utils
+from tensorflow.python.framework import func_graph as func_graph_lib
+from tensorflow.python.framework import function_def_to_graph as function_def_lib
+from tensorflow.python.framework import op_def_registry
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import custom_gradient
+from tensorflow.python.ops import default_gradient
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+
+
+def _is_tensor(t):
+  return isinstance(
+      t, (tensor.Tensor, resource_variable_ops.BaseResourceVariable))
+
+
+# TODO(b/205016027): Update this to just use ConcreteFunction.__call__ with the
+# structured signature.
+def _call_concrete_function(function, inputs):
+  """Calls a restored Function with structured inputs.
+
+  This differs from `function.__call__` in that inputs and outputs are
+  structured and that it casts inputs to tensors if needed.
+
+  Note: this does not checks that non-tensor inputs match. That should be
+  done before via `_concrete_function_callable_with`.
+
+  Args:
+    function: ConcreteFunction to call.
+    inputs: Structured inputs compatible with
+      `function.graph.structured_input_signature`.
+
+  Returns:
+    The structured function output.
+  """
+  expected_structure = function.graph.structured_input_signature
+  flatten_inputs = nest.flatten_up_to(
+      expected_structure, inputs, expand_composites=True)
+  flatten_expected = nest.flatten(expected_structure, expand_composites=True)
+  tensor_inputs = []
+  for arg, expected in zip(flatten_inputs, flatten_expected):
+    if isinstance(expected, tensor.TensorSpec):
+      tensor_inputs.append(
+          ops.convert_to_tensor(arg, dtype_hint=expected.dtype))
+    elif isinstance(expected, resource_variable_ops.VariableSpec):
+      tensor_inputs.append(arg.handle)
+  result = function._call_flat(tensor_inputs, function.captured_inputs)  # pylint: disable=protected-access
+  if isinstance(result, ops.Operation):
+    return None
+  return result
+
+
+def _try_convert_to_tensor_spec(arg, dtype_hint):
+  """Returns None or TensorSpec obtained if `arg` is converted to tensor."""
+  try:
+    # Note: try conversion in a FuncGraph to avoid polluting current context.
+    with func_graph_lib.FuncGraph(name="guess_conversion").as_default():
+      result = ops.convert_to_tensor(arg, dtype_hint=dtype_hint)
+      return tensor.TensorSpec(shape=result.shape, dtype=result.dtype)
+  except (TypeError, ValueError):
+    return None
+
+
+def _concrete_function_callable_with(function, inputs, allow_conversion):
+  """Returns whether concrete `function` can be called with `inputs`."""
+  expected_structure = function.graph.structured_input_signature
+  try:
+    flatten_inputs = nest.flatten_up_to(expected_structure, inputs)
+  except (TypeError, ValueError):
+    return False
+
+  for arg, expected in zip(flatten_inputs, nest.flatten(expected_structure)):
+    if isinstance(expected, tensor.TensorSpec):
+      if allow_conversion:
+        arg = _try_convert_to_tensor_spec(arg, dtype_hint=expected.dtype)
+      if not _is_tensor(arg) and not isinstance(arg, tensor.TensorSpec):
+        return False
+      if arg.dtype != expected.dtype:
+        return False
+      if not expected.shape.is_compatible_with(arg.shape):
+        return False
+    elif isinstance(expected, type_spec.TypeSpec):
+      if not expected.is_compatible_with(arg):
+        return False
+    elif _is_tensor(arg):
+      if id(arg) != id(expected):
+        return False
+    else:
+      if arg != expected:
+        return False
+  return True
+
+
+def _deserialize_function_spec_as_nonmethod(function_spec_proto):
+  """Deserialize a FunctionSpec object from its proto representation."""
+  typeless_fullargspec = nested_structure_coder.decode_proto(
+      function_spec_proto.fullargspec)
+
+  # Convert a method function into a non method.
+  if function_spec_proto.is_method or (
+      typeless_fullargspec.args and typeless_fullargspec.args[0] == "self"
+  ):
+    if not typeless_fullargspec.args:
+      raise NotImplementedError(
+          "Cannot deserialize a method function without a named "
+          "'self' argument.")
+    args = typeless_fullargspec.args[1:]
+  else:
+    args = typeless_fullargspec.args
+
+  fullargspec = tf_inspect.FullArgSpec(
+      args=args,
+      varargs=typeless_fullargspec.varargs,
+      varkw=typeless_fullargspec.varkw,
+      defaults=typeless_fullargspec.defaults,
+      kwonlyargs=typeless_fullargspec.kwonlyargs,
+      kwonlydefaults=typeless_fullargspec.kwonlydefaults,
+      annotations=typeless_fullargspec.annotations)
+  input_signature = nested_structure_coder.decode_proto(
+      function_spec_proto.input_signature)
+
+  # See `tf.function` and the JitCompile proto for details.
+  jit_compile = {
+      saved_object_graph_pb2.FunctionSpec.JitCompile.DEFAULT: None,
+      saved_object_graph_pb2.FunctionSpec.JitCompile.ON: True,
+      saved_object_graph_pb2.FunctionSpec.JitCompile.OFF: False,
+  }.get(function_spec_proto.jit_compile)
+
+  return function_type_utils.FunctionSpec.from_fullargspec_and_signature(
+      fullargspec=fullargspec,
+      input_signature=input_signature,
+      jit_compile=jit_compile)
+
+
+# TODO(b/203440205): Set FunctionType with ConcreteFunction constructor.
+def set_preinitialized_function_spec(concrete_fn, spec):
+  """Set the FunctionType of the ConcreteFunction using FunctionSpec."""
+  if spec is None:
+    concrete_fn._function_type = None  # pylint: disable=protected-access
+    return
+
+  unconstrained_type = function_type_lib.FunctionType(
+      [
+          function_type_lib.Parameter(p.name, p.kind, p.optional, None)
+          for p in spec.function_type.parameters.values()
+      ]
+  )
+  arg_specs, kwarg_specs = concrete_fn.structured_input_signature
+
+  input_function_type, _ = function_type_lib.canonicalize_to_monomorphic(
+      arg_specs,
+      {
+          function_type_lib.sanitize_arg_name(k): v
+          for k, v in kwarg_specs.items()
+      },
+      spec.default_values,
+      {},
+      unconstrained_type,
+  )
+
+  output_type = trace_type.from_value(concrete_fn.graph.structured_outputs)
+  # Captures are restored later so we will update it then.
+  function_type = function_type_lib.FunctionType(
+      input_function_type.parameters.values(),
+      return_annotation=output_type,
+  )
+  concrete_fn._function_type = function_type  # pylint: disable=protected-access
+
+
+# TODO(b/205016761): The fact that we can't derive ConcreteFunction calling
+# conventions from the serialized input spec right now is unfortunate. Merging
+# these would be good, maybe by adding TensorSpec names to cache keys so renamed
+# keyword arguments would yield different ConcreteFunctions.
+def setup_bare_concrete_function(saved_bare_concrete_function,
+                                 concrete_functions):
+  """Makes a restored bare concrete function callable."""
+  concrete_function = concrete_functions[
+      saved_bare_concrete_function.concrete_function_name]
+  # pylint: disable=protected-access
+  concrete_function._arg_keywords = (
+      saved_bare_concrete_function.argument_keywords)
+  concrete_function._num_positional_args = (
+      saved_bare_concrete_function.allowed_positional_arguments)
+  if saved_bare_concrete_function.HasField("function_spec"):
+    function_spec = _deserialize_function_spec_as_nonmethod(
+        saved_bare_concrete_function.function_spec)
+    set_preinitialized_function_spec(concrete_function, function_spec)
+  # pylint: enable=protected-access
+  concrete_function.add_to_graph()
+  return concrete_function
+
+
+class RestoredFunction(def_function.Function):
+  """Wrapper class for a function that has been restored from saved state.
+
+  See `def_function.Function`.
+  """
+
+  def __init__(self, python_function, name, function_spec, concrete_functions):
+    # TODO(b/205016819): We may enable autograph once exceptions are supported.
+    super(RestoredFunction, self).__init__(
+        python_function,
+        name,
+        autograph=False,
+        jit_compile=function_spec.jit_compile)
+    self.concrete_functions = concrete_functions
+    self._function_type = function_spec.function_type
+    self._default_values = function_spec.default_values
+
+    # Prevent RestoredFunction from spamming users with frequent tracing
+    # warnings.
+    self._omit_frequent_tracing_warning = True
+
+  @property
+  def _run_functions_eagerly(self):
+    # We do not have access to the original python function, and thus, we
+    # cannot meaningfully do anything but call our concrete function graphs
+    # under the hood.
+    #
+    # Attempting to call our bespoke python function (i.e.
+    # `restored_function_body`) will work so long as the user passes in all
+    # required and optional arguments. If an optional argument is missing,
+    # however, the call will break. For this reason, we instead skip the
+    # eager call path altogether if a user has enabled eager function execution
+    # via `tf.config.run_functions_eagerly`.
+    return False
+
+  def _list_all_concrete_functions(self):
+    return self.concrete_functions
+
+  def _list_all_concrete_functions_for_serialization(self):
+    return self.concrete_functions
+
+
+def recreate_function(saved_function, concrete_functions):
+  """Creates a `Function` from a `SavedFunction`.
+
+  Args:
+    saved_function: `SavedFunction` proto.
+    concrete_functions: map from function name to `ConcreteFunction`. As a side
+      effect of this function, the `FunctionSpec` from `saved_function` is added
+      to each `ConcreteFunction` in this map.
+
+  Returns:
+    A `Function`.
+  """
+  # TODO(b/205017389): Construct a `Function` with the cache populated
+  # instead of creating a new `Function` backed by a Python layer to
+  # glue things together. Current approach is nesting functions deeper for each
+  # serialization cycle.
+
+  # Note: handling method functions is tricky since make_decorator does not
+  # allows control of "ismethod". Additionally since restored functions do
+  # not behave as methods i.e. they always use the same captured tensors
+  # independent of the object they are bound to, there is little value on
+  # propagating that correctly.
+  #
+  # Ideally this conversion should happen at serialization time. But since
+  # there are SavedModels which have "ismethod" populated and have an extra
+  # argument that they expect to be ignored, we do it at deserialization.
+  function_spec = _deserialize_function_spec_as_nonmethod(
+      saved_function.function_spec)
+
+  def restored_function_body(*args, **kwargs):
+    """Calls a restored function or raises an error if no matching function."""
+    if not saved_function.concrete_functions:
+      raise ValueError("Found zero restored functions for caller function.")
+    # This is the format of function.graph.structured_input_signature. At this
+    # point, the args and kwargs have already been canonicalized.
+    inputs = (args, kwargs)
+
+    # First try to find a concrete function that can be called without input
+    # conversions. This allows one to pick a more specific trace in case there
+    # was also a more expensive one that supported tensors.
+    for allow_conversion in [False, True]:
+      for function_name in saved_function.concrete_functions:
+        function = concrete_functions[function_name]
+        if any([inp is None for inp in function.captured_inputs]):
+          raise ValueError("Looks like you are trying to run a loaded "
+                           "non-Keras model that was trained using "
+                           "tf.distribute.experimental.ParameterServerStrategy "
+                           "with variable partitioning, which is not currently "
+                           "supported. Try using Keras to define your model "
+                           "if possible.")
+        if _concrete_function_callable_with(function, inputs, allow_conversion):
+          return _call_concrete_function(function, inputs)
+
+    signature_descriptions = []
+
+    def _pretty_format_positional(positional):
+      return "Positional arguments ({} total):\n    * {}".format(
+          len(positional),
+          "\n    * ".join(pprint.pformat(a) for a in positional))
+
+    for index, function_name in enumerate(saved_function.concrete_functions):
+      concrete_function = concrete_functions[function_name]
+      positional, keyword = concrete_function.structured_input_signature
+      signature_descriptions.append(
+          "Option {}:\n  {}\n  Keyword arguments: {}".format(
+              index + 1, _pretty_format_positional(positional), keyword))
+    raise ValueError(
+        "Could not find matching concrete function to call loaded from the "
+        f"SavedModel. Got:\n  {_pretty_format_positional(args)}\n  Keyword "
+        f"arguments: {kwargs}\n\n Expected these arguments to match one of the "
+        f"following {len(saved_function.concrete_functions)} option(s):\n\n"
+        f"{(chr(10)+chr(10)).join(signature_descriptions)}")
+
+  concrete_function_objects = []
+  for concrete_function_name in saved_function.concrete_functions:
+    concrete_function_objects.append(concrete_functions[concrete_function_name])
+
+  for cf in concrete_function_objects:
+    set_preinitialized_function_spec(cf, function_spec)
+
+  restored_function = RestoredFunction(restored_function_body,
+                                       restored_function_body.__name__,
+                                       function_spec, concrete_function_objects)
+
+  return tf_decorator.make_decorator(
+      restored_function_body,
+      restored_function,
+      decorator_argspec=function_spec.fullargspec)
+
+
+def load_function_def_library(library,
+                              saved_object_graph=None,
+                              load_shared_name_suffix=None,
+                              wrapper_function=None):
+  """Load a set of functions as concrete functions without captured inputs.
+
+  Functions names are manipulated during load such that they do not overlap
+  with previously created ones.
+
+  Gradients are re-registered under new names. Ops that reference the gradients
+  are updated to reflect the new registered names.
+
+  Args:
+    library: FunctionDefLibrary proto message.
+    saved_object_graph: SavedObjectGraph proto message. If not passed in,
+      concrete function structured signatures and outputs will not be set.
+    load_shared_name_suffix: If specified, used to uniquify shared names.
+      Otherwise, a unique name is generated.
+    wrapper_function: An object that will be wrapped on newly created functions.
+
+  Returns:
+    Map of original function names in the library to instances of
+    `ConcreteFunction` without captured inputs.
+
+  Raises:
+    ValueError: if functions dependencies have a cycle.
+  """
+  library_function_names = set(fdef.signature.name for fdef in library.function)
+  functions = {}
+  renamed_functions = {}
+
+  # Our graph building code currently requires functions to be registered with
+  # some tf.Graph in order to import functions using the
+  # op-name-is-function-name calling convention. To avoid leaking memory into
+  # the global default graph when executing eagerly, we create a temporary
+  # Graph.
+  #
+  # TODO(b/205023033): Make this Graph creation unnecessary when executing
+  # eagerly by fixing function_def_to_graph_def.
+  if ops.executing_eagerly_outside_functions():
+    graph = ops.Graph()
+  else:
+    graph = ops.get_default_graph()
+
+  if load_shared_name_suffix is None:
+    load_shared_name_suffix = "_load_{}".format(ops.uid())
+
+  # Custom gradient functions must be re-registered under new UIDs.
+  library_gradient_names = {}  # Maps old op type to old function name
+  new_gradient_op_types = {}  # Maps old gradient op type to new op type.
+  gradients_to_register = {}  # Maps old function name to new op type
+  for gdef in library.registered_gradients:
+    if gdef.registered_op_type:
+      new_op_type = custom_gradient.generate_name()
+      old_op_type = compat.as_bytes(gdef.registered_op_type)
+
+      library_gradient_names[old_op_type] = gdef.gradient_func
+      new_gradient_op_types[old_op_type] = new_op_type
+      gradients_to_register[gdef.gradient_func] = new_op_type
+
+  function_deps = {}
+  for fdef in library.function:
+    function_deps[fdef.signature.name] = _list_function_deps(
+        fdef, library_function_names, library_gradient_names)
+
+  loaded_gradients = {}
+  for fdef in _sort_function_defs(library, function_deps):
+    orig_name = _fix_fdef_in_place(fdef, functions, load_shared_name_suffix,
+                                   new_gradient_op_types)
+
+    # Setup function signatures and outputs
+    #
+    # When concrete functions are created normally (i.e. when they're originally
+    # created and not loaded via saved model), the inputs and outputs are
+    # calculated based on the values passed in by the user and returned from the
+    # original function, respectively. We don't have access to those anymore at
+    # restore time, so we must instead pass them to the FuncGraph explicitly.
+    structured_input_signature = None
+    structured_outputs = None
+    if (saved_object_graph is not None and
+        orig_name in saved_object_graph.concrete_functions):
+      # TODO(b/204324043): Offload the deserialization of the protos to the
+      # first class objects by passing the actual protos. This is blocked on
+      # importing `nested_structure_coder` in function.py causing a circular
+      # dependency.
+      proto = saved_object_graph.concrete_functions[orig_name]
+      structured_input_signature = nested_structure_coder.decode_proto(
+          proto.canonicalized_input_signature)
+      structured_outputs = nested_structure_coder.decode_proto(
+          proto.output_signature)
+
+    # There is no need to copy all functions into the function def graph. It
+    # leads to a O(n^2) increase of memory when importing functions and the
+    # extra function definitions are a no-op since they already imported as a
+    # function before and passed in explicitly (due to the topologic sort
+    # import).
+    with graph.as_default():
+      func_graph = function_def_lib.function_def_to_graph(
+          fdef,
+          structured_input_signature=structured_input_signature,
+          structured_outputs=structured_outputs)
+    # Restores gradients for function-call ops (not the same as ops that use
+    # custom gradients)
+    _restore_gradient_functions(func_graph, renamed_functions, loaded_gradients)
+
+    for dep in function_deps[orig_name]:
+      functions[dep].add_to_graph(func_graph)
+
+    # We do not initialize the new ConcreteFunction's function_spec and/or
+    # arg_keywords here (which are used to parse the structured and flat
+    # signatures, respectively). ConcreteFunction that are part of a saved
+    # function is set up later by recreate_function(); and bare ConcreteFunction
+    # is set up by by setup_bare_concrete_function().
+    # However, we copy the FunctionDef attributes to the new ConcreteFunction,
+    # excluding the "_input_shapes", which may cause an error during input shape
+    # initialization at a later stage.
+    if "_input_shapes" in fdef.attr:
+      del fdef.attr["_input_shapes"]
+    function_type = function_type_lib.from_structured_signature(
+        func_graph.structured_input_signature,
+        func_graph.structured_outputs,
+        func_graph.function_captures.capture_types,
+    )
+    func = function_lib.ConcreteFunction.from_func_graph(
+        func_graph, function_type, attrs=fdef.attr)
+    if wrapper_function:
+      func = wrapper_function(func)
+    func.add_to_graph(graph)
+
+    functions[orig_name] = func
+    renamed_functions[func.name] = func
+    if any(op.type == "TRTEngineOp" for op in func_graph.get_operations()):
+      # TODO(b/150708051): Remove this hack once TensorRT SavedModel integration
+      # is fixed. Currently it's leaking memory to maintain bug compatibility
+      # with previous behavior.
+      func.add_to_graph(ops.get_default_graph())
+
+    if orig_name in gradients_to_register:
+      gradient_op_type = gradients_to_register[orig_name]
+      loaded_gradients[compat.as_bytes(gradient_op_type)] = func
+      ops.RegisterGradient(gradient_op_type)(_gen_gradient_func(func))
+
+  return functions
+
+
+def _gen_gradient_func(func):
+  """Wraps a deserialized function."""
+
+  def gradient_func(unused_op, *result_grads):
+    # Replace all `None` arguments, because the traced custom gradient function
+    # expects tensors. Replacing with zeros is correct since the `None` values
+    # occur when the gradient is unconnected, and thus the gradient is
+    # "statically proven to be zero." See `tf.UnconnectedGradients` for details.
+
+    def none_to_zero(x, t):
+      if x is not None:
+        return x
+
+      shape, dtype = default_gradient.shape_and_dtype(t)
+
+      if shape.is_fully_defined():
+        return default_gradient.zeros_like(t)
+
+      dims = []
+      if shape.rank is not None:
+        dims = [1 if d is None else d for d in shape.as_list()]
+
+      return array_ops.zeros(dims, dtype)
+
+    result_grads = [
+        none_to_zero(x, t) for (x, t) in zip(result_grads, func.graph.inputs)
+    ]
+
+    return func(*result_grads)
+
+  return gradient_func
+
+
+def _restore_gradient_functions(func_graph, renamed_functions,
+                                loaded_gradients):
+  """Populate function op's _gradient_function with default gradient."""
+  for op in func_graph.get_operations():
+    # TODO(b/205024208): This code assumes that the gradient registered for this
+    # function call is the default gradient for the function and not a custom
+    # one.
+    if op.type in ["StatefulPartitionedCall", "PartitionedCall"]:
+      function = renamed_functions[compat.as_bytes(
+          op.node_def.attr["f"].func.name)]
+      op._gradient_function = function._get_gradient_function()  # pylint: disable=protected-access
+    try:
+      gradient_op_type = op.get_attr("_gradient_op_type")
+    except ValueError:
+      pass
+    else:
+      if gradient_op_type in loaded_gradients:
+        grad_fn = loaded_gradients[gradient_op_type]
+        grad_fn._num_positional_args = len(op.inputs)  # pylint: disable=protected-access
+        grad_fn._arg_keywords = [inp.name for inp in op.inputs]  # pylint: disable=protected-access
+
+
+def _sort_function_defs(library, function_deps):
+  """Return a topologic sort of FunctionDefs in a library."""
+  edges = collections.defaultdict(list)
+  in_count = collections.defaultdict(lambda: 0)
+
+  for fname, deps in function_deps.items():
+    for dep in deps:
+      edges[dep].append(fname)
+      in_count[fname] += 1
+  ready = [
+      fdef.signature.name
+      for fdef in library.function
+      if in_count[fdef.signature.name] == 0
+  ]
+  output = []
+  while ready:
+    node = ready.pop()
+    output.append(node)
+    for dest in edges[node]:
+      in_count[dest] -= 1
+      if not in_count[dest]:
+        ready.append(dest)
+
+  if len(output) != len(library.function):
+    failed_to_resolve = sorted(set(in_count.keys()) - set(output))
+    raise ValueError("There is a cyclic dependency between functions. ",
+                     f"Could not resolve {failed_to_resolve}.")
+
+  reverse = {fdef.signature.name: fdef for fdef in library.function}
+  return [reverse[x] for x in output]
+
+
+def _get_gradient_op_type(node_def):
+  """Returns the custom gradient op type."""
+  if ("_gradient_op_type" in node_def.attr and
+      node_def.op not in ["StatefulPartitionedCall", "PartitionedCall"]):
+    return node_def.attr["_gradient_op_type"].s
+  return None
+
+
+def fix_node_def(node_def, functions, shared_name_suffix):
+  """Replace functions calls and shared names in `node_def`."""
+  if node_def.op in functions:
+    node_def.op = functions[node_def.op].name
+  for _, attr_value in node_def.attr.items():
+    if attr_value.WhichOneof("value") == "func":
+      attr_value.func.name = functions[attr_value.func.name].name
+    elif attr_value.WhichOneof("value") == "list":
+      for fn in attr_value.list.func:
+        fn.name = functions[fn.name].name
+
+  # Fix old table creation bug.
+  if node_def.op == "HashTableV2":
+    if ("use_node_name_sharing" not in node_def.attr or
+        not node_def.attr["use_node_name_sharing"].b):
+      node_def.attr["use_node_name_sharing"].b = True
+      # We are turning on node mame sharing, so have to make sure we don't
+      # accidentally share a table resource.
+      shared_name_suffix += "_{}".format(ops.uid())
+
+  # TODO(b/124205571): Avoid accidental sharing and destruction of restored
+  # resources. For now uniquify "shared_name" when loading functions to avoid
+  # sharing.
+  # TODO: Add regression test for b/150826922.
+  op_def = op_def_registry.get(node_def.op)
+  if op_def:
+    attr = next((a for a in op_def.attr if a.name == "shared_name"), None)
+    if attr:
+      shared_name = None
+      if "shared_name" in node_def.attr and node_def.attr["shared_name"].s:
+        shared_name = node_def.attr["shared_name"].s
+      elif attr.default_value.s:
+        shared_name = compat.as_bytes(attr.default_value.s)
+      if not shared_name:
+        shared_name = compat.as_bytes(node_def.name)
+
+      node_def.attr["shared_name"].s = (
+          shared_name + compat.as_bytes(shared_name_suffix))
+
+
+def _fix_fdef_in_place(fdef, functions, shared_name_suffix,
+                       new_gradient_op_types):
+  """Fixes a FunctionDef proto to be loaded in current context.
+
+  In particular, when loading a function library into an eager context, one
+  must rename the functions to avoid conflicts with existent functions.
+
+  Args:
+    fdef: FunctionDef proto to fix. It is mutated in-place.
+    functions: map from function name to a ConcreteFunction instance.
+    shared_name_suffix: A unique string for this load which helps to avoid
+      `shared_name` collisions across loads. Two functions from the same load
+      using the same `shared_name` still need to share, but functions from
+      different loads with the same `shared_name` should not.
+    new_gradient_op_types: map from old gradient op type to newly generated op
+      type.
+
+  Returns:
+    orig_name: original value of fdef.signature.name
+  """
+  orig_name = fdef.signature.name
+  contains_unsaved_custom_gradients = False
+
+  for node_def in fdef.node_def:
+    fix_node_def(node_def, functions, shared_name_suffix)
+    op_type = _get_gradient_op_type(node_def)
+    if op_type is not None:
+      if op_type in new_gradient_op_types:
+        node_def.attr["_gradient_op_type"].s = compat.as_bytes(
+            new_gradient_op_types[op_type])
+      else:
+        contains_unsaved_custom_gradients = True
+  if contains_unsaved_custom_gradients:
+    logging.warning(
+        "Importing a function (%s) with ops with unsaved custom gradients. Will"
+        " likely fail if a gradient is requested.", fdef.signature.name)
+
+  fdef.signature.name = _clean_function_name(fdef.signature.name)
+  return orig_name
+
+
+def _list_function_deps(fdef, library_function_names, library_gradient_names):
+  """Find functions referenced in `fdef`."""
+  # TODO(b/205023953): Recurse into list attributes and into NameAttrList attrs
+  # both when listing deps and when fixing them. `function_def_to_graph` also
+  # requires fixes.
+  deps = set()
+  for node_def in fdef.node_def:
+    grad_op_type = _get_gradient_op_type(node_def)
+    if node_def.op in library_function_names:
+      deps.add(node_def.op)
+    elif grad_op_type and grad_op_type in library_gradient_names:
+      deps.add(library_gradient_names[grad_op_type])
+    else:
+      for _, attr_value in node_def.attr.items():
+        if attr_value.WhichOneof("value") == "func":
+          deps.add(attr_value.func.name)
+        elif attr_value.WhichOneof("value") == "list":
+          for fn in attr_value.list.func:
+            deps.add(fn.name)
+
+  return deps
+
+
+_FUNCTION_WRAPPER_NAME_REGEX = r"^%s(.*)_\d+$" % (function_lib._INFERENCE_PREFIX
+                                                 )  # pylint:disable=protected-access
+
+
+def _clean_function_name(name):
+  """Vanity function to keep the function names comprehensible."""
+  # Note: each time a function is wrapped into `function_lib.ConcreteFunction`
+  # its name becomes "__inference_<orig>_xyz".
+  match = re.search(_FUNCTION_WRAPPER_NAME_REGEX, name)
+  if match:
+    return match.group(1)
+  else:
+    return name

tensorflow/python/saved_model/function_serialization.py

@@ -0,0 +1,220 @@
+# 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.
+# ==============================================================================
+"""Tools for serializing `Function`s."""
+
+from tensorflow.core.function.polymorphism import function_type as function_type_lib
+from tensorflow.core.protobuf import saved_object_graph_pb2
+from tensorflow.python.eager import function as defun
+from tensorflow.python.eager import wrap_function as wrap_function_lib
+from tensorflow.python.eager.polymorphic_function import function_type_utils
+from tensorflow.python.framework import func_graph as func_graph_module
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.util import nest
+
+
+def _serialize_function_spec(function_spec):
+  """Serialize a FunctionSpec object into its proto representation."""
+  if (
+      function_spec.fullargspec.args
+      and function_spec.fullargspec.args[0] == "self"
+  ):
+    raise TypeError(
+        "Can not serialize tf.function with unbound 'self' parameter."
+    )
+
+  proto = saved_object_graph_pb2.FunctionSpec()
+
+  # Intentionally skip encoding annotations of a function because function
+  # annotations are mainly for optional type checking during development
+  # and does not affect runtime behavior.
+  # https://www.python.org/dev/peps/pep-3107/
+  # https://docs.python.org/3/library/inspect.html#inspect.getfullargspec
+  proto.fullargspec.CopyFrom(
+      nested_structure_coder.encode_structure(
+          function_spec.fullargspec._replace(annotations={})))
+
+  proto.is_method = False
+  proto.input_signature.CopyFrom(
+      nested_structure_coder.encode_structure(function_spec.input_signature))
+
+  # See `tf.function` and the JitCompile proto for details.
+  proto.jit_compile = {
+      None: saved_object_graph_pb2.FunctionSpec.JitCompile.DEFAULT,
+      True: saved_object_graph_pb2.FunctionSpec.JitCompile.ON,
+      False: saved_object_graph_pb2.FunctionSpec.JitCompile.OFF,
+  }.get(function_spec.jit_compile)
+
+  return proto
+
+
+def serialize_concrete_function(concrete_function, node_ids):
+  """Build a SavedConcreteFunction."""
+  bound_inputs = []
+  try:
+    for capture in concrete_function.captured_inputs:
+      bound_inputs.append(node_ids[capture])
+  except KeyError:
+    raise KeyError(
+        f"Failed to add concrete function '{concrete_function.name}' to object-"
+        f"based SavedModel as it captures tensor {capture!r} which is unsupported"
+        " or not reachable from root. "
+        "One reason could be that a stateful object or a variable that the "
+        "function depends on is not assigned to an attribute of the serialized "
+        "trackable object (see SaveTest.test_captures_unreachable_variable).")
+  concrete_function_proto = saved_object_graph_pb2.SavedConcreteFunction()
+  structured_outputs = func_graph_module.convert_structure_to_signature(
+      concrete_function.structured_outputs)
+  concrete_function_proto.canonicalized_input_signature.CopyFrom(
+      nested_structure_coder.encode_structure(
+          concrete_function.structured_input_signature))
+  concrete_function_proto.output_signature.CopyFrom(
+      nested_structure_coder.encode_structure(structured_outputs))
+  concrete_function_proto.bound_inputs.extend(bound_inputs)
+  return concrete_function_proto
+
+
+# TODO(b/203440205): Support FunctionType directly.
+def get_preinitialized_function_spec(concrete_function):
+  """Generates an unconstrained FunctionSpec from FunctionType."""
+  # TODO(b/203440205): SavedModel does not support FunctionType on its own
+  # without a FuncGraph signature.
+  # WrappedFunctions are not supposed to have FunctionSpecs.
+  if concrete_function.structured_input_signature is None or isinstance(
+      concrete_function, wrap_function_lib.WrappedFunction
+  ):
+    return None
+
+  function_type = concrete_function.function_type
+  if function_type is None:
+    return None
+
+  unconstrained_type = function_type_lib.FunctionType(
+      [
+          function_type_lib.Parameter(p.name, p.kind, p.optional, None)
+          for p in function_type.parameters.values()
+      ]
+  )
+  default_values = {
+      p.default for p in function_type.parameters.values() if p.optional
+  }
+  return function_type_utils.FunctionSpec(
+      unconstrained_type,
+      default_values,
+      False,
+      name=concrete_function.name,
+  )
+
+
+def serialize_bare_concrete_function(concrete_function):
+  """Build a SavedBareConcreteFunction."""
+  # pylint: disable=protected-access
+  proto = saved_object_graph_pb2.SavedBareConcreteFunction(
+      concrete_function_name=concrete_function.name,
+      allowed_positional_arguments=concrete_function._num_positional_args,
+      argument_keywords=concrete_function._arg_keywords)
+  function_spec = get_preinitialized_function_spec(concrete_function)
+  if function_spec is not None:
+    proto.function_spec.CopyFrom(_serialize_function_spec(function_spec))
+  return proto
+  # pylint: enable=protected-access
+
+
+def serialize_function(function, concrete_functions):
+  """Build a SavedFunction proto."""
+  proto = saved_object_graph_pb2.SavedFunction()
+
+  function_spec_proto = _serialize_function_spec(function.function_spec)
+  proto.function_spec.CopyFrom(function_spec_proto)
+  for concrete_function in concrete_functions:
+    proto.concrete_functions.append(concrete_function.name)
+  return proto
+
+
+def wrap_cached_variables(concrete_function):
+  """Wraps the concrete function if it uses cached read tensors.
+
+  This function creates a new concrete function that captures variables
+  instead of the cached read tensors.
+
+  Args:
+    concrete_function: A Concrete function that maybe captures cached read
+      tensors.
+
+  Returns:
+    A concrete function that wraps the original concrete function, which
+    captures variables instead. If the original function did not capture any
+    cached values, then the function is not wrapped and the original object is
+    returned.
+  """
+  outer_graph = func_graph_module.FuncGraph(
+      "{}_no_cache".format(concrete_function.graph.name))
+  mapped_captures = None
+  remapped_captures = {}
+
+  # Update the external captures to use read tensors generated in the outer
+  # graph.
+  with outer_graph.as_default():
+    for capture, placeholder in concrete_function.graph.captures:
+      cached_variable = getattr(capture, "_cached_variable", None)
+      if cached_variable is None:
+        continue
+      cached_variable = cached_variable()
+      new_cached_value = cached_variable.read_value()
+      key = id(capture)
+      external = concrete_function.graph.function_captures.by_val_external[key]
+      internal = concrete_function.graph.function_captures.by_val_internal[key]
+      remapped_captures[key] = [external, internal]
+      concrete_function.graph.function_captures.add_or_replace(
+          key=key,
+          external=new_cached_value,
+          internal=placeholder,
+          is_by_ref=False)
+      mapped_captures = True
+
+  if not mapped_captures:
+    return concrete_function
+
+  inner_concrete = defun.ConcreteFunction.from_func_graph(
+      concrete_function.graph, concrete_function.function_type, {}
+  )
+
+  def wrap_function(*args):
+    return inner_concrete._call_flat(list(args), inner_concrete.captured_inputs)  # pylint:disable=protected-access
+
+  args = nest.flatten(concrete_function.structured_input_signature,
+                      expand_composites=True)
+  func_graph_module.func_graph_from_py_func(
+      None, wrap_function, args=tuple(args), kwargs={},
+      func_graph=outer_graph)
+
+  # Create concrete function, and copy the attributes necessary to serialize
+  # the function.
+  # pylint: disable=protected-access
+  fn = defun.ConcreteFunction.from_func_graph(
+      outer_graph, concrete_function.function_type, {}
+  )
+  fn._arg_keywords = concrete_function._arg_keywords
+  fn._num_positional_args = concrete_function._num_positional_args
+  # pylint: enable=protected-access
+
+  # Return the captures to their original values
+  for key, capture in remapped_captures.items():
+    external, internal = capture
+    concrete_function.graph._function_captures.add_or_replace(  # pylint: disable=protected-access
+        key=key,
+        external=external,
+        internal=internal,
+        is_by_ref=False)
+  return fn

tensorflow/python/saved_model/keras_injection_test.py

@@ -0,0 +1,38 @@
+# 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.
+# ==============================================================================
+"""Keras injection tests."""
+
+import tensorflow as tf
+
+from tensorflow.python.eager import test
+
+
+# Some of the Keras code load should be triggered so that it will inject proper
+# functionality like registering the optimizer class for SavedModel.
+class KerasInjectionTest(tf.test.TestCase):
+
+  def test_keras_optimizer_injected(self):
+    save_path = test.test_src_dir_path(
+        'cc/saved_model/testdata/OptimizerSlotVariableModule')
+    _ = tf.saved_model.load(save_path)
+    # Make sure keras optimizers are registed without accessing keras code
+    # when loading a model with optimizers
+    self.assertIn(
+        'optimizer', tf.__internal__.saved_model.load.registered_identifiers()
+    )
+
+
+if __name__ == '__main__':
+  tf.test.main()

tensorflow/python/saved_model/load_optimizer_test.py

@@ -0,0 +1,36 @@
+# 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 loading SavedModels with optimizers."""
+
+
+from tensorflow.python.eager import test
+from tensorflow.python.ops import variables
+from tensorflow.python.saved_model import load
+
+
+class LoadOptimizerTest(test.TestCase):
+
+  def test_load_optimizer_without_keras(self):
+    # Make sure that a SavedModel w/ optimizer can be loaded without the Keras
+    # module imported.
+    save_path = test.test_src_dir_path(
+        "cc/saved_model/testdata/OptimizerSlotVariableModule")
+    loaded = load.load(save_path)
+    self.assertIsInstance(
+        loaded.opt.get_slot(loaded.v, "v"), variables.Variable)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/load_options.py

@@ -0,0 +1,125 @@
+# 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 SavedModels."""
+
+from tensorflow.python.saved_model import save_options
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("saved_model.LoadOptions", v1=[])
+class LoadOptions(object):
+  """Options for loading a SavedModel.
+
+  This function may be used in the `options` argument in functions that
+  load a SavedModel (`tf.saved_model.load`, `tf.keras.models.load_model`).
+  """
+
+  # Define object attributes in __slots__ for improved memory and performance.
+  __slots__ = ("allow_partial_checkpoint", "experimental_io_device",
+               "experimental_skip_checkpoint", "experimental_variable_policy",
+               "experimental_load_function_aliases")
+
+  def __init__(self,
+               allow_partial_checkpoint=False,
+               experimental_io_device=None,
+               experimental_skip_checkpoint=False,
+               experimental_variable_policy=None,
+               experimental_load_function_aliases=False):
+    """Creates an object that stores options for SavedModel loading.
+
+    *When to set `allow_partial_checkpoint=True`?*
+
+    This can be used when loading a Keras model (`tf.keras.models.load_model`)
+    with custom objects. When new variables are added to the custom object
+    class, loading will fail the assertion check that all loaded variables have
+    been restored, because the SavedModel checkpoint only contains the variables
+    that were in original the custom object.
+    See the following example:
+
+    ```
+    class Custom(tf.keras.Model):
+      def __init__(self):
+        super(Custom, self).__init__()
+        self.v = tf.Variable(...)
+
+      def call(self, inputs):
+        return ...
+
+    model = Custom()
+    model.save(...)
+    ```
+
+    After saving, say that `Custom` is updated to include an additional
+    variable.
+
+    ```
+    class Custom(tf.keras.Model):
+      def __init__(self):
+        super(Custom, self).__init__()
+        self.v = tf.Variable(...)
+        self.w = tf.Variable(...)
+
+      def call(self, inputs):
+        return ...
+    ```
+
+    `tf.keras.models.load_model(path, custom_objects={'Custom': Custom})` fails
+    to load since `Custom.w` does not exist in the SavedModel checkpoint. To
+    acknowledge that there are variables that are not restored from the
+    checkpoint and successfully load the model, call:
+
+    ```
+    tf.keras.models.load_model(
+      path, custom_objects={'Custom': Custom},
+      options=tf.saved_model.LoadOptions(allow_partial_checkpoint=True))
+    ```
+
+    Args:
+      allow_partial_checkpoint: bool. Defaults to `False`. When enabled, allows
+        the SavedModel checkpoint to not entirely match the loaded object.
+      experimental_io_device: string. Applies in a distributed setting.
+        Tensorflow device to use to access the filesystem. If `None` (default)
+        then for each variable the filesystem is accessed from the CPU:0 device
+        of the host where that variable is assigned. If specified, the
+        filesystem is instead accessed from that device for all variables.
+        This is for example useful if you want to load from a local directory,
+        such as "/tmp" when running in a distributed setting. In that case
+        pass a device for the host where the "/tmp" directory is accessible.
+      experimental_skip_checkpoint: bool. Defaults to `False`. If set to `True`,
+        checkpoints will not be restored. Note that this in the majority of
+        cases will generate an unusable model.
+      experimental_variable_policy: string. The policy to apply to variables
+        when loading. This is either a `saved_model.experimental.VariablePolicy`
+        enum instance or one of its value strings (case is not important). See
+        that enum documentation for details. A value of `None` corresponds to
+        the default policy.
+      experimental_load_function_aliases: bool. Defaults to `False`. If set to
+        `True`, a `function_aliases` attribute will be added to the loaded
+        SavedModel object.
+
+    Example:
+
+      load_options = tf.saved_model.LoadOptions(experimental_io_device=
+        '/job:localhost')
+      restoredmodel = tf.keras.models.load_model(saved_model_path,
+                                                 options=load_options)
+
+    """
+    self.experimental_io_device = experimental_io_device
+    self.allow_partial_checkpoint = allow_partial_checkpoint
+    self.experimental_skip_checkpoint = experimental_skip_checkpoint
+    self.experimental_variable_policy = (
+        save_options.VariablePolicy.from_obj(experimental_variable_policy))
+    self.experimental_load_function_aliases = experimental_load_function_aliases

tensorflow/python/saved_model/load_v1_in_v2.py

@@ -0,0 +1,320 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Import a TF v1-style SavedModel when executing eagerly."""
+
+import functools
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import lift_to_graph
+from tensorflow.python.eager import wrap_function
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import func_graph
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import function_deserialization
+from tensorflow.python.saved_model import loader_impl
+from tensorflow.python.saved_model import signature_serialization
+from tensorflow.python.saved_model.pywrap_saved_model import metrics
+from tensorflow.python.trackable import asset
+from tensorflow.python.trackable import autotrackable
+from tensorflow.python.trackable import resource
+from tensorflow.python.training import monitored_session
+from tensorflow.python.training import saver as tf_saver
+from tensorflow.python.util import nest
+
+# API label for SavedModel metrics.
+_LOAD_V1_V2_LABEL = "load_v1_in_v2"
+
+
+class _Initializer(resource.CapturableResource):
+  """Represents an initialization operation restored from a SavedModel.
+
+  Without this object re-export of imported 1.x SavedModels would omit the
+  original SavedModel's initialization procedure.
+
+  Created when `tf.saved_model.load` loads a TF 1.x-style SavedModel with an
+  initialization op. This object holds a function that runs the
+  initialization. It does not require any manual user intervention;
+  `tf.saved_model.save` will see this object and automatically add it to the
+  exported SavedModel, and `tf.saved_model.load` runs the initialization
+  function automatically.
+  """
+
+  def __init__(self, init_fn, asset_paths):
+    super(_Initializer, self).__init__()
+    self._asset_paths = asset_paths
+    self._init_fn = init_fn
+
+  def _create_resource(self):
+    # Return a constant here so that when re-saved, the traced `create_resource`
+    # has valid returns.
+    return constant_op.constant(1.0)
+
+  def _initialize(self):
+    return self._init_fn(*[path.asset_path for path in self._asset_paths])
+
+
+class _EagerSavedModelLoader(loader_impl.SavedModelLoader):
+  """Loads a SavedModel without using Sessions."""
+
+  def get_meta_graph_def_from_tags(self, tags):
+    """Override to support implicit one-MetaGraph loading with tags=None."""
+    if tags is None:
+      if len(self._saved_model.meta_graphs) != 1:
+        tag_sets = [
+            mg.meta_info_def.tags for mg in self._saved_model.meta_graphs
+        ]
+        raise ValueError(
+            "Importing a SavedModel with `tf.saved_model.load` requires a "
+            "`tags=` argument if there is more than one MetaGraph. Got "
+            f"`tags=None`, but there are {len(self._saved_model.meta_graphs)} "
+            f"MetaGraphs in the SavedModel with tag sets: {tag_sets}. Pass a "
+            "`tags=` argument to load this SavedModel."
+        )
+      return self._saved_model.meta_graphs[0]
+    return super(_EagerSavedModelLoader, self).get_meta_graph_def_from_tags(
+        tags
+    )
+
+  def load_graph(self, returns, meta_graph_def):
+    """Called from wrap_function to import `meta_graph_def`."""
+    # pylint: disable=protected-access
+    saver, _ = tf_saver._import_meta_graph_with_return_elements(meta_graph_def)
+    # pylint: enable=protected-access
+    returns[0] = saver
+
+  def _extract_saver_restore(self, wrapped, saver):
+    if saver is None:
+      return None
+    saver_def = saver.saver_def
+    filename_tensor = wrapped.graph.as_graph_element(
+        saver_def.filename_tensor_name
+    )
+    # We both feed and fetch filename_tensor so we have an operation to use to
+    # feed into variable initializers (only relevant for v1 graph building).
+    return wrapped.prune(
+        feeds=[filename_tensor],
+        fetches=[
+            filename_tensor,
+            wrapped.graph.as_graph_element(saver_def.restore_op_name),
+        ],
+    )
+
+  def restore_variables(self, wrapped, restore_from_saver):
+    """Restores variables from the checkpoint."""
+    if restore_from_saver is not None:
+      initializer, _ = restore_from_saver(
+          constant_op.constant(self._variables_path)
+      )
+      if not ops.executing_eagerly_outside_functions():
+        # Add the initialization operation to the "saved_model_initializers"
+        # collection in case we don't have any lifted variables to attach it to.
+        ops.add_to_collection("saved_model_initializers", initializer)
+        one_unlifted = False
+
+        for variable in wrapped.graph.get_collection_ref(
+            ops.GraphKeys.GLOBAL_VARIABLES
+        ):
+          if variable.graph is wrapped.graph:
+            one_unlifted = True
+          # pylint: disable=protected-access
+          variable._initializer_op = initializer
+          # pylint: enable=protected-access
+        if one_unlifted:
+          logging.warning(
+              "Some variables could not be lifted out of a loaded function. "
+              "Please run "
+              '`sess.run(tf.get_collection("saved_model_initializers"))`to '
+              "restore these variables."
+          )
+
+  def _extract_signatures(self, wrapped, meta_graph_def):
+    """Creates ConcreteFunctions for signatures in `meta_graph_def`."""
+    signature_functions = {}
+    for signature_key, signature_def in meta_graph_def.signature_def.items():
+      if signature_def.inputs:
+        input_items = sorted(
+            signature_def.inputs.items(), key=lambda item: item[0]
+        )
+        original_input_names, input_specs = zip(*input_items)
+      else:
+        original_input_names = []
+        input_specs = []
+
+      # TODO(b/205015292): Support optional arguments
+      feeds = [
+          wrap_function._get_element_from_tensor_info(input_spec, wrapped.graph)  # pylint: disable=protected-access
+          for input_spec in input_specs
+      ]
+      input_names = []
+      input_tensors = []
+      for original_input_name, feed in zip(original_input_names, feeds):
+        if isinstance(feed, sparse_tensor.SparseTensor):
+          # We have to give explicit name for SparseTensor arguments, because
+          # these are not present in the TensorInfo.
+          indices_name = "%s_indices" % original_input_name
+          values_name = "%s_values" % original_input_name
+          dense_shape_name = "%s_dense_shape" % original_input_name
+          input_names.extend([indices_name, values_name, dense_shape_name])
+          input_tensors.extend([feed.indices, feed.values, feed.dense_shape])
+        elif isinstance(feed, composite_tensor.CompositeTensor):
+          component_tensors = nest.flatten(feed, expand_composites=True)
+          input_names.extend(
+              "%s_component_%d" % (original_input_name, n)
+              for n in range(len(component_tensors))
+          )
+          input_tensors.extend(component_tensors)
+        else:
+          input_names.append(original_input_name)
+          input_tensors.append(feed)
+      fetches = {name: out for name, out in signature_def.outputs.items()}
+      input_signature = (
+          (),
+          func_graph.convert_structure_to_signature(
+              dict(zip(input_names, input_tensors))
+          ),
+      )
+      try:
+        signature_fn = wrapped.prune(
+            feeds=feeds,
+            fetches=fetches,
+            input_signature=input_signature,
+            are_keyword_args_also_positional=True,
+        )
+      except lift_to_graph.UnliftableError as ex:
+        # Mutate the exception to add a bit more detail.
+        args = ex.args
+        if not args:
+          message = ""
+        else:
+          message = args[0]
+        message = (
+            "A SavedModel signature needs an input for each placeholder the "
+            "signature's outputs use. An output for signature '{}' depends on "
+            "a placeholder which is not an input (i.e. the placeholder is not "
+            "fed a value).\n\n"
+        ).format(signature_key) + message
+        ex.args = (message,) + args[1:]
+        raise
+
+      # pylint: disable=protected-access
+      signature_fn._arg_keywords = input_names
+
+      if len(input_names) == 1:
+        # Allowing positional arguments does not create any ambiguity if there's
+        # only one.
+        signature_fn._num_positional_args = 1
+      else:
+        signature_fn._num_positional_args = 0
+      # pylint: enable=protected-access
+      signature_functions[signature_key] = signature_fn
+    return signature_functions
+
+  def load(self, tags, skip_restoring_checkpoint=False):
+    """Creates an object from the MetaGraph identified by `tags`."""
+    meta_graph_def = self.get_meta_graph_def_from_tags(tags)
+    load_shared_name_suffix = "_load_{}".format(ops.uid())
+    functions = function_deserialization.load_function_def_library(
+        meta_graph_def.graph_def.library,
+        load_shared_name_suffix=load_shared_name_suffix,
+    )
+    # Replace existing functions in the MetaGraphDef with renamed functions so
+    # we don't have duplicates or name collisions.
+    meta_graph_def.graph_def.library.Clear()
+    for function in functions.values():
+      meta_graph_def.graph_def.library.function.add().CopyFrom(
+          function.function_def
+      )
+    # We've renamed functions and shared names. We need the same operation on
+    # the GraphDef itself for consistency.
+    for node_def in meta_graph_def.graph_def.node:
+      function_deserialization.fix_node_def(
+          node_def, functions, load_shared_name_suffix
+      )
+
+    load_graph_returns = [None]
+    wrapped = wrap_function.wrap_function(
+        functools.partial(self.load_graph, load_graph_returns, meta_graph_def),
+        signature=[],
+    )
+    (saver,) = load_graph_returns
+    restore_from_saver = self._extract_saver_restore(wrapped, saver)
+
+    if not skip_restoring_checkpoint:
+      self.restore_variables(wrapped, restore_from_saver)
+
+    with wrapped.graph.as_default():
+      init_op = (
+          loader_impl.get_init_op(meta_graph_def)
+          or monitored_session.Scaffold.default_local_init_op()
+      )
+      # Add a dummy Tensor we know we can fetch to add control dependencies to.
+      init_anchor = constant_op.constant(0.0, name="dummy_fetch")
+
+    root = autotrackable.AutoTrackable()
+    if restore_from_saver is not None:
+      root.restore = lambda path: restore_from_saver(constant_op.constant(path))
+    asset_feed_tensors = []
+    asset_paths = []
+    for tensor_name, value in loader_impl.get_asset_tensors(
+        self._export_dir, meta_graph_def
+    ).items():
+      asset_feed_tensors.append(wrapped.graph.as_graph_element(tensor_name))
+      asset_paths.append(asset.Asset(value))
+    init_fn = wrapped.prune(
+        feeds=asset_feed_tensors,
+        fetches=[init_anchor, wrapped.graph.as_graph_element(init_op)],
+    )
+    initializer = _Initializer(init_fn, asset_paths)
+    # pylint: disable=protected-access
+    local_init_op, _ = initializer._initialize()
+    # pylint: enable=protected-access
+    with ops.init_scope():
+      if not context.executing_eagerly():
+        ops.add_to_collection(ops.GraphKeys.TABLE_INITIALIZERS, local_init_op)
+        for variable in wrapped.graph.get_collection_ref(
+            ops.GraphKeys.LOCAL_VARIABLES
+        ):
+          # pylint: disable=protected-access
+          variable._initializer_op = local_init_op
+          # pylint: enable=protected-access
+    root.initializer = initializer
+    root.asset_paths = asset_paths
+    signature_functions = self._extract_signatures(wrapped, meta_graph_def)
+
+    root.signatures = signature_serialization.create_signature_map(
+        signature_functions
+    )
+    root.variables = list(wrapped.graph.variables)
+    root.tensorflow_version = meta_graph_def.meta_info_def.tensorflow_version
+    root.tensorflow_git_version = (
+        meta_graph_def.meta_info_def.tensorflow_git_version
+    )
+    root.graph = wrapped.graph
+    root.prune = wrapped.prune
+    return root
+
+
+def load(export_dir, tags, skip_restoring_checkpoint=False):
+  """Load a v1-style SavedModel as an object."""
+  metrics.IncrementReadApi(_LOAD_V1_V2_LABEL)
+  loader = _EagerSavedModelLoader(export_dir)
+  result = loader.load(
+      tags=tags, skip_restoring_checkpoint=skip_restoring_checkpoint
+  )
+  metrics.IncrementRead(write_version="1")
+  return result

tensorflow/python/saved_model/load_v1_in_v2_test.py

@@ -0,0 +1,961 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for importing a TF v1-style SavedModel when executing eagerly."""
+
+import os
+import shutil
+
+from absl.testing import parameterized
+
+from tensorflow.core.framework import variable_pb2
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import lift_to_graph
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function as framework_function
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework import versions
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import partitioned_variables
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import ref_variable
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.ops import while_loop
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.saved_model import builder_impl
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import save
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import simple_save
+from tensorflow.python.saved_model import tag_constants
+from tensorflow.python.saved_model import utils_impl
+from tensorflow.python.training import saver
+
+
+class LoadTest(test.TestCase, parameterized.TestCase):
+
+  def _v1_single_metagraph_saved_model(self, use_resource):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      start = array_ops.placeholder(
+          shape=None, dtype=dtypes.float32, name="start"
+      )
+      if use_resource:
+        distractor = ref_variable.RefVariable(-1.0, name="distractor")
+        v = resource_variable_ops.ResourceVariable(3.0, name="v")
+      else:
+        # "distractor" gets saved in the checkpoint and so used in the restore
+        # function, but not in the pruned function for the signature. This tests
+        # node naming: it needs to be consistent (and ideally always the same as
+        # the node in the original GraphDef) for the resource manager to find
+        # the right variable.
+        distractor = ref_variable.RefVariable(-1.0, name="distractor")
+        v = ref_variable.RefVariable(3.0, name="v")
+      local_variable = variable_v1.VariableV1(
+          1.0,
+          collections=[ops.GraphKeys.LOCAL_VARIABLES],
+          trainable=False,
+          use_resource=True,
+      )
+      output = array_ops.identity(start * v * local_variable, name="output")
+      with session_lib.Session() as session:
+        session.run(
+            [v.initializer, distractor.initializer, local_variable.initializer]
+        )
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session,
+            path,
+            inputs={"start": start},
+            outputs={"output": output},
+            legacy_init_op=local_variable.initializer,
+        )
+    return path
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_pretty_printed_signature(self):
+    imported = load.load(
+        self._v1_single_metagraph_saved_model(use_resource=True)
+    )
+    self.evaluate(variables.global_variables_initializer())
+    self.evaluate(variables.local_variables_initializer())
+    concrete_fn = imported.signatures["serving_default"]
+
+    summary = (
+        "(start: TensorSpec(shape=<unknown>, dtype=tf.float32,"
+        " name='start')) -> Dict[['output', TensorSpec(shape=<unknown>,"
+        " dtype=tf.float32, name=None)]]"
+    )
+    details = (
+        r"Input Parameters:\n"
+        r"  start \(POSITIONAL_OR_KEYWORD\): TensorSpec\(shape=<unknown>,"
+        r" dtype=tf\.float32, name='start'\)\n"
+        r"Output Type:\n"
+        r"  Dict\[\['output', TensorSpec\(shape=<unknown>,"
+        r" dtype=tf\.float32, name=None\)\]\]\n"
+        r"Captures:\n"
+        r"  \d+: TensorSpec\(shape=\(\), dtype=tf\.resource, name=None\)\n"
+        r"  \d+: TensorSpec\(shape=\(\), dtype=tf\.resource, name=None\)"
+    )
+    self.assertEqual(
+        concrete_fn.pretty_printed_signature(verbose=False), summary
+    )
+    self.assertRegex(
+        concrete_fn.pretty_printed_signature(verbose=True), details
+    )
+    self.assertRegex(repr(concrete_fn), r"<ConcreteFunction .* at .*")
+    self.assertRegex(str(concrete_fn), r"ConcreteFunction " + details)
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_resource_variable_import(self):
+    imported = load.load(
+        self._v1_single_metagraph_saved_model(use_resource=True)
+    )
+    self.evaluate(variables.global_variables_initializer())
+    self.evaluate(variables.local_variables_initializer())
+    fn = imported.signatures["serving_default"]
+    self.assertEqual(
+        {"output": 6.0}, self.evaluate(fn(constant_op.constant(2.0)))
+    )
+    self.assertAllEqual([3.0, 1.0], self.evaluate(imported.variables))
+    self.evaluate(imported.variables[0].assign(4.0))
+    self.assertEqual(
+        {"output": 8.0}, self.evaluate(fn(start=constant_op.constant(2.0)))
+    )
+    self.evaluate(imported.variables[1].assign(2.0))
+    self.assertEqual(
+        {"output": 24.0}, self.evaluate(fn(start=constant_op.constant(3.0)))
+    )
+    self.assertTrue(imported.variables[0].trainable)
+    self.assertFalse(imported.variables[1].trainable)
+    with backprop.GradientTape() as tape:
+      output = fn(start=constant_op.constant(4.0))
+    self.assertEqual(imported.variables[:1], list(tape.watched_variables()))
+    self.assertEqual(
+        8.0, self.evaluate(tape.gradient(output, imported.variables[0]))
+    )
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_ref_variable_import(self):
+    saved = self._v1_single_metagraph_saved_model(use_resource=False)
+    imported = load.load(saved)
+    fn = imported.signatures["serving_default"]
+    self.evaluate(lookup_ops.tables_initializer())
+    self.evaluate(ops.get_collection("saved_model_initializers"))
+    self.assertEqual(
+        6.0, self.evaluate(fn(start=constant_op.constant(2.0))["output"])
+    )
+
+  def _v1_output_shape_saved_model(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      start = array_ops.placeholder(
+          shape=[None], dtype=dtypes.float32, name="start"
+      )
+      output = array_ops.identity(start, name="output")
+      output.set_shape([1])  # Ok to use [1] because shape is only informational
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        builder = builder_impl.SavedModelBuilder(path)
+        builder.add_meta_graph_and_variables(
+            session,
+            tags=[tag_constants.SERVING],
+            signature_def_map={
+                "serving_default": signature_def_utils.build_signature_def(
+                    {"start": utils_impl.build_tensor_info(start)},
+                    {"output": utils_impl.build_tensor_info(output)},
+                )
+            },
+        )
+        builder.save()
+    return path
+
+  def test_restore_output_shapes(self):
+    saved = self._v1_output_shape_saved_model()
+    imported = load.load(saved)
+    fn = imported.signatures["serving_default"]
+    self.assertEqual(tensor_shape.TensorShape([1]), fn.outputs[0].shape)
+
+  def _v1_multi_metagraph_saved_model(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      start = array_ops.placeholder(
+          shape=[None], dtype=dtypes.float32, name="start"
+      )
+      v = resource_variable_ops.ResourceVariable(21.0)
+      first_output = array_ops.identity(start * v, name="first_output")
+      second_output = array_ops.identity(v, name="second_output")
+      with session_lib.Session() as session:
+        session.run(v.initializer)
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        builder = builder_impl.SavedModelBuilder(path)
+        builder.add_meta_graph_and_variables(
+            session,
+            tags=["first"],
+            signature_def_map={
+                "first_key": signature_def_utils.build_signature_def(
+                    {"first_start": utils_impl.build_tensor_info(start)},
+                    {
+                        "first_output": utils_impl.build_tensor_info(
+                            first_output
+                        )
+                    },
+                )
+            },
+        )
+        builder.add_meta_graph(
+            tags=["second"],
+            signature_def_map={
+                "second_key": signature_def_utils.build_signature_def(
+                    {"second_start": utils_impl.build_tensor_info(start)},
+                    {
+                        "second_output": utils_impl.build_tensor_info(
+                            second_output
+                        )
+                    },
+                )
+            },
+        )
+        builder.save()
+    return path
+
+  def test_multi_meta_graph_loading(self):
+    with self.assertRaisesRegex(ValueError, "2 MetaGraphs"):
+      load.load(self._v1_multi_metagraph_saved_model())
+    first_imported = load.load(
+        self._v1_multi_metagraph_saved_model(), tags=["first"]
+    )
+    self.assertEqual(
+        {"first_output": 42.0},
+        self.evaluate(
+            first_imported.signatures["first_key"](
+                first_start=constant_op.constant(2.0)
+            )
+        ),
+    )
+    second_imported = load.load(
+        self._v1_multi_metagraph_saved_model(), tags=set(["second"])
+    )
+    with self.assertRaisesRegex(TypeError, "second_start"):
+      second_imported.signatures["second_key"](x=constant_op.constant(2.0))
+    with self.assertRaisesRegex(TypeError, "second_start"):
+      second_imported.signatures["second_key"](
+          second_start=constant_op.constant(2.0), x=constant_op.constant(2.0)
+      )
+    self.assertEqual(
+        {"second_output": 21.0},
+        self.evaluate(
+            second_imported.signatures["second_key"](
+                second_start=constant_op.constant(2.0)
+            )
+        ),
+    )
+
+  def _v1_asset_saved_model(self, clear_shared_name):
+    export_graph = ops.Graph()
+    vocab_path = os.path.join(self.get_temp_dir(), "vocab.txt")
+    with open(vocab_path, "w") as f:
+      f.write("alpha\nbeta\ngamma\n")
+    with export_graph.as_default():
+      initializer = lookup_ops.TextFileInitializer(
+          vocab_path,
+          key_dtype=dtypes.string,
+          key_index=lookup_ops.TextFileIndex.WHOLE_LINE,
+          value_dtype=dtypes.int64,
+          value_index=lookup_ops.TextFileIndex.LINE_NUMBER,
+      )
+      table = lookup_ops.HashTable(initializer, default_value=-1)
+      start = array_ops.placeholder(shape=None, dtype=dtypes.string, name="in")
+      output = table.lookup(start, name="out")
+      if clear_shared_name:
+        export_graph.get_operation_by_name("hash_table")._clear_attr(
+            "shared_name"
+        )
+      with session_lib.Session() as session:
+        session.run([table.initializer])
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session,
+            path,
+            inputs={"start": start},
+            outputs={"output": output},
+            legacy_init_op=table.initializer,
+        )
+    file_io.delete_file(vocab_path)
+    return path
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_asset_loading(self):
+    first_path = self._v1_asset_saved_model(clear_shared_name=False)
+    imported = load.load(first_path)
+    self.evaluate(lookup_ops.tables_initializer())
+    fn = imported.signatures["serving_default"]
+    self.assertAllClose(
+        {"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))
+    )
+    second_path = os.path.join(
+        self.get_temp_dir(), "saved_model", str(ops.uid())
+    )
+    save.save(imported, second_path, signatures=imported.signatures)
+    shutil.rmtree(first_path)
+    del ops.get_collection_ref(ops.GraphKeys.TABLE_INITIALIZERS)[:]
+    second_import = load.load(second_path)
+    self.evaluate(lookup_ops.tables_initializer())
+    fn = second_import.signatures["serving_default"]
+    self.assertAllClose(
+        {"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))
+    )
+
+    third_path = os.path.join(
+        self.get_temp_dir(), "saved_model", str(ops.uid())
+    )
+    save.save(second_import, third_path, signatures=second_import.signatures)
+    shutil.rmtree(second_path)
+    del ops.get_collection_ref(ops.GraphKeys.TABLE_INITIALIZERS)[:]
+    third_import = load.load(third_path)
+    self.evaluate(lookup_ops.tables_initializer())
+    fn = third_import.signatures["serving_default"]
+    self.assertAllClose(
+        {"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))
+    )
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_node_name_sharing(self):
+    fourth_path = self._v1_asset_saved_model(clear_shared_name=True)
+    fourth_import = load.load(fourth_path)
+    self.evaluate(lookup_ops.tables_initializer())
+    fn = fourth_import.signatures["serving_default"]
+    self.assertAllClose(
+        {"output": [2, 0]}, fn(start=constant_op.constant(["gamma", "alpha"]))
+    )
+
+  def _v1_cond_saved_model(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      branch_selector = array_ops.placeholder(
+          name="branch_selector", shape=[], dtype=dtypes.bool
+      )
+      output = cond.cond(
+          branch_selector,
+          lambda: array_ops.ones([]),
+          lambda: array_ops.zeros([]),
+      )
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session,
+            path,
+            inputs={"branch_selector": branch_selector},
+            outputs={"output": output},
+        )
+    return path
+
+  def test_cond(self):
+    first_path = self._v1_cond_saved_model()
+    imported = load.load(first_path)
+    function = imported.signatures["serving_default"]
+    self.assertAllClose({"output": 1.0}, function(constant_op.constant(True)))
+    self.assertAllClose({"output": 0.0}, function(constant_op.constant(False)))
+
+  def _v1_while_saved_model(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      loop_iterations = array_ops.placeholder(
+          name="loop_iterations", shape=[], dtype=dtypes.int32
+      )
+      _, output = while_loop.while_loop(
+          lambda index, accum: index <= loop_iterations,
+          lambda index, accum: (index + 1, accum + index),
+          [constant_op.constant(0),
+           constant_op.constant(0)],
+      )
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session,
+            path,
+            inputs={"loop_iterations": loop_iterations},
+            outputs={"output": output},
+        )
+    return path
+
+  def test_while(self):
+    first_path = self._v1_while_saved_model()
+    imported = load.load(first_path)
+    function = imported.signatures["serving_default"]
+    self.assertAllClose({"output": 10}, function(constant_op.constant(4)))
+    self.assertAllClose({"output": 15}, function(constant_op.constant(5)))
+
+  def _v1_nested_while_saved_model(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+
+      def _inner_while(loop_iterations):
+        _, output = while_loop.while_loop(
+            lambda index, accum: index <= loop_iterations,
+            lambda index, accum: (index + 1, accum + index),
+            [constant_op.constant(0),
+             constant_op.constant(0)],
+        )
+        return output
+
+      loop_iterations = array_ops.placeholder(
+          name="loop_iterations", shape=[], dtype=dtypes.int32
+      )
+      _, output = while_loop.while_loop(
+          lambda index, accum: index <= loop_iterations,
+          lambda index, accum: (index + 1, accum + _inner_while(index)),
+          [constant_op.constant(0),
+           constant_op.constant(0)],
+      )
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session,
+            path,
+            inputs={"loop_iterations": loop_iterations},
+            outputs={"output": output},
+        )
+    return path
+
+  def test_nested_while(self):
+    first_path = self._v1_nested_while_saved_model()
+    imported = load.load(first_path)
+    function = imported.signatures["serving_default"]
+    self.assertAllClose({"output": 20}, function(constant_op.constant(4)))
+    self.assertAllClose({"output": 35}, function(constant_op.constant(5)))
+
+  def _no_signatures_model(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      inp = array_ops.placeholder(name="x", shape=[], dtype=dtypes.float32)
+      array_ops.identity(inp + 1.0, name="out")
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        b = builder_impl.SavedModelBuilder(path)
+        b.add_meta_graph_and_variables(
+            session,
+            tags=[tag_constants.SERVING],
+            signature_def_map={},
+            assets_collection=ops.get_collection(ops.GraphKeys.ASSET_FILEPATHS),
+        )
+        b.save()
+    return path
+
+  def test_no_signature(self):
+    path = self._no_signatures_model()
+    imported = load.load(path)
+    self.assertEqual([], list(imported.signatures.keys()))
+
+  def _signature_with_no_inputs(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      array_ops.placeholder(name="x", shape=[], dtype=dtypes.float32)
+      output = random_ops.random_normal([2])
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        b = builder_impl.SavedModelBuilder(path)
+        b.add_meta_graph_and_variables(
+            session,
+            tags=[tag_constants.SERVING],
+            signature_def_map={
+                "key": signature_def_utils.build_signature_def(
+                    {}, dict(value=utils_impl.build_tensor_info(output))
+                )
+            },
+        )
+        b.save()
+    return path
+
+  def test_signature_with_no_inputs(self):
+    path = self._signature_with_no_inputs()
+    imported = load.load(path)
+    self.assertEqual([2], imported.signatures["key"]()["value"].shape)
+
+  def test_version_info(self):
+    path = self._signature_with_no_inputs()
+    imported = load.load(path)
+    self.assertEqual(versions.__version__, imported.tensorflow_version)
+    self.assertEqual(versions.__git_version__, imported.tensorflow_git_version)
+
+  def _unfed_placeholder_signature(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      x = array_ops.placeholder(name="x", shape=[], dtype=dtypes.float32)
+      output = x * random_ops.random_normal([2])
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        b = builder_impl.SavedModelBuilder(path)
+        b.add_meta_graph_and_variables(
+            session,
+            tags=[tag_constants.SERVING],
+            signature_def_map={
+                "key": signature_def_utils.build_signature_def(
+                    {}, dict(value=utils_impl.build_tensor_info(output))
+                )
+            },
+        )
+        b.save()
+    return path
+
+  def test_unfed_placeholder_exception(self):
+    path = self._unfed_placeholder_signature()
+    with self.assertRaisesRegex(
+        lift_to_graph.UnliftableError,
+        "signature needs an input for each placeholder.*\n\nUnable to lift",
+    ):
+      load.load(path)
+
+  def test_custom_pruning(self):
+    path = self._no_signatures_model()
+    root = load.load(path)
+    fn = root.prune("x:0", "out:0")
+    self.assertEqual(2.0, self.evaluate(fn(x=array_ops.ones([]))))
+    root.graph.as_graph_element("x:0")
+
+  def _no_trainable_variable_attribute(self, trainable):
+    """A SavedModel where the VariableDef has no 'trainable' (it's false)."""
+
+    class _MissingFieldsVariable(resource_variable_ops.ResourceVariable):
+
+      def to_proto(self, export_scope=None):
+        full_proto = super(_MissingFieldsVariable, self).to_proto(export_scope)
+        return variable_pb2.VariableDef(
+            variable_name=full_proto.variable_name,
+            initial_value_name=full_proto.initial_value_name,
+            initializer_name=full_proto.snapshot_name,
+            save_slice_info_def=full_proto.save_slice_info_def,
+            is_resource=full_proto.is_resource,
+        )
+
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      v = _MissingFieldsVariable(3.0, trainable=trainable)
+      with session_lib.Session() as session:
+        session.run([v.initializer])
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        b = builder_impl.SavedModelBuilder(path)
+        b.add_meta_graph_and_variables(
+            session, tags=[tag_constants.SERVING], signature_def_map={}
+        )
+        b.save()
+
+    return path
+
+  def test_trainable_not_set_in_proto(self):
+    """If a VariableDef has no 'trainable', we fall back to collections."""
+    real_tf_version = versions.__version__
+    # Pretend to be exported from an older version of TensorFlow, so trainable
+    # will follow collections instead of checking VariableDefs.
+    versions.__version__ = "1.7.0"
+    path = self._no_trainable_variable_attribute(trainable=True)
+    root = load.load(path)
+    self.assertTrue(root.variables[0].trainable)
+    path = self._no_trainable_variable_attribute(trainable=False)
+    root = load.load(path)
+    self.assertFalse(root.variables[0].trainable)
+    versions.__version__ = real_tf_version
+
+  def _export_variable(self, **kwargs_for_variable):
+    """A 1.x SavedModel with a single variable."""
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      v = resource_variable_ops.ResourceVariable(3.0, **kwargs_for_variable)
+      with session_lib.Session() as session:
+        session.run([v.initializer])
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        b = builder_impl.SavedModelBuilder(path)
+        b.add_meta_graph_and_variables(
+            session, tags=[tag_constants.SERVING], signature_def_map={}
+        )
+        b.save()
+
+    return path
+
+  def test_trainable_in_proto(self):
+    """If a VariableDef has a trainable property, we do not use collections."""
+    path = self._export_variable(
+        trainable=True, collections=[ops.GraphKeys.GLOBAL_VARIABLES]
+    )
+    root = load.load(path)
+    self.assertTrue(root.variables[0].trainable)
+    path = self._export_variable(
+        trainable=False,
+        collections=[
+            ops.GraphKeys.GLOBAL_VARIABLES,
+            ops.GraphKeys.TRAINABLE_VARIABLES,
+        ],
+    )
+    root = load.load(path)
+    self.assertFalse(root.variables[0].trainable)
+
+  def _model_with_sparse_output(self):
+    """Generate a graph with a SparseTensor output and serialize in V1 format"""
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      in_placeholder = array_ops.placeholder(dtype=dtypes.int64, shape=[1])
+      out_sparse_tensor = (
+          sparse_tensor.SparseTensor(
+              indices=[[0]], values=in_placeholder, dense_shape=[1]
+          )
+          * 2
+      )
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session,
+            path,
+            inputs={"start": in_placeholder},
+            outputs={"output": out_sparse_tensor},
+        )
+    return path
+
+  def test_load_sparse_outputs(self):
+    path = self._model_with_sparse_output()
+    imported = load.load(path)
+    imported_fn = imported.signatures["serving_default"]
+    forty_two = constant_op.constant([42], dtype=dtypes.int64)
+    self.assertEqual([84], imported_fn(forty_two)["output"].values.numpy())
+
+  def _model_with_sparse_input(self):
+    """Generate a graph with a SparseTensor input and serialize in V1 format."""
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      in_sparse_placeholder = array_ops.sparse_placeholder(
+          dtype=dtypes.int64, shape=[2, 2]
+      )
+      out_sparse_tensor = (
+          sparse_tensor.SparseTensor(
+              indices=in_sparse_placeholder.indices,
+              values=in_sparse_placeholder.values,
+              dense_shape=in_sparse_placeholder.dense_shape,
+          )
+          * 2
+      )
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session,
+            path,
+            inputs={"start": in_sparse_placeholder},
+            outputs={"output": out_sparse_tensor},
+        )
+    return path
+
+  def test_load_sparse_inputs(self):
+    path = self._model_with_sparse_input()
+    imported = load.load(path)
+    imported_fn = imported.signatures["serving_default"]
+    indices = constant_op.constant([[0, 0], [0, 1], [1, 1]], dtype=dtypes.int64)
+    values = constant_op.constant([42, 43, 44], dtype=dtypes.int64)
+    dense_shape = constant_op.constant([2, 2], dtype=dtypes.int64)
+    result = imported_fn(
+        start_indices=indices,
+        start_values=values,
+        start_dense_shape=dense_shape,
+    )
+    self.assertAllEqual([84, 86, 88], result["output"].values.numpy())
+
+  def _model_with_ragged_input(self):
+    """Generate a graph with a RaggedTensor input and serialize in V1 format."""
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      x = ragged_factory_ops.placeholder(dtypes.float32, 1, [])
+      y = x * 2
+      with session_lib.Session() as sess:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(sess, path, inputs={"x": x}, outputs={"y": y})
+    return path
+
+  def test_load_ragged_inputs(self):
+    path = self._model_with_ragged_input()
+    imported = load.load(path)
+    imported_fn = imported.signatures["serving_default"]
+    x = ragged_factory_ops.constant([[10.0, 20.0], [30.0]])
+    result = imported_fn(x_component_0=x.values, x_component_1=x.row_splits)
+    self.assertAllEqual(result["y"], [[20.0, 40.0], [60.0]])
+
+  def _model_with_defun(self):
+    """Generate a graph with a Defun and serialize in V1 format."""
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+
+      @framework_function.Defun(dtypes.int64)
+      def z(x):
+        return x + 1
+
+      @framework_function.Defun(dtypes.int64)
+      def g(x):
+        return z(x) + 1
+
+      @framework_function.Defun(dtypes.int64)
+      def f(x):
+        return g(x) + 1
+
+      in_placeholder = array_ops.placeholder(dtype=dtypes.int64, shape=[1])
+      out = f(in_placeholder)
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session,
+            path,
+            inputs={"start": in_placeholder},
+            outputs={"output": out},
+        )
+    return path
+
+  def test_load_defun(self):
+    path = self._model_with_defun()
+    imported = load.load(path)
+    imported_fn = imported.signatures["serving_default"]
+    forty_two = constant_op.constant([42], dtype=dtypes.int64)
+    self.assertEqual([45], imported_fn(forty_two)["output"].numpy())
+
+  def test_load_and_restore_partitioned_variables(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      partitioned_var = variable_scope.get_variable(
+          "a",
+          shape=[6],
+          initializer=init_ops.constant_initializer(13),
+          partitioner=partitioned_variables.fixed_size_partitioner(2),
+          use_resource=True,
+      )
+      x = array_ops.placeholder(shape=[], dtype=dtypes.float32)
+      y = x * partitioned_var
+      with session_lib.Session() as session:
+        session.run(variables.global_variables_initializer())
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session, path, inputs={"x": x}, outputs={"y": y}
+        )
+
+        # Create a name-based checkpoint with different values.
+        session.run(partitioned_var.assign([[5, 4, 3], [2, 1, 0]]))
+        ckpt_path = os.path.join(self.get_temp_dir(), "restore_ckpt")
+        saver.Saver().save(session, ckpt_path)
+
+    imported = load.load(path)
+    self.assertAllClose(
+        self.evaluate(imported.variables), [[13, 13, 13], [13, 13, 13]]
+    )
+
+    self.evaluate(imported.restore(ckpt_path))
+    self.assertAllClose(
+        self.evaluate(imported.variables), [[5, 4, 3], [2, 1, 0]]
+    )
+    self.assertAllClose(
+        self.evaluate(
+            imported.signatures["serving_default"](constant_op.constant(2.0))
+        ),
+        {"y": [10, 8, 6, 4, 2, 0]},
+    )
+
+  def test_structured_input_signature(self):
+    path = self._v1_single_metagraph_saved_model(False)
+    imported = load.load(path)
+    args, kwargs = imported.signatures[
+        "serving_default"
+    ].structured_input_signature
+    self.assertEqual(args, ())
+    self.assertAllEqual(
+        kwargs, {"start": tensor_spec.TensorSpec(shape=None, name="start")}
+    )
+
+  def _model_with_multiple_inputs(self, input_names, compute_fn, var_value):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      inputs = tuple(
+          array_ops.placeholder(shape=(), dtype=dtypes.float32, name=name)
+          for name in input_names
+      )
+      v = resource_variable_ops.ResourceVariable(var_value)
+      output = array_ops.identity(compute_fn(inputs, v), name="output")
+      with session_lib.Session() as session:
+        session.run(v.initializer)
+        path = os.path.join(
+            self.get_temp_dir(), "tf1_saved_model", str(ops.uid())
+        )
+        builder = builder_impl.SavedModelBuilder(path)
+        feeds = {
+            name: utils_impl.build_tensor_info(input)
+            for name, input in zip(input_names, inputs)
+        }
+        builder.add_meta_graph_and_variables(
+            session,
+            tags=[tag_constants.SERVING],
+            signature_def_map={
+                "serving_default": signature_def_utils.build_signature_def(
+                    feeds,
+                    {"output": utils_impl.build_tensor_info(output)},
+                )
+            },
+        )
+        builder.save()
+    return path
+
+  @parameterized.named_parameters(
+      (f"_{input_names_idx}_{reverse}", input_names, reverse)  # pylint: disable=g-complex-comprehension
+      for reverse in (False, True)
+      for input_names_idx, input_names in enumerate((
+          ("input1", "input2"),
+          ("input1", "input./-"),
+      ))
+  )
+  def test_multiple_inputs(self, input_names, reverse):
+    if reverse:
+      input_names = tuple(reversed(input_names))
+
+    def compute_fn(ls, a):
+      result = a
+      for x in ls:
+        result = (result + x) * a
+      return result
+
+    var_value = 21.0
+    path = self._model_with_multiple_inputs(
+        input_names, compute_fn=compute_fn, var_value=21.0
+    )
+    imported = load.load(path)
+    sorted_with_idx = sorted(
+        zip(range(len(input_names)), input_names), key=lambda x: x[1]
+    )
+    fn = imported.signatures["serving_default"]
+    for i, (_, input_name) in enumerate(sorted_with_idx):
+      self.assertEqual(fn.inputs[i].name, f"{input_name}:0")
+    inputs = tuple(i + 2.0 for i in range(len(input_names)))
+    expected_output = compute_fn(inputs, var_value)
+    # Call `fn`` with keyword arguments
+    self.assertEqual(
+        self.evaluate(
+            fn(**{
+                name: constant_op.constant(v)
+                for name, v in zip(input_names, inputs)
+            })["output"]
+        ),
+        expected_output,
+    )
+    # Call `fn`` with positional arguments
+    self.assertEqual(
+        self.evaluate(fn(*(inputs[i] for i, _ in sorted_with_idx))["output"]),
+        expected_output,
+    )
+
+    # Test saving the model again in TF2
+    path2 = os.path.join(self.get_temp_dir(), "tf2_saved_model", str(ops.uid()))
+    save.save(imported, path2, imported.signatures)
+
+    imported2 = load.load(path2)
+    fn = imported2.signatures["serving_default"]
+    # Call `fn`` with keyword arguments
+    self.assertEqual(
+        self.evaluate(
+            fn(**{
+                name: constant_op.constant(v)
+                for name, v in zip(input_names, inputs)
+            })["output"]
+        ),
+        expected_output,
+    )
+    # `fn` can no longer be called with positional arguments, because
+    # during TF2 saving, those positional-keyword-hybrid arguments are
+    # converted to keyword-only arguments.
+
+  def _v1_multi_input_saved_model(self):
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      input1 = array_ops.placeholder(
+          shape=[None], dtype=dtypes.float32, name="input1"
+      )
+      input2 = array_ops.placeholder(
+          shape=[None], dtype=dtypes.float32, name="input2"
+      )
+      v = resource_variable_ops.ResourceVariable(21.0)
+      output = array_ops.identity(input1 * v + input2, name="output")
+      with session_lib.Session() as session:
+        session.run(v.initializer)
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        builder = builder_impl.SavedModelBuilder(path)
+        builder.add_meta_graph_and_variables(
+            session,
+            tags=[tag_constants.SERVING],
+            signature_def_map={
+                "serving_default": signature_def_utils.build_signature_def(
+                    {
+                        "input1": utils_impl.build_tensor_info(input1),
+                        "input2": utils_impl.build_tensor_info(input2),
+                    },
+                    {"output": utils_impl.build_tensor_info(output)},
+                )
+            },
+        )
+        builder.save()
+    return path
+
+  def test_v1_input_ordered(self):
+    path = self._v1_multi_input_saved_model()
+    imported = load.load(path)
+    self.assertEqual(
+        imported.signatures["serving_default"].inputs[0].name, "input1:0"
+    )
+    self.assertEqual(
+        imported.signatures["serving_default"].inputs[1].name, "input2:0"
+    )
+
+  def test_resave_signature(self):
+    # Tests that signatures saved using TF1 can be resaved with TF2.
+    # See b/211666001 for context.
+    export_graph = ops.Graph()
+    with export_graph.as_default():
+      a = array_ops.placeholder(
+          shape=[None, 1], dtype=dtypes.float32, name="input_2"
+      )
+      b = array_ops.placeholder(
+          shape=[None, 2], dtype=dtypes.float32, name="input_1"
+      )
+      c = array_ops.identity(a)
+      with session_lib.Session() as session:
+        path = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+        simple_save.simple_save(
+            session, path, inputs={"a": a, "b": b}, outputs={"c": c}
+        )
+    imported = load.load(path)
+    path2 = os.path.join(self.get_temp_dir(), "saved_model", str(ops.uid()))
+    save.save(imported, path2, imported.signatures)
+
+    imported2 = load.load(path2)
+    self.assertEqual(
+        imported2.signatures["serving_default"](
+            a=constant_op.constant([5.0]), b=constant_op.constant([1.0, 3.0])
+        )["c"].numpy(),
+        5.0,
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/loader.py

@@ -0,0 +1,65 @@
+# 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.
+# ==============================================================================
+"""Loader functionality for SavedModel with hermetic, language-neutral exports.
+
+Load and restore capability for a SavedModel, which may include multiple meta
+graph defs. Each SavedModel is associated with a single checkpoint. Each meta
+graph def is saved with one or more tags, which are used to identify the exact
+meta graph def to load.
+
+The `load` operation requires the session in which to restore the graph
+definition and variables, the tags used to identify the meta graph def to
+load and the location of the SavedModel.
+
+Upon a load, the subset of variables and assets supplied as part of the specific
+meta graph def, will be restored into the supplied session. The values of the
+variables though will correspond to the saved values from the first meta graph
+added to the SavedModel using `add_meta_graph_and_variables(...)` in
+`builder.py`.
+
+Typical usage:
+
+```python
+...
+builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_dir)
+
+with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+  ...
+  builder.add_meta_graph_and_variables(sess,
+                                       ["foo-tag"],
+                                       signature_def_map=foo_signatures,
+                                       assets_collection=foo_assets)
+...
+
+with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+  ...
+  builder.add_meta_graph(["bar-tag", "baz-tag"],
+                         assets_collection=bar_baz_assets)
+...
+
+builder.save()
+
+...
+with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+  tf.compat.v1.saved_model.loader.load(sess, ["foo-tag"], export_dir)
+  ...
+
+```
+"""
+
+# pylint: disable=unused-import
+from tensorflow.python.saved_model.loader_impl import load
+from tensorflow.python.saved_model.loader_impl import maybe_saved_model_directory
+# pylint: enable=unused-import

tensorflow/python/saved_model/loader_impl.py

@@ -0,0 +1,516 @@
+# 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.
+# ==============================================================================
+"""Loader implementation for SavedModel with hermetic, language-neutral exports.
+"""
+
+import os
+import sys
+
+from google.protobuf import message
+from google.protobuf import text_format
+
+from tensorflow.core.framework import graph_debug_info_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.core.protobuf import saved_model_pb2
+from tensorflow.python.framework import ops
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import path_helpers
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import utils_impl as saved_model_utils
+# Placeholder for protosplitter merger import.
+from tensorflow.python.saved_model.pywrap_saved_model import metrics
+from tensorflow.python.training import saver as tf_saver
+from tensorflow.python.util import compat
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+# API label for SavedModel metrics.
+_LOADER_LABEL = "loader"
+
+
+def parse_saved_model_with_debug_info(export_dir):
+  """Reads the savedmodel as well as the graph debug info.
+
+  Args:
+    export_dir: Directory containing the SavedModel and GraphDebugInfo files.
+
+  Returns:
+    `SavedModel` and `GraphDebugInfo` protocol buffers.
+
+  Raises:
+    IOError: If the saved model file does not exist, or cannot be successfully
+    parsed. Missing graph debug info file is fine.
+  """
+  saved_model = parse_saved_model(export_dir)
+
+  debug_info_path = file_io.join(
+      path_helpers.get_debug_dir(export_dir),
+      constants.DEBUG_INFO_FILENAME_PB)
+  debug_info = graph_debug_info_pb2.GraphDebugInfo()
+  if file_io.file_exists(debug_info_path):
+    with file_io.FileIO(debug_info_path, "rb") as debug_file:
+      try:
+        debug_info.ParseFromString(debug_file.read())
+      except message.DecodeError as e:
+        raise IOError(f"Cannot parse file {debug_info_path}: {e}.")
+
+  return (saved_model, debug_info)
+
+
+@tf_export("__internal__.saved_model.parse_saved_model", v1=[])
+def parse_saved_model(export_dir):
+  """Reads the savedmodel.pb or savedmodel.pbtxt file containing `SavedModel`.
+
+  Args:
+    export_dir: String or Pathlike, path to the directory containing the
+    SavedModel file.
+
+  Returns:
+    A `SavedModel` protocol buffer.
+
+  Raises:
+    IOError: If the file does not exist, or cannot be successfully parsed.
+  """
+  # Build the path to the SavedModel in pbtxt format.
+  path_to_pbtxt = file_io.join(
+      compat.as_bytes(compat.path_to_str(export_dir)),
+      compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
+  # Build the path to the SavedModel in pb format.
+  path_to_pb = file_io.join(
+      compat.as_bytes(compat.path_to_str(export_dir)),
+      compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
+  # Build the path to the SavedModel in cpb format.
+  path_to_cpb = file_io.join(
+      compat.as_bytes(compat.path_to_str(export_dir)),
+      compat.as_bytes(constants.SAVED_MODEL_FILENAME_CPB))
+
+  # Parse the SavedModel protocol buffer.
+  saved_model = saved_model_pb2.SavedModel()
+  if file_io.file_exists(path_to_pb):
+    with file_io.FileIO(path_to_pb, "rb") as f:
+      file_content = f.read()
+    try:
+      saved_model.ParseFromString(file_content)
+    except message.DecodeError as e:
+      raise IOError(f"Cannot parse file {path_to_pb}: {str(e)}.") from e
+  elif file_io.file_exists(path_to_pbtxt):
+    with file_io.FileIO(path_to_pbtxt, "rb") as f:
+      file_content = f.read()
+    try:
+      text_format.Parse(file_content.decode("utf-8"), saved_model)
+    except text_format.ParseError as e:
+      raise IOError(f"Cannot parse file {path_to_pbtxt}: {str(e)}.") from e
+  else:
+    raise IOError(
+        f"SavedModel file does not exist at: {export_dir}{os.path.sep}"
+        f"{{{constants.SAVED_MODEL_FILENAME_PBTXT}|"
+        f"{constants.SAVED_MODEL_FILENAME_PB}}}")
+  return saved_model
+
+
+def get_asset_tensors(export_dir, meta_graph_def_to_load, import_scope=None):
+  """Gets the asset tensors, if defined in the meta graph def to load.
+
+  Args:
+    export_dir: Directory where the SavedModel is located.
+    meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.
+    import_scope: Optional `string` -- if specified, prepend this followed by
+        '/' to all returned asset tensor names.
+
+  Returns:
+    A dictionary of asset tensors, keyed by the name of the asset tensor. The
+    value in the map corresponds to the absolute path of the asset file.
+  """
+  # Collection-def that may contain the assets key.
+  collection_def = meta_graph_def_to_load.collection_def
+
+  asset_tensor_dict = {}
+  asset_protos = []
+
+  if meta_graph_def_to_load.asset_file_def:
+    asset_protos = meta_graph_def_to_load.asset_file_def
+  elif constants.ASSETS_KEY in collection_def:
+    assets_any_proto = collection_def[constants.ASSETS_KEY].any_list.value
+    for asset_any_proto in assets_any_proto:
+      asset_proto = meta_graph_pb2.AssetFileDef()
+      asset_any_proto.Unpack(asset_proto)
+      asset_protos.append(asset_proto)
+
+  # Location of the assets for SavedModel.
+  assets_directory = file_io.join(
+      compat.as_bytes(export_dir), compat.as_bytes(constants.ASSETS_DIRECTORY))
+  # Process each asset and add it to the asset tensor dictionary.
+  for asset_proto in asset_protos:
+    tensor_name = asset_proto.tensor_info.name
+    if import_scope:
+      tensor_name = "%s/%s" % (import_scope, tensor_name)
+    asset_tensor_dict[tensor_name] = file_io.join(
+        compat.as_bytes(assets_directory),
+        compat.as_bytes(asset_proto.filename))
+
+  return asset_tensor_dict
+
+
+def _get_main_op_tensor(
+    meta_graph_def_to_load, init_op_key=constants.MAIN_OP_KEY):
+  """Gets the main op tensor, if one exists.
+
+  Args:
+    meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.
+    init_op_key: name of the collection to check; should be one of MAIN_OP_KEY
+      or the deprecated LEGACY_INIT_OP_KEY
+
+  Returns:
+    The main op tensor, if it exists and `None` otherwise.
+
+  Raises:
+    RuntimeError: If the collection def corresponding to the main op key has
+        other than exactly one tensor.
+  """
+  collection_def = meta_graph_def_to_load.collection_def
+  init_op = None
+  if init_op_key in collection_def:
+    init_op_list = collection_def[init_op_key].node_list.value
+    if len(init_op_list) != 1:
+      raise RuntimeError("Expected exactly one SavedModel init op. "
+                         f"Found {len(init_op_list)}: {init_op_list}.")
+    init_op = ops.get_collection(init_op_key)[0]
+  return init_op
+
+
+def _get_op_from_collection(meta_graph_def, op_key):
+  return _get_main_op_tensor(meta_graph_def, op_key)
+
+
+def _get_op_from_signature_def(meta_graph_def, op_signature_key, import_scope):
+  """Retrieve op stored in the imported meta graph's signature def."""
+  if op_signature_key in meta_graph_def.signature_def:
+    return signature_def_utils.load_op_from_signature_def(
+        meta_graph_def.signature_def[op_signature_key], op_signature_key,
+        import_scope)
+  else:
+    return None
+
+
+def get_init_op(meta_graph_def, import_scope=None):
+  return (_get_op_from_signature_def(
+      meta_graph_def, constants.INIT_OP_SIGNATURE_KEY, import_scope) or
+          _get_op_from_collection(meta_graph_def, constants.MAIN_OP_KEY) or
+          _get_op_from_collection(meta_graph_def, constants.LEGACY_INIT_OP_KEY))
+
+
+def get_train_op(meta_graph_def, import_scope=None):
+  train_op = _get_op_from_signature_def(
+      meta_graph_def, constants.TRAIN_OP_SIGNATURE_KEY, import_scope)
+  if train_op is None:
+    train_op = _get_op_from_collection(meta_graph_def, constants.TRAIN_OP_KEY)
+  return train_op
+
+
+@tf_export(v1=[
+    "saved_model.contains_saved_model",
+    "saved_model.maybe_saved_model_directory",
+    "saved_model.loader.maybe_saved_model_directory"
+])
+@deprecation.deprecated_endpoints(
+    "saved_model.loader.maybe_saved_model_directory")
+def maybe_saved_model_directory(export_dir):
+  """Checks whether the provided export directory could contain a SavedModel.
+
+  Note that the method does not load any data by itself. If the method returns
+  `false`, the export directory definitely does not contain a SavedModel. If the
+  method returns `true`, the export directory may contain a SavedModel but
+  provides no guarantee that it can be loaded.
+
+  Args:
+    export_dir: Absolute string path to possible export location. For example,
+                '/my/foo/model'.
+
+  Returns:
+    True if the export directory contains SavedModel files, False otherwise.
+  """
+  txt_path = file_io.join(export_dir, constants.SAVED_MODEL_FILENAME_PBTXT)
+  pb_path = file_io.join(export_dir, constants.SAVED_MODEL_FILENAME_PB)
+  cpb_path = file_io.join(export_dir, constants.SAVED_MODEL_FILENAME_CPB)
+  return (
+      file_io.file_exists(txt_path)
+      or file_io.file_exists(pb_path)
+      or file_io.file_exists(cpb_path)
+  )
+
+
+@tf_export("saved_model.contains_saved_model", v1=[])
+def contains_saved_model(export_dir):
+  """Checks whether the provided export directory could contain a SavedModel.
+
+  Note that the method does not load any data by itself. If the method returns
+  `false`, the export directory definitely does not contain a SavedModel. If the
+  method returns `true`, the export directory may contain a SavedModel but
+  provides no guarantee that it can be loaded.
+
+  Args:
+    export_dir: Absolute path to possible export location. For example,
+                '/my/foo/model'.
+
+  Returns:
+    True if the export directory contains SavedModel files, False otherwise.
+  """
+  if isinstance(export_dir, os.PathLike):
+    export_dir = os.fspath(export_dir)
+  return maybe_saved_model_directory(export_dir)
+
+
+@tf_export(v1=["saved_model.load", "saved_model.loader.load"])
+@deprecation.deprecated(
+    None,
+    "Use `tf.saved_model.load` instead.")
+def load(sess, tags, export_dir, import_scope=None, **saver_kwargs):
+  """Loads the model from a SavedModel as specified by tags.
+
+  Args:
+    sess: The TensorFlow session to restore the variables.
+    tags: Set of string tags to identify the required MetaGraphDef. These should
+        correspond to the tags used when saving the variables using the
+        SavedModel `save()` API.
+    export_dir: Directory in which the SavedModel protocol buffer and variables
+        to be loaded are located.
+    import_scope: Optional `string` -- if specified, prepend this string
+        followed by '/' to all loaded tensor names. This scope is applied to
+        tensor instances loaded into the passed session, but it is *not* written
+        through to the static `MetaGraphDef` protocol buffer that is returned.
+    **saver_kwargs: Optional keyword arguments passed through to Saver.
+
+  Returns:
+    The `MetaGraphDef` protocol buffer loaded in the provided session. This
+    can be used to further extract signature-defs, collection-defs, etc.
+
+  Raises:
+    RuntimeError: MetaGraphDef associated with the tags cannot be found.
+
+  @compatibility(TF2)
+
+  `tf.compat.v1.saved_model.load` or `tf.compat.v1.saved_model.loader.load` is
+  not compatible with eager execution. Please use `tf.saved_model.load` instead
+  to load your model. You can refer to the [SavedModel guide]
+  (https://www.tensorflow.org/guide/saved_model) for more information as well as
+  "Importing SavedModels from TensorFlow 1.x" in the [`tf.saved_model.load`]
+  (https://www.tensorflow.org/api_docs/python/tf/saved_model/load) docstring.
+
+  #### How to Map Arguments
+
+  | TF1 Arg Name          | TF2 Arg Name    | Note                       |
+  | :-------------------- | :-------------- | :------------------------- |
+  | `sess`                | Not supported   | -                          |
+  | `tags`                | `tags`          | -                          |
+  | `export_dir`          | `export_dir`    | -                          |
+  | `import_scope`        | Not supported   | Name scopes are not needed.
+  :                       :                 : By default, variables are  :
+  :                       :                 : associated with the loaded :
+  :                       :                 : object and function names  :
+  :                       :                 : are deduped.               :
+  | `saver_kwargs`        | Not supported   | -                          |
+
+  #### Before & After Usage Example
+
+  Before:
+
+  ```
+  with tf.compat.v1.Session(graph=tf.Graph()) as sess:
+    tf.compat.v1.saved_model.loader.load(sess, ["foo-tag"], export_dir)
+  ```
+
+  After:
+
+  ```
+  model = tf.saved_model.load(export_dir, tags=["foo-tag"])
+  ```
+  @end_compatibility
+  """
+  loader = SavedModelLoader(export_dir)
+  return loader.load(sess, tags, import_scope, **saver_kwargs)
+
+
+class SavedModelLoader(object):
+  """Load graphs and restore variable values from a `SavedModel`."""
+
+  def __init__(self, export_dir):
+    """Creates a `SavedModelLoader`.
+
+    Args:
+      export_dir: Directory in which the SavedModel protocol buffer and
+        variables to be loaded are located.
+    """
+    self._export_dir = export_dir
+    self._variables_path = path_helpers.get_variables_path(export_dir)
+    self._saved_model = parse_saved_model(export_dir)
+
+  @property
+  def export_dir(self):
+    """Directory containing the SavedModel."""
+    return self._export_dir
+
+  @property
+  def variables_path(self):
+    """Path to variable checkpoint files."""
+    return self._variables_path
+
+  @property
+  def saved_model(self):
+    """SavedModel object parsed from the export directory."""
+    return self._saved_model
+
+  def get_meta_graph_def_from_tags(self, tags):
+    """Return MetaGraphDef with the exact specified tags.
+
+    Args:
+      tags: A list or set of string tags that identify the MetaGraphDef.
+
+    Returns:
+      MetaGraphDef with the same tags.
+
+    Raises:
+      RuntimeError: if no metagraphs were found with the associated tags.
+    """
+    found_match = False
+    meta_graph_def_to_load = None
+    available_tags = []
+    for meta_graph_def in self._saved_model.meta_graphs:
+      available_tags.append(set(meta_graph_def.meta_info_def.tags))
+      if set(meta_graph_def.meta_info_def.tags) == set(tags):
+        meta_graph_def_to_load = meta_graph_def
+        found_match = True
+        break
+
+    if not found_match:
+      raise RuntimeError(
+          f"MetaGraphDef associated with tags {str(tags).strip('[]')} "
+          "could not be found in SavedModel, with available tags "
+          f"'{available_tags}'. To inspect available tag-sets in"
+          " the SavedModel, please use the SavedModel CLI: `saved_model_cli`.")
+    return meta_graph_def_to_load
+
+  def load_graph(self, graph, tags, import_scope=None, **saver_kwargs):
+    """Load ops and nodes from SavedModel MetaGraph into graph.
+
+    Args:
+      graph: tf.Graph object.
+      tags: a set of string tags identifying a MetaGraphDef.
+      import_scope: Optional `string` -- if specified, prepend this string
+        followed by '/' to all loaded tensor names. This scope is applied to
+        tensor instances loaded into the passed session, but it is *not* written
+        through to the static `MetaGraphDef` protocol buffer that is returned.
+      **saver_kwargs: keyword arguments to pass to tf.train.import_meta_graph.
+
+    Returns:
+      A tuple of
+        * Saver defined by the MetaGraph, which can be used to restore the
+          variable values.
+        * List of `Operation`/`Tensor` objects returned from
+          `tf.import_graph_def` (may be `None`).
+    """
+    meta_graph_def = self.get_meta_graph_def_from_tags(tags)
+    if sys.byteorder == "big":
+      saved_model_utils.swap_function_tensor_content(meta_graph_def, "little",
+                                                     "big")
+    with graph.as_default():
+      return tf_saver._import_meta_graph_with_return_elements(  # pylint: disable=protected-access
+          meta_graph_def, import_scope=import_scope, **saver_kwargs)
+
+  def restore_variables(self, sess, saver, import_scope=None):
+    """Restore SavedModel variable values into the session.
+
+    Args:
+      sess: tf.compat.v1.Session to restore variable values.
+      saver: a tf.compat.v1.train.Saver object. Can be None if there are no
+        variables in graph. This may be the saver returned by the load_graph()
+        function, or a default `tf.compat.v1.train.Saver()`.
+      import_scope: Optional `string` -- if specified, prepend this string
+        followed by '/' to all loaded tensor names. This scope is applied to
+        tensor instances loaded into the passed session, but it is *not* written
+        through to the static `MetaGraphDef` protocol buffer that is returned.
+
+    Raises:
+      ValueError: if no saver was passed to the saver argument, and there are
+        variables in the graph.
+    """
+    with sess.graph.as_default():
+      if (saver is None and
+          not variables._all_saveable_objects(scope=import_scope)):  # pylint: disable=protected-access
+        tf_logging.info("The specified SavedModel has no variables; no "
+                        "checkpoints were restored.")
+      elif isinstance(saver, tf_saver.Saver):
+        saver.restore(sess, self._variables_path)
+      else:
+        raise ValueError(
+            "No tf.train.Saver object was passed to the function "
+            "`SavedModelLoader.restore_variables`. Since there are variables in"
+            " the graph, a saver is required.")
+
+  def run_init_ops(self, sess, tags, import_scope=None):
+    """Run initialization ops defined in the `MetaGraphDef`.
+
+    Args:
+      sess: tf.compat.v1.Session to restore variable values.
+      tags: a set of string tags identifying a MetaGraphDef.
+      import_scope: Optional `string` -- if specified, prepend this string
+        followed by '/' to all loaded tensor names. This scope is applied to
+        tensor instances loaded into the passed session, but it is *not* written
+        through to the static `MetaGraphDef` protocol buffer that is returned.
+    """
+    meta_graph_def = self.get_meta_graph_def_from_tags(tags)
+    with sess.graph.as_default():
+      # Get asset tensors, if any.
+      asset_tensors_dictionary = get_asset_tensors(
+          self._export_dir, meta_graph_def, import_scope=import_scope)
+
+      init_op = get_init_op(meta_graph_def, import_scope)
+      if init_op is not None:
+        sess.run(fetches=[init_op], feed_dict=asset_tensors_dictionary)
+
+  def load(self, sess, tags, import_scope=None, **saver_kwargs):
+    """Load the MetaGraphDef graph and restore variable values into the session.
+
+    Args:
+      sess: tf.compat.v1.Session to restore variable values.
+      tags: a set of string tags identifying a MetaGraphDef.
+      import_scope: Optional `string` -- if specified, prepend this string
+        followed by '/' to all loaded tensor names. This scope is applied to
+        tensor instances loaded into the passed session, but it is *not* written
+        through to the static `MetaGraphDef` protocol buffer that is returned.
+      **saver_kwargs: keyword arguments to pass to tf.train.import_meta_graph.
+
+    Returns:
+      `MetagraphDef` proto of the graph that was loaded.
+    """
+    saved_model_proto = parse_saved_model(self._export_dir)
+    metrics.IncrementReadApi(_LOADER_LABEL)
+
+    with sess.graph.as_default():
+      saver, _ = self.load_graph(sess.graph, tags, import_scope,
+                                 **saver_kwargs)
+      self.restore_variables(sess, saver, import_scope)
+      self.run_init_ops(sess, tags, import_scope)
+    meta_graph_def = self.get_meta_graph_def_from_tags(tags)
+
+    if (len(saved_model_proto.meta_graphs) == 1 and
+        saved_model_proto.meta_graphs[0].HasField("object_graph_def")):
+      metrics.IncrementRead(write_version="2")
+    else:
+      metrics.IncrementRead(write_version="1")
+
+    return meta_graph_def

tensorflow/python/saved_model/loader_test.py

@@ -0,0 +1,301 @@
+# 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.
+# ==============================================================================
+"""Tests for SavedModelLoader class."""
+
+import os
+import shutil
+
+from absl.testing import parameterized
+
+from tensorflow.python.client import session
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import resource_variables_toggle
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_v1
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import builder as saved_model_builder
+from tensorflow.python.saved_model import loader_impl
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import utils
+from tensorflow.python.training import saver as tf_saver
+
+
+def _get_export_dir(label):
+  return os.path.join(test.get_temp_dir(), label)
+
+
+def _tensor_name(name):
+  if resource_variables_toggle.resource_variables_enabled():
+    return name + "/Read/ReadVariableOp:0"
+  return name + ":0"
+
+
+SIMPLE_ADD_SAVED_MODEL = _get_export_dir("simple_add_saved_model")
+SAVED_MODEL_WITH_MAIN_OP = _get_export_dir("saved_model_with_main_op")
+
+
+def build_graph_helper():
+  g = ops.Graph()
+  with g.as_default():
+    x = variable_v1.VariableV1(5, name="x")
+    y = variable_v1.VariableV1(11, name="y")
+    z = x + y
+
+    foo_sig_def = signature_def_utils.build_signature_def({
+        "foo_input": utils.build_tensor_info(x)
+    }, {"foo_output": utils.build_tensor_info(z)})
+    bar_sig_def = signature_def_utils.build_signature_def({
+        "bar_x": utils.build_tensor_info(x),
+        "bar_y": utils.build_tensor_info(y)
+    }, {"bar_z": utils.build_tensor_info(z)})
+  return g, {"foo": foo_sig_def, "bar": bar_sig_def}, y
+
+
+@parameterized.parameters((saved_model_builder.SavedModelBuilder,),
+                          (saved_model_builder._SavedModelBuilder,))
+class SavedModelLoaderTest(test.TestCase, parameterized.TestCase):
+
+  def export_simple_graph(self, builder_cls):
+    g, sig_def_map, _ = build_graph_helper()
+    with session.Session(graph=g) as sess:
+      self.evaluate(variables.global_variables_initializer())
+      builder = builder_cls(SIMPLE_ADD_SAVED_MODEL)
+      builder.add_meta_graph_and_variables(sess, ["foo_graph"], sig_def_map)
+      builder.save()
+
+  def export_graph_with_main_op(self, builder_cls):
+    g, sig_def_map, y = build_graph_helper()
+    with session.Session(graph=g) as sess:
+      self.evaluate(variables.global_variables_initializer())
+      assign_op = control_flow_ops.group(state_ops.assign(y, 7))
+
+      builder = builder_cls(SAVED_MODEL_WITH_MAIN_OP)
+
+      if builder_cls == saved_model_builder._SavedModelBuilder:
+        builder.add_meta_graph_and_variables(
+            sess, ["foo_graph"], sig_def_map, init_op=assign_op)
+      else:
+        builder.add_meta_graph_and_variables(
+            sess, ["foo_graph"], sig_def_map, main_op=assign_op)
+      builder.save()
+
+  def tearDown(self):
+    super(SavedModelLoaderTest, self).tearDown()
+    shutil.rmtree(test.get_temp_dir(), ignore_errors=True)
+
+  def test_load_function(self, builder_cls):
+    # Force test to run in graph mode.
+    # The SavedModelLoader.load method is a v1-only API that requires a session
+    # to work.
+    with ops.Graph().as_default():
+      self.export_simple_graph(builder_cls)
+      loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
+      with self.session(graph=ops.Graph()) as sess:
+        loader.load(sess, ["foo_graph"])
+        self.assertEqual(5, sess.run(_tensor_name("x")))
+        self.assertEqual(11, sess.run(_tensor_name("y")))
+
+      self.export_graph_with_main_op(builder_cls)
+      loader2 = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
+      with self.session(graph=ops.Graph()) as sess:
+        loader2.load(sess, ["foo_graph"])
+        self.assertEqual(5, sess.run(_tensor_name("x")))
+        self.assertEqual(7, sess.run(_tensor_name("y")))
+
+  def test_load_graph(self, builder_cls):
+    self.export_simple_graph(builder_cls)
+    loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
+    graph = ops.Graph()
+    loader.load_graph(graph, ["foo_graph"])
+
+    x = graph.get_tensor_by_name(_tensor_name("x"))
+    y = graph.get_tensor_by_name(_tensor_name("y"))
+
+    with self.assertRaises(KeyError):
+      graph.get_tensor_by_name(_tensor_name("z"))
+
+    with graph.as_default(), self.session():
+      # Check that x and y are not initialized
+      with self.assertRaises(errors.FailedPreconditionError):
+        self.evaluate(x)
+      with self.assertRaises(errors.FailedPreconditionError):
+        self.evaluate(y)
+
+  def test_load_with_import_scope(self, builder_cls):
+    # Force test to run in graph mode.
+    # The SavedModelLoader.restore_variables and SavedModelLoader.run_init_ops
+    # methods are v1-only APIs that require a session to work.
+    with ops.Graph().as_default():
+      self.export_graph_with_main_op(builder_cls)
+      loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
+      with self.session(graph=ops.Graph()) as sess:
+        saver, _ = loader.load_graph(
+            sess.graph, ["foo_graph"], import_scope="baz")
+
+        # The default saver should not work when the import scope is set.
+        with self.assertRaises(errors.NotFoundError):
+          loader.restore_variables(sess, tf_saver.Saver())
+
+        loader.restore_variables(sess, saver)
+
+        if builder_cls == saved_model_builder._SavedModelBuilder:
+          with self.assertRaises(errors.NotFoundError):
+            loader.run_init_ops(sess, ["foo_graph"])
+          loader.run_init_ops(sess, ["foo_graph"], import_scope="baz")
+        else:
+          loader.run_init_ops(sess, ["foo_graph"])
+
+        self.assertEqual(5, sess.run(_tensor_name("baz/x")))
+        self.assertEqual(7, sess.run(_tensor_name("baz/y")))
+
+      # Test combined load function.
+      loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
+      with self.session(graph=ops.Graph()) as sess:
+        loader.load(sess, ["foo_graph"], import_scope="baa")
+        self.assertEqual(5, sess.run(_tensor_name("baa/x")))
+        self.assertEqual(7, sess.run(_tensor_name("baa/y")))
+
+  def test_restore_variables(self, builder_cls):
+    # Force test to run in graph mode.
+    # The SavedModelLoader.restore_variables method is a v1-only API requiring a
+    # session to work.
+    with ops.Graph().as_default():
+      self.export_graph_with_main_op(builder_cls)
+      loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
+      with self.session() as sess:
+        x = variable_v1.VariableV1(0, name="x")
+        y = variable_v1.VariableV1(0, name="y")
+        z = x * y
+
+        self.evaluate(variables.global_variables_initializer())
+
+        # There are variables to restore, so a saver must be created.
+        with self.assertRaises(ValueError):
+          loader.restore_variables(sess, None)
+
+        loader.restore_variables(sess, tf_saver.Saver())
+        self.assertEqual(55, self.evaluate(z))
+
+  def test_run_init_op(self, builder_cls):
+    # Force test to run in graph mode.
+    # The SavedModelLoader.restore_variables and SavedModelLoader.run_init_ops
+    # methods are v1-only APIs that require a session to work.
+    with ops.Graph().as_default():
+      self.export_graph_with_main_op(builder_cls)
+      loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
+      graph = ops.Graph()
+      saver, _ = loader.load_graph(graph, ["foo_graph"])
+      with self.session(graph=graph) as sess:
+        loader.restore_variables(sess, saver)
+        self.assertEqual(5, sess.run(_tensor_name("x")))
+        self.assertEqual(11, sess.run(_tensor_name("y")))
+
+        loader.run_init_ops(sess, ["foo_graph"])
+        self.assertEqual(5, sess.run(_tensor_name("x")))
+        self.assertEqual(7, sess.run(_tensor_name("y")))
+
+  def test_parse_saved_model(self, builder_cls):
+    self.export_simple_graph(builder_cls)
+    loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
+    meta_graph = loader.get_meta_graph_def_from_tags(["foo_graph"])
+    self.assertIsNotNone(meta_graph)
+    self.assertIn("foo", meta_graph.signature_def)
+    self.assertIn("bar", meta_graph.signature_def)
+
+  def test_load_invalid_meta_graph(self, builder_cls):
+    self.export_simple_graph(builder_cls)
+    loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
+    with self.assertRaises(RuntimeError):
+      loader.get_meta_graph_def_from_tags([])
+    with self.assertRaises(RuntimeError):
+      loader.get_meta_graph_def_from_tags([""])
+    with self.assertRaises(RuntimeError):
+      loader.get_meta_graph_def_from_tags(["not_a_graph"])
+
+  def test_load_saved_model_with_no_variables(self, builder_cls):
+    """Test that SavedModel runs saver when there appear to be no variables.
+
+    When no variables are detected, this may mean that the variables were saved
+    to different collections, or the collections weren't saved to the
+    SavedModel. If the SavedModel MetaGraphDef contains a saver, it should still
+    run in either of these cases.
+
+    Args:
+      builder_cls: SavedModelBuilder or _SavedModelBuilder class
+    """
+    # Force test to run in graph mode.
+    # The SavedModelBuilder.add_meta_graph_and_variables and
+    # SavedModelLoader.load methods are v1-only APIs that require a session to
+    # work.
+    with ops.Graph().as_default():
+      path = _get_export_dir("no_variable_saved_model")
+      with session.Session(graph=ops.Graph()) as sess:
+        x = variable_v1.VariableV1(
+            5, name="x", collections=["not_global_variable"])
+        y = variable_v1.VariableV1(
+            11, name="y", collections=["not_global_variable"])
+        self.assertFalse(variables._all_saveable_objects())
+        z = x + y
+        self.evaluate(variables.variables_initializer([x, y]))
+
+        foo_sig_def = signature_def_utils.build_signature_def(
+            {"foo_input": utils.build_tensor_info(x)},
+            {"foo_output": utils.build_tensor_info(z)})
+
+        builder = saved_model_builder.SavedModelBuilder(path)
+        builder.add_meta_graph_and_variables(
+            sess, ["foo_graph"], {"foo": foo_sig_def},
+            saver=tf_saver.Saver([x, y]))
+        builder.save()
+
+      loader = loader_impl.SavedModelLoader(path)
+      with self.session(graph=ops.Graph()) as sess:
+        saver, _ = loader.load_graph(sess.graph, ["foo_graph"])
+        self.assertFalse(variables._all_saveable_objects())
+        self.assertIsNotNone(saver)
+
+      with self.session(graph=ops.Graph()) as sess:
+        loader.load(sess, ["foo_graph"])
+        self.assertEqual(5, sess.run(_tensor_name("x")))
+        self.assertEqual(11, sess.run(_tensor_name("y")))
+
+  def test_load_saved_model_graph_with_return_elements(self, builder_cls):
+    """Ensure that the correct elements are returned."""
+    self.export_simple_graph(builder_cls)
+    loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
+    graph = ops.Graph()
+    _, ret = loader.load_graph(graph, ["foo_graph"],
+                               return_elements=["y:0", "x:0"])
+
+    self.assertEqual(graph.get_tensor_by_name("y:0"), ret[0])
+    self.assertEqual(graph.get_tensor_by_name("x:0"), ret[1])
+
+    with self.assertRaisesRegex(ValueError, "not found in graph"):
+      loader.load_graph(graph, ["foo_graph"], return_elements=["z:0"])
+
+  def test_parse_saved_model_exception(self, builder_cls):
+    """Test that error message for not exist model have OS-depend delimiter in path"""
+    path = _get_export_dir("not_existing_dir")
+    pattern = os.path.sep + "{"
+    with self.assertRaises(IOError) as err:
+      loader_impl.parse_saved_model(path)
+    self.assertTrue(pattern in str(err.exception))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/main_op.py

@@ -0,0 +1,24 @@
+# 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.
+# ==============================================================================
+"""SavedModel main op.
+
+Builds a main op that defines the sequence of ops to be run as part of the
+SavedModel load/restore operations.
+"""
+
+# pylint: disable=unused-import
+from tensorflow.python.saved_model.main_op_impl import main_op
+from tensorflow.python.saved_model.main_op_impl import main_op_with_restore
+# pylint: enable=unused-import

tensorflow/python/saved_model/main_op_impl.py

@@ -0,0 +1,66 @@
+# 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.
+# ==============================================================================
+"""SavedModel main op implementation."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+_DEPRECATION_MSG = (
+    'This API was designed for TensorFlow v1. See '
+    'https://www.tensorflow.org/guide/migrate for instructions on how to '
+    'migrate your code to TensorFlow v2.')
+
+
+@tf_export(v1=['saved_model.main_op.main_op'])
+@deprecation.deprecated(None, _DEPRECATION_MSG)
+def main_op():
+  """Returns a main op to init variables and tables.
+
+  Returns the main op including the group of ops that initializes all
+  variables, initializes local variables and initialize all tables.
+
+  Returns:
+    The set of ops to be run as part of the main op upon the load operation.
+  """
+  init = variables.global_variables_initializer()
+  init_local = variables.local_variables_initializer()
+  init_tables = lookup_ops.tables_initializer()
+  return control_flow_ops.group(init, init_local, init_tables)
+
+
+# TODO(sukritiramesh): Integrate with Saver for complete restore functionality.
+@tf_export(v1=['saved_model.main_op_with_restore',
+               'saved_model.main_op.main_op_with_restore'])
+@deprecation.deprecated(None, _DEPRECATION_MSG)
+def main_op_with_restore(restore_op_name):
+  """Returns a main op to init variables, tables and restore the graph.
+
+  Returns the main op including the group of ops that initializes all
+  variables, initialize local variables, initialize all tables and the restore
+  op name.
+
+  Args:
+    restore_op_name: Name of the op to use to restore the graph.
+
+  Returns:
+    The set of ops to be run as part of the main op upon the load operation.
+  """
+  with ops.control_dependencies([main_op()]):
+    main_op_with_restore = control_flow_ops.group(restore_op_name)
+  return main_op_with_restore

tensorflow/python/saved_model/method_name_updater.py

@@ -0,0 +1,143 @@
+# 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.
+# ==============================================================================
+"""SignatureDef method name utility functions.
+
+Utility functions for manipulating signature_def.method_names.
+"""
+
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.platform import tf_logging
+from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import loader_impl as loader
+from tensorflow.python.util import compat
+from tensorflow.python.util.tf_export import tf_export
+
+
+# TODO(jdchung): Consider integrated this into the saved_model_cli so that users
+# could do this from the command line directly.
+@tf_export(v1=["saved_model.signature_def_utils.MethodNameUpdater"])
+class MethodNameUpdater(object):
+  """Updates the method name(s) of the SavedModel stored in the given path.
+
+  The `MethodNameUpdater` class provides the functionality to update the method
+  name field in the signature_defs of the given SavedModel. For example, it
+  can be used to replace the `predict` `method_name` to `regress`.
+
+  Typical usages of the `MethodNameUpdater`
+  ```python
+  ...
+  updater = tf.compat.v1.saved_model.signature_def_utils.MethodNameUpdater(
+      export_dir)
+  # Update all signature_defs with key "foo" in all meta graph defs.
+  updater.replace_method_name(signature_key="foo", method_name="regress")
+  # Update a single signature_def with key "bar" in the meta graph def with
+  # tags ["serve"]
+  updater.replace_method_name(signature_key="bar", method_name="classify",
+                              tags="serve")
+  updater.save(new_export_dir)
+  ```
+
+  Note: This function will only be available through the v1 compatibility
+  library as tf.compat.v1.saved_model.builder.MethodNameUpdater.
+  """
+
+  def __init__(self, export_dir):
+    """Creates an MethodNameUpdater object.
+
+    Args:
+      export_dir: Directory containing the SavedModel files.
+
+    Raises:
+      IOError: If the saved model file does not exist, or cannot be successfully
+      parsed.
+    """
+    self._export_dir = export_dir
+    self._saved_model = loader.parse_saved_model(export_dir)
+
+  def replace_method_name(self, signature_key, method_name, tags=None):
+    """Replaces the method_name in the specified signature_def.
+
+    This will match and replace multiple sig defs iff tags is None (i.e when
+    multiple `MetaGraph`s have a signature_def with the same key).
+    If tags is not None, this will only replace a single signature_def in the
+    `MetaGraph` with matching tags.
+
+    Args:
+      signature_key: Key of the signature_def to be updated.
+      method_name: new method_name to replace the existing one.
+      tags: A tag or sequence of tags identifying the `MetaGraph` to update. If
+          None, all meta graphs will be updated.
+    Raises:
+      ValueError: if signature_key or method_name are not defined or
+          if no metagraphs were found with the associated tags or
+          if no meta graph has a signature_def that matches signature_key.
+    """
+    if not signature_key:
+      raise ValueError("`signature_key` must be defined.")
+    if not method_name:
+      raise ValueError("`method_name` must be defined.")
+
+    if (tags is not None and not isinstance(tags, list)):
+      tags = [tags]
+    found_match = False
+    for meta_graph_def in self._saved_model.meta_graphs:
+      if tags is None or set(tags) == set(meta_graph_def.meta_info_def.tags):
+        if signature_key not in meta_graph_def.signature_def:
+          raise ValueError(
+              f"MetaGraphDef associated with tags {tags} "
+              f"does not have a signature_def with key: '{signature_key}'. "
+              "This means either you specified the wrong signature key or "
+              "forgot to put the signature_def with the corresponding key in "
+              "your SavedModel.")
+        meta_graph_def.signature_def[signature_key].method_name = method_name
+        found_match = True
+
+    if not found_match:
+      raise ValueError(
+          f"MetaGraphDef associated with tags {tags} could not be found in "
+          "SavedModel. This means either you specified invalid tags or your "
+          "SavedModel does not have a MetaGraphDef with the specified tags.")
+
+  def save(self, new_export_dir=None):
+    """Saves the updated `SavedModel`.
+
+    Args:
+      new_export_dir: Path where the updated `SavedModel` will be saved. If
+          None, the input `SavedModel` will be overriden with the updates.
+
+    Raises:
+      errors.OpError: If there are errors during the file save operation.
+    """
+
+    is_input_text_proto = file_io.file_exists(
+        file_io.join(
+            compat.as_bytes(self._export_dir),
+            compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT)))
+    if not new_export_dir:
+      new_export_dir = self._export_dir
+
+    if is_input_text_proto:
+      # TODO(jdchung): Add a util for the path creation below.
+      path = file_io.join(
+          compat.as_bytes(new_export_dir),
+          compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
+      file_io.write_string_to_file(path, str(self._saved_model))
+    else:
+      path = file_io.join(
+          compat.as_bytes(new_export_dir),
+          compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
+      file_io.write_string_to_file(
+          path, self._saved_model.SerializeToString(deterministic=True))
+    tf_logging.info("SavedModel written to: %s", compat.as_text(path))

tensorflow/python/saved_model/method_name_updater_test.py

@@ -0,0 +1,373 @@
+# 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.
+# ==============================================================================
+"""Tests for method name utils."""
+
+import os
+import tempfile
+
+from google.protobuf import text_format
+from tensorflow.core.protobuf import saved_model_pb2
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import loader_impl as loader
+from tensorflow.python.saved_model import method_name_updater
+from tensorflow.python.util import compat
+
+_SAVED_MODEL_PROTO = text_format.Parse("""
+saved_model_schema_version: 1
+meta_graphs {
+  meta_info_def {
+    tags: "serve"
+  }
+  signature_def: {
+    key: "serving_default"
+    value: {
+      inputs: {
+        key: "inputs"
+        value { name: "input_node:0" }
+      }
+      method_name: "predict"
+      outputs: {
+        key: "outputs"
+        value {
+          dtype: DT_FLOAT
+          tensor_shape {
+            dim { size: -1 }
+            dim { size: 100 }
+          }
+        }
+      }
+    }
+  }
+  signature_def: {
+    key: "foo"
+    value: {
+      inputs: {
+        key: "inputs"
+        value { name: "input_node:0" }
+      }
+      method_name: "predict"
+      outputs: {
+        key: "outputs"
+        value {
+          dtype: DT_FLOAT
+          tensor_shape { dim { size: 1 } }
+        }
+      }
+    }
+  }
+}
+meta_graphs {
+  meta_info_def {
+    tags: "serve"
+    tags: "gpu"
+  }
+  signature_def: {
+    key: "serving_default"
+    value: {
+      inputs: {
+        key: "inputs"
+        value { name: "input_node:0" }
+      }
+      method_name: "predict"
+      outputs: {
+        key: "outputs"
+        value {
+          dtype: DT_FLOAT
+          tensor_shape {
+            dim { size: -1 }
+          }
+        }
+      }
+    }
+  }
+  signature_def: {
+    key: "bar"
+    value: {
+      inputs: {
+        key: "inputs"
+        value { name: "input_node:0" }
+      }
+      method_name: "predict"
+      outputs: {
+        key: "outputs"
+        value {
+          dtype: DT_FLOAT
+          tensor_shape { dim { size: 1 } }
+        }
+      }
+    }
+  }
+}
+""", saved_model_pb2.SavedModel())
+
+
+class MethodNameUpdaterTest(test.TestCase):
+
+  def setUp(self):
+    super(MethodNameUpdaterTest, self).setUp()
+    self._saved_model_path = tempfile.mkdtemp(prefix=test.get_temp_dir())
+
+  def testBasic(self):
+    path = os.path.join(
+        compat.as_bytes(self._saved_model_path),
+        compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
+    file_io.write_string_to_file(
+        path, _SAVED_MODEL_PROTO.SerializeToString(deterministic=True))
+
+    updater = method_name_updater.MethodNameUpdater(self._saved_model_path)
+    updater.replace_method_name(
+        signature_key="serving_default", method_name="classify")
+    updater.save()
+
+    actual = loader.parse_saved_model(self._saved_model_path)
+    self.assertProtoEquals(
+        actual,
+        text_format.Parse(
+            """
+        saved_model_schema_version: 1
+        meta_graphs {
+          meta_info_def {
+            tags: "serve"
+          }
+          signature_def: {
+            key: "serving_default"
+            value: {
+              inputs: {
+                key: "inputs"
+                value { name: "input_node:0" }
+              }
+              method_name: "classify"
+              outputs: {
+                key: "outputs"
+                value {
+                  dtype: DT_FLOAT
+                  tensor_shape {
+                    dim { size: -1 }
+                    dim { size: 100 }
+                  }
+                }
+              }
+            }
+          }
+          signature_def: {
+            key: "foo"
+            value: {
+              inputs: {
+                key: "inputs"
+                value { name: "input_node:0" }
+              }
+              method_name: "predict"
+              outputs: {
+                key: "outputs"
+                value {
+                  dtype: DT_FLOAT
+                  tensor_shape { dim { size: 1 } }
+                }
+              }
+            }
+          }
+        }
+        meta_graphs {
+          meta_info_def {
+            tags: "serve"
+            tags: "gpu"
+          }
+          signature_def: {
+            key: "serving_default"
+            value: {
+              inputs: {
+                key: "inputs"
+                value { name: "input_node:0" }
+              }
+              method_name: "classify"
+              outputs: {
+                key: "outputs"
+                value {
+                  dtype: DT_FLOAT
+                  tensor_shape {
+                    dim { size: -1 }
+                  }
+                }
+              }
+            }
+          }
+          signature_def: {
+            key: "bar"
+            value: {
+              inputs: {
+                key: "inputs"
+                value { name: "input_node:0" }
+              }
+              method_name: "predict"
+              outputs: {
+                key: "outputs"
+                value {
+                  dtype: DT_FLOAT
+                  tensor_shape { dim { size: 1 } }
+                }
+              }
+            }
+          }
+        }
+    """, saved_model_pb2.SavedModel()))
+
+  def testTextFormatAndNewExportDir(self):
+    path = os.path.join(
+        compat.as_bytes(self._saved_model_path),
+        compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
+    file_io.write_string_to_file(path, str(_SAVED_MODEL_PROTO))
+
+    updater = method_name_updater.MethodNameUpdater(self._saved_model_path)
+    updater.replace_method_name(
+        signature_key="foo", method_name="regress", tags="serve")
+    updater.replace_method_name(
+        signature_key="bar", method_name="classify", tags=["gpu", "serve"])
+
+    new_export_dir = tempfile.mkdtemp(prefix=test.get_temp_dir())
+    updater.save(new_export_dir)
+
+    self.assertTrue(
+        file_io.file_exists(
+            os.path.join(
+                compat.as_bytes(new_export_dir),
+                compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))))
+    actual = loader.parse_saved_model(new_export_dir)
+    self.assertProtoEquals(
+        actual,
+        text_format.Parse(
+            """
+        saved_model_schema_version: 1
+        meta_graphs {
+          meta_info_def {
+            tags: "serve"
+          }
+          signature_def: {
+            key: "serving_default"
+            value: {
+              inputs: {
+                key: "inputs"
+                value { name: "input_node:0" }
+              }
+              method_name: "predict"
+              outputs: {
+                key: "outputs"
+                value {
+                  dtype: DT_FLOAT
+                  tensor_shape {
+                    dim { size: -1 }
+                    dim { size: 100 }
+                  }
+                }
+              }
+            }
+          }
+          signature_def: {
+            key: "foo"
+            value: {
+              inputs: {
+                key: "inputs"
+                value { name: "input_node:0" }
+              }
+              method_name: "regress"
+              outputs: {
+                key: "outputs"
+                value {
+                  dtype: DT_FLOAT
+                  tensor_shape { dim { size: 1 } }
+                }
+              }
+            }
+          }
+        }
+        meta_graphs {
+          meta_info_def {
+            tags: "serve"
+            tags: "gpu"
+          }
+          signature_def: {
+            key: "serving_default"
+            value: {
+              inputs: {
+                key: "inputs"
+                value { name: "input_node:0" }
+              }
+              method_name: "predict"
+              outputs: {
+                key: "outputs"
+                value {
+                  dtype: DT_FLOAT
+                  tensor_shape {
+                    dim { size: -1 }
+                  }
+                }
+              }
+            }
+          }
+          signature_def: {
+            key: "bar"
+            value: {
+              inputs: {
+                key: "inputs"
+                value { name: "input_node:0" }
+              }
+              method_name: "classify"
+              outputs: {
+                key: "outputs"
+                value {
+                  dtype: DT_FLOAT
+                  tensor_shape { dim { size: 1 } }
+                }
+              }
+            }
+          }
+        }
+    """, saved_model_pb2.SavedModel()))
+
+  def testExceptions(self):
+    with self.assertRaises(IOError):
+      updater = method_name_updater.MethodNameUpdater(
+          tempfile.mkdtemp(prefix=test.get_temp_dir()))
+
+    path = os.path.join(
+        compat.as_bytes(self._saved_model_path),
+        compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
+    file_io.write_string_to_file(
+        path, _SAVED_MODEL_PROTO.SerializeToString(deterministic=True))
+    updater = method_name_updater.MethodNameUpdater(self._saved_model_path)
+
+    with self.assertRaisesRegex(ValueError, "`signature_key` must be defined"):
+      updater.replace_method_name(
+          signature_key=None, method_name="classify")
+
+    with self.assertRaisesRegex(ValueError, "`method_name` must be defined"):
+      updater.replace_method_name(
+          signature_key="foobar", method_name="")
+
+    with self.assertRaisesRegex(
+        ValueError,
+        r"MetaGraphDef associated with tags \['gpu'\] could not be found"):
+      updater.replace_method_name(
+          signature_key="bar", method_name="classify", tags=["gpu"])
+
+    with self.assertRaisesRegex(
+        ValueError, r"MetaGraphDef associated with tags \['serve'\] does not "
+        r"have a signature_def with key: 'baz'"):
+      updater.replace_method_name(
+          signature_key="baz", method_name="classify", tags=["serve"])
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/metrics_test.py

@@ -0,0 +1,177 @@
+# 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 SavedModel instrumentation for Python reading/writing code.
+
+These tests verify that the counters are incremented correctly after SavedModel
+API calls.
+"""
+
+import os
+
+from google.protobuf import json_format
+
+from tensorflow.python.checkpoint.sharding import sharding_policies
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.saved_model import builder
+from tensorflow.python.saved_model import builder_impl
+from tensorflow.python.saved_model import fingerprinting
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import load_v1_in_v2
+from tensorflow.python.saved_model import loader_impl
+from tensorflow.python.saved_model import save
+from tensorflow.python.saved_model.pywrap_saved_model import metrics
+from tensorflow.python.trackable import autotrackable
+
+
+class MetricsTests(test.TestCase):
+
+  def _create_save_v2_model(self):
+    root = autotrackable.AutoTrackable()
+    save_dir = os.path.join(self.get_temp_dir(), "saved_model")
+    save.save(root, save_dir)
+    return save_dir
+
+  def _create_save_v1_model(self):
+    save_dir = os.path.join(self.get_temp_dir(), "builder")
+    builder_ = builder.SavedModelBuilder(save_dir)
+
+    with ops.Graph().as_default():
+      with self.session(graph=ops.Graph()) as sess:
+        constant_op.constant(5.0)
+        builder_.add_meta_graph_and_variables(sess, ["foo"])
+      builder_.save()
+    return save_dir
+
+  def test_python_save(self):
+    write_count = metrics.GetWrite(write_version="2")
+    save_api_count = metrics.GetWriteApi(save._SAVE_V2_LABEL)
+    _ = self._create_save_v2_model()
+
+    self.assertEqual(
+        metrics.GetWriteApi(save._SAVE_V2_LABEL), save_api_count + 1)
+    self.assertEqual(metrics.GetWrite(write_version="2"), write_count + 1)
+
+  def test_builder_save(self):
+    write_count = metrics.GetWrite(write_version="1")
+    save_builder_count = metrics.GetWriteApi(builder_impl._SAVE_BUILDER_LABEL)
+    _ = self._create_save_v1_model()
+
+    self.assertEqual(
+        metrics.GetWriteApi(builder_impl._SAVE_BUILDER_LABEL),
+        save_builder_count + 1)
+    self.assertEqual(metrics.GetWrite(write_version="1"), write_count + 1)
+
+  def test_load_v2(self):
+    save_dir = self._create_save_v2_model()
+
+    read_count = metrics.GetRead(write_version="2")
+    load_v2_count = metrics.GetReadApi(load._LOAD_V2_LABEL)
+    load.load(save_dir)
+
+    self.assertEqual(metrics.GetReadApi(load._LOAD_V2_LABEL), load_v2_count + 1)
+    self.assertEqual(metrics.GetRead(write_version="2"), read_count + 1)
+
+  def test_load_v1_in_v2(self):
+    save_dir = self._create_save_v1_model()
+    read_v1_count = metrics.GetRead(write_version="1")
+    read_v2_count = metrics.GetRead(write_version="2")
+    load_v2_count = metrics.GetReadApi(load._LOAD_V2_LABEL)
+    load_v1_v2_count = metrics.GetReadApi(load_v1_in_v2._LOAD_V1_V2_LABEL)
+    load.load(save_dir)
+
+    # Check that `load_v2` was *not* incremented.
+    self.assertEqual(metrics.GetReadApi(load._LOAD_V2_LABEL), load_v2_count)
+    self.assertEqual(metrics.GetRead(write_version="2"), read_v2_count)
+
+    self.assertEqual(
+        metrics.GetReadApi(load_v1_in_v2._LOAD_V1_V2_LABEL),
+        load_v1_v2_count + 1)
+    self.assertEqual(metrics.GetRead(write_version="1"), read_v1_count + 1)
+
+  def test_loader_v1(self):
+    ops.disable_eager_execution()
+    save_dir = self._create_save_v1_model()
+
+    read_count = metrics.GetRead(write_version="1")
+    loader = loader_impl.SavedModelLoader(save_dir)
+    with self.session(graph=ops.Graph()) as sess:
+      loader.load(sess, ["foo"])
+    ops.enable_eager_execution()
+
+    self.assertEqual(metrics.GetReadApi(loader_impl._LOADER_LABEL), 1)
+    self.assertEqual(metrics.GetRead(write_version="1"), read_count + 1)
+
+  def test_save_sets_write_fingerprint_metric(self):
+    exported_dir = self._create_save_v2_model()
+    fingerprint = fingerprinting.read_fingerprint(exported_dir)
+    fingerprint_metric = fingerprinting.Fingerprint.from_proto(
+        json_format.Parse(metrics.GetWriteFingerprint(),
+                          fingerprinting.fingerprint_pb2.FingerprintDef()))
+    self.assertEqual(fingerprint, fingerprint_metric)
+
+  def test_load_sets_read_fingerprint_metric(self):
+    exported_dir = self._create_save_v2_model()
+    load.load(exported_dir)
+    fingerprint = fingerprinting.read_fingerprint(exported_dir)
+    fingerprint_metric = fingerprinting.Fingerprint.from_proto(
+        json_format.Parse(metrics.GetReadFingerprint(),
+                          fingerprinting.fingerprint_pb2.FingerprintDef()))
+    self.assertEqual(fingerprint, fingerprint_metric)
+
+  def test_save_sets_write_path_metric(self):
+    exported_dir = self._create_save_v2_model()
+
+    self.assertEqual(metrics.GetWritePath(), exported_dir)
+
+  def test_load_sets_read_path_metric(self):
+    exported_dir = self._create_save_v2_model()
+    load.load(exported_dir)
+
+    self.assertEqual(metrics.GetReadPath(), exported_dir)
+
+  def test_save_sets_write_path_and_singleprint_metric(self):
+    exported_dir = self._create_save_v2_model()
+    singleprint = fingerprinting.read_fingerprint(exported_dir).singleprint()
+    path_and_singleprint_metric = metrics.GetWritePathAndSingleprint()
+    self.assertEqual(path_and_singleprint_metric, (exported_dir, singleprint))
+
+  def test_save_sets_read_path_and_singleprint_metric(self):
+    exported_dir = self._create_save_v2_model()
+    load.load(exported_dir)
+    singleprint = fingerprinting.read_fingerprint(exported_dir).singleprint()
+    path_and_singleprint_metric = metrics.GetReadPathAndSingleprint()
+    self.assertEqual(path_and_singleprint_metric, (exported_dir, singleprint))
+
+  def test_save_sets_sharding_callback_duration_metric(self):
+    self._create_save_v2_model()
+    sharding_callback_duration_metric = metrics.GetShardingCallbackDuration()
+    self.assertGreater(sharding_callback_duration_metric, 0)
+
+  def test_save_sets_num_checkpoint_shards_written_metric(self):
+    self._create_save_v2_model()
+    num_shards_written_metric = metrics.GetNumCheckpointShardsWritten()
+    self.assertGreater(num_shards_written_metric, 0)
+
+  def test_save_sets_sharding_callback_description_metric(self):
+    self._create_save_v2_model()
+    callback_description_metric = metrics.GetShardingCallbackDescription()
+    self.assertEqual(callback_description_metric,
+                     sharding_policies.ShardByTaskPolicy().description)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/model_utils/BUILD

@@ -0,0 +1,130 @@
+# Copyright 2016 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.
+# ==============================================================================
+
+# Description:
+#   Keras saving and loading libraries.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:__subpackages__"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "model_utils",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+    deps = [
+        ":export_output",
+        ":export_utils",
+        ":mode_keys",
+    ],
+)
+
+py_library(
+    name = "export_output",
+    srcs = ["export_output.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/saved_model:signature_def_utils",
+    ],
+)
+
+py_test(
+    name = "export_output_test",
+    srcs = ["export_output_test.py"],
+    strict_deps = True,
+    deps = [
+        ":export_output",
+        "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/ops:metrics",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:signature_constants",
+    ],
+)
+
+py_library(
+    name = "export_utils",
+    srcs = ["export_utils.py"],
+    strict_deps = True,
+    deps = [
+        ":export_output",
+        ":mode_keys",
+        "//tensorflow/python/lib/io:file_io",
+        "//tensorflow/python/ops:op_selector",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:signature_def_utils",
+        "//tensorflow/python/saved_model:tag_constants",
+        "//tensorflow/python/saved_model:utils",
+        "//tensorflow/python/util:compat",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:object_identity",
+    ],
+)
+
+py_test(
+    name = "export_test",
+    srcs = ["export_test.py"],
+    strict_deps = True,
+    deps = [
+        ":export_output",
+        ":export_utils",
+        ":mode_keys",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:signature_constants",
+        "//tensorflow/python/saved_model:signature_def_utils",
+    ],
+)
+
+py_library(
+    name = "mode_keys",
+    srcs = ["mode_keys.py"],
+    strict_deps = True,
+    deps = ["//tensorflow/python/util:compat"],
+)
+
+py_test(
+    name = "mode_keys_test",
+    srcs = ["mode_keys_test.py"],
+    strict_deps = True,
+    deps = [
+        ":mode_keys",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/saved_model/model_utils/__init__.py

@@ -0,0 +1,27 @@
+# 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.
+# ==============================================================================
+# LINT.IfChange
+"""Utils for saving a Keras Model to the SavedModel format."""
+# pylint: disable=wildcard-import
+from tensorflow.python.saved_model.model_utils.export_output import *
+from tensorflow.python.saved_model.model_utils.export_utils import build_all_signature_defs
+from tensorflow.python.saved_model.model_utils.export_utils import export_outputs_for_mode
+from tensorflow.python.saved_model.model_utils.export_utils import EXPORT_TAG_MAP
+from tensorflow.python.saved_model.model_utils.export_utils import get_export_outputs
+from tensorflow.python.saved_model.model_utils.export_utils import get_temp_export_dir
+from tensorflow.python.saved_model.model_utils.export_utils import get_timestamped_export_dir
+from tensorflow.python.saved_model.model_utils.export_utils import SIGNATURE_KEY_MAP
+# pylint: enable=wildcard-import
+# LINT.ThenChange(//tensorflow/python/keras/saving/utils_v1/__init__.py)

tensorflow/python/saved_model/model_utils/export_output.py

@@ -0,0 +1,425 @@
+# 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.
+# ==============================================================================
+# LINT.IfChange
+"""Classes for different types of export output."""
+
+import abc
+
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.saved_model import signature_def_utils
+
+
+class ExportOutput:
+  """Represents an output of a model that can be served.
+
+  These typically correspond to model heads.
+  """
+
+  __metaclass__ = abc.ABCMeta
+
+  _SEPARATOR_CHAR = '/'
+
+  @abc.abstractmethod
+  def as_signature_def(self, receiver_tensors):
+    """Generate a SignatureDef proto for inclusion in a MetaGraphDef.
+
+    The SignatureDef will specify outputs as described in this ExportOutput,
+    and will use the provided receiver_tensors as inputs.
+
+    Args:
+      receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
+        input nodes that will be fed.
+    """
+    pass
+
+  def _check_output_key(self, key, error_label):
+    # For multi-head models, the key can be a tuple.
+    if isinstance(key, tuple):
+      key = self._SEPARATOR_CHAR.join(key)
+
+    if not isinstance(key, str):
+      raise ValueError(
+          '{} output key must be a string; got {}.'.format(error_label, key))
+    return key
+
+  def _wrap_and_check_outputs(
+      self, outputs, single_output_default_name, error_label=None):
+    """Wraps raw tensors as dicts and checks type.
+
+    Note that we create a new dict here so that we can overwrite the keys
+    if necessary.
+
+    Args:
+      outputs: A `Tensor` or a dict of string to `Tensor`.
+      single_output_default_name: A string key for use in the output dict
+        if the provided `outputs` is a raw tensor.
+      error_label: descriptive string for use in error messages. If none,
+        single_output_default_name will be used.
+
+    Returns:
+      A dict of tensors
+
+    Raises:
+      ValueError: if the outputs dict keys are not strings or tuples of strings
+        or the values are not Tensors.
+    """
+    if not isinstance(outputs, dict):
+      outputs = {single_output_default_name: outputs}
+
+    output_dict = {}
+    for key, value in outputs.items():
+      error_name = error_label or single_output_default_name
+      key = self._check_output_key(key, error_name)
+      if not isinstance(value, tensor.Tensor):
+        raise ValueError(
+            '{} output value must be a Tensor; got {}.'.format(
+                error_name, value))
+
+      output_dict[key] = value
+    return output_dict
+
+
+class ClassificationOutput(ExportOutput):
+  """Represents the output of a classification head.
+
+  Either classes or scores or both must be set.
+
+  The classes `Tensor` must provide string labels, not integer class IDs.
+
+  If only classes is set, it is interpreted as providing top-k results in
+  descending order.
+
+  If only scores is set, it is interpreted as providing a score for every class
+  in order of class ID.
+
+  If both classes and scores are set, they are interpreted as zipped, so each
+  score corresponds to the class at the same index.  Clients should not depend
+  on the order of the entries.
+  """
+
+  def __init__(self, scores=None, classes=None):
+    """Constructor for `ClassificationOutput`.
+
+    Args:
+      scores: A float `Tensor` giving scores (sometimes but not always
+          interpretable as probabilities) for each class.  May be `None`, but
+          only if `classes` is set.  Interpretation varies-- see class doc.
+      classes: A string `Tensor` giving predicted class labels.  May be `None`,
+          but only if `scores` is set.  Interpretation varies-- see class doc.
+
+    Raises:
+      ValueError: if neither classes nor scores is set, or one of them is not a
+          `Tensor` with the correct dtype.
+    """
+    if (scores is not None
+        and not (isinstance(scores, tensor.Tensor)
+                 and scores.dtype.is_floating)):
+      raise ValueError('Classification scores must be a float32 Tensor; '
+                       'got {}'.format(scores))
+    if (classes is not None
+        and not (isinstance(classes, tensor.Tensor)
+                 and dtypes.as_dtype(classes.dtype) == dtypes.string)):
+      raise ValueError('Classification classes must be a string Tensor; '
+                       'got {}'.format(classes))
+    if scores is None and classes is None:
+      raise ValueError('Cannot create a ClassificationOutput with empty '
+                       'arguments. At least one of `scores` and `classes` '
+                       'must be defined.')
+    self._scores = scores
+    self._classes = classes
+
+  @property
+  def scores(self):
+    return self._scores
+
+  @property
+  def classes(self):
+    return self._classes
+
+  def as_signature_def(self, receiver_tensors):
+    if len(receiver_tensors) != 1:
+      raise ValueError(
+          'Classification signatures can only accept a single tensor input of '
+          'type tf.string. Please check to make sure that you have structured '
+          'the serving_input_receiver_fn so that it creates a single string '
+          'placeholder. If your model function expects multiple inputs, then '
+          'use `tf.io.parse_example()` to parse the string into multiple '
+          f'tensors.\n Received: {receiver_tensors}')
+    (_, examples), = receiver_tensors.items()
+    if dtypes.as_dtype(examples.dtype) != dtypes.string:
+      raise ValueError(
+          'Classification signatures can only accept a single tensor input of '
+          'type tf.string. Please check to make sure that you have structured '
+          'the serving_input_receiver_fn so that it creates a single string '
+          'placeholder. If your model function expects multiple inputs, then '
+          'use `tf.io.parse_example()` to parse the string into multiple '
+          f'tensors.\n Received: {receiver_tensors}')
+    return signature_def_utils.classification_signature_def(
+        examples, self.classes, self.scores)
+
+
+class RegressionOutput(ExportOutput):
+  """Represents the output of a regression head."""
+
+  def __init__(self, value):
+    """Constructor for `RegressionOutput`.
+
+    Args:
+      value: a float `Tensor` giving the predicted values.  Required.
+
+    Raises:
+      ValueError: if the value is not a `Tensor` with dtype tf.float32.
+    """
+    if not (isinstance(value, tensor.Tensor) and value.dtype.is_floating):
+      raise ValueError('Regression output value must be a float32 Tensor; '
+                       'got {}'.format(value))
+    self._value = value
+
+  @property
+  def value(self):
+    return self._value
+
+  def as_signature_def(self, receiver_tensors):
+    if len(receiver_tensors) != 1:
+      raise ValueError(
+          'Regression signatures can only accept a single tensor input of '
+          'type tf.string. Please check to make sure that you have structured '
+          'the serving_input_receiver_fn so that it creates a single string '
+          'placeholder. If your model function expects multiple inputs, then '
+          'use `tf.io.parse_example()` to parse the string into multiple '
+          f'tensors.\n Received: {receiver_tensors}')
+    (_, examples), = receiver_tensors.items()
+    if dtypes.as_dtype(examples.dtype) != dtypes.string:
+      raise ValueError(
+          'Regression signatures can only accept a single tensor input of '
+          'type tf.string. Please check to make sure that you have structured '
+          'the serving_input_receiver_fn so that it creates a single string '
+          'placeholder. If your model function expects multiple inputs, then '
+          'use `tf.io.parse_example()` to parse the string into multiple '
+          f'tensors.\n Received: {receiver_tensors}')
+    return signature_def_utils.regression_signature_def(examples, self.value)
+
+
+class PredictOutput(ExportOutput):
+  """Represents the output of a generic prediction head.
+
+  A generic prediction need not be either a classification or a regression.
+
+  Named outputs must be provided as a dict from string to `Tensor`,
+  """
+  _SINGLE_OUTPUT_DEFAULT_NAME = 'output'
+
+  def __init__(self, outputs):
+    """Constructor for PredictOutput.
+
+    Args:
+      outputs: A `Tensor` or a dict of string to `Tensor` representing the
+        predictions.
+
+    Raises:
+      ValueError: if the outputs is not dict, or any of its keys are not
+          strings, or any of its values are not `Tensor`s.
+    """
+
+    self._outputs = self._wrap_and_check_outputs(
+        outputs, self._SINGLE_OUTPUT_DEFAULT_NAME, error_label='Prediction')
+
+  @property
+  def outputs(self):
+    return self._outputs
+
+  def as_signature_def(self, receiver_tensors):
+    return signature_def_utils.predict_signature_def(receiver_tensors,
+                                                     self.outputs)
+
+
+class _SupervisedOutput(ExportOutput):
+  """Represents the output of a supervised training or eval process."""
+  __metaclass__ = abc.ABCMeta
+
+  LOSS_NAME = 'loss'
+  PREDICTIONS_NAME = 'predictions'
+  METRICS_NAME = 'metrics'
+
+  METRIC_VALUE_SUFFIX = 'value'
+  METRIC_UPDATE_SUFFIX = 'update_op'
+
+  _loss = None
+  _predictions = None
+  _metrics = None
+
+  def __init__(self, loss=None, predictions=None, metrics=None):
+    """Constructor for SupervisedOutput (ie, Train or Eval output).
+
+    Args:
+      loss: dict of Tensors or single Tensor representing calculated loss.
+      predictions: dict of Tensors or single Tensor representing model
+        predictions.
+      metrics: Dict of metric results keyed by name.
+        The values of the dict can be one of the following:
+        (1) instance of `Metric` class.
+        (2) (metric_value, update_op) tuples, or a single tuple.
+        metric_value must be a Tensor, and update_op must be a Tensor or Op.
+
+    Raises:
+      ValueError: if any of the outputs' dict keys are not strings or tuples of
+        strings or the values are not Tensors (or Operations in the case of
+        update_op).
+    """
+
+    if loss is not None:
+      loss_dict = self._wrap_and_check_outputs(loss, self.LOSS_NAME)
+      self._loss = self._prefix_output_keys(loss_dict, self.LOSS_NAME)
+    if predictions is not None:
+      pred_dict = self._wrap_and_check_outputs(
+          predictions, self.PREDICTIONS_NAME)
+      self._predictions = self._prefix_output_keys(
+          pred_dict, self.PREDICTIONS_NAME)
+    if metrics is not None:
+      self._metrics = self._wrap_and_check_metrics(metrics)
+
+  def _prefix_output_keys(self, output_dict, output_name):
+    """Prepend output_name to the output_dict keys if it doesn't exist.
+
+    This produces predictable prefixes for the pre-determined outputs
+    of SupervisedOutput.
+
+    Args:
+      output_dict: dict of string to Tensor, assumed valid.
+      output_name: prefix string to prepend to existing keys.
+
+    Returns:
+      dict with updated keys and existing values.
+    """
+
+    new_outputs = {}
+    for key, val in output_dict.items():
+      key = self._prefix_key(key, output_name)
+      new_outputs[key] = val
+    return new_outputs
+
+  def _prefix_key(self, key, output_name):
+    if key.find(output_name) != 0:
+      key = output_name + self._SEPARATOR_CHAR + key
+    return key
+
+  def _wrap_and_check_metrics(self, metrics):
+    """Handle the saving of metrics.
+
+    Metrics is either a tuple of (value, update_op), or a dict of such tuples.
+    Here, we separate out the tuples and create a dict with names to tensors.
+
+    Args:
+      metrics: Dict of metric results keyed by name.
+        The values of the dict can be one of the following:
+        (1) instance of `Metric` class.
+        (2) (metric_value, update_op) tuples, or a single tuple.
+        metric_value must be a Tensor, and update_op must be a Tensor or Op.
+
+    Returns:
+      dict of output_names to tensors
+
+    Raises:
+      ValueError: if the dict key is not a string, or the metric values or ops
+        are not tensors.
+    """
+    if not isinstance(metrics, dict):
+      metrics = {self.METRICS_NAME: metrics}
+
+    outputs = {}
+    for key, value in metrics.items():
+      if isinstance(value, tuple):
+        metric_val, metric_op = value
+      else:  # value is a keras.Metrics object
+        metric_val = value.result()
+        assert len(value.updates) == 1  # We expect only one update op.
+        metric_op = value.updates[0]
+      key = self._check_output_key(key, self.METRICS_NAME)
+      key = self._prefix_key(key, self.METRICS_NAME)
+
+      val_name = key + self._SEPARATOR_CHAR + self.METRIC_VALUE_SUFFIX
+      op_name = key + self._SEPARATOR_CHAR + self.METRIC_UPDATE_SUFFIX
+      if not isinstance(metric_val, tensor.Tensor):
+        raise ValueError(
+            '{} output value must be a Tensor; got {}.'.format(
+                key, metric_val))
+      if not (tensor_util.is_tf_type(metric_op) or
+              isinstance(metric_op, ops.Operation)):
+        raise ValueError(
+            '{} update_op must be a Tensor or Operation; got {}.'.format(
+                key, metric_op))
+
+      # We must wrap any ops (or variables) in a Tensor before export, as the
+      # SignatureDef proto expects tensors only. See b/109740581
+      metric_op_tensor = metric_op
+      if not isinstance(metric_op, tensor.Tensor):
+        with ops.control_dependencies([metric_op]):
+          metric_op_tensor = constant_op.constant([], name='metric_op_wrapper')
+
+      outputs[val_name] = metric_val
+      outputs[op_name] = metric_op_tensor
+
+    return outputs
+
+  @property
+  def loss(self):
+    return self._loss
+
+  @property
+  def predictions(self):
+    return self._predictions
+
+  @property
+  def metrics(self):
+    return self._metrics
+
+  @abc.abstractmethod
+  def _get_signature_def_fn(self):
+    """Returns a function that produces a SignatureDef given desired outputs."""
+    pass
+
+  def as_signature_def(self, receiver_tensors):
+    signature_def_fn = self._get_signature_def_fn()
+    return signature_def_fn(
+        receiver_tensors, self.loss, self.predictions, self.metrics)
+
+
+class TrainOutput(_SupervisedOutput):
+  """Represents the output of a supervised training process.
+
+  This class generates the appropriate signature def for exporting
+  training output by type-checking and wrapping loss, predictions, and metrics
+  values.
+  """
+
+  def _get_signature_def_fn(self):
+    return signature_def_utils.supervised_train_signature_def
+
+
+class EvalOutput(_SupervisedOutput):
+  """Represents the output of a supervised eval process.
+
+  This class generates the appropriate signature def for exporting
+  eval output by type-checking and wrapping loss, predictions, and metrics
+  values.
+  """
+
+  def _get_signature_def_fn(self):
+    return signature_def_utils.supervised_eval_signature_def

tensorflow/python/saved_model/model_utils/export_output_test.py

@@ -0,0 +1,400 @@
+# 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.
+# ==============================================================================
+"""Tests for export."""
+
+from tensorflow.core.framework import tensor_shape_pb2
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import metrics as metrics_module
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model.model_utils import export_output as export_output_lib
+
+
+class ExportOutputTest(test.TestCase):
+
+  def test_regress_value_must_be_float(self):
+    with context.graph_mode():
+      value = array_ops.placeholder(dtypes.string, 1, name='output-tensor-1')
+      with self.assertRaisesRegex(
+          ValueError, 'Regression output value must be a float32 Tensor'):
+        export_output_lib.RegressionOutput(value)
+
+  def test_classify_classes_must_be_strings(self):
+    with context.graph_mode():
+      classes = array_ops.placeholder(dtypes.float32, 1, name='output-tensor-1')
+      with self.assertRaisesRegex(
+          ValueError, 'Classification classes must be a string Tensor'):
+        export_output_lib.ClassificationOutput(classes=classes)
+
+  def test_classify_scores_must_be_float(self):
+    with context.graph_mode():
+      scores = array_ops.placeholder(dtypes.string, 1, name='output-tensor-1')
+      with self.assertRaisesRegex(
+          ValueError, 'Classification scores must be a float32 Tensor'):
+        export_output_lib.ClassificationOutput(scores=scores)
+
+  def test_classify_requires_classes_or_scores(self):
+    with self.assertRaisesRegex(
+        ValueError,
+        'Cannot create a ClassificationOutput with empty arguments'):
+      export_output_lib.ClassificationOutput()
+
+  def test_build_standardized_signature_def_regression(self):
+    with context.graph_mode():
+      input_tensors = {
+          'input-1':
+              array_ops.placeholder(
+                  dtypes.string, 1, name='input-tensor-1')
+      }
+      value = array_ops.placeholder(dtypes.float32, 1, name='output-tensor-1')
+
+      export_output = export_output_lib.RegressionOutput(value)
+      actual_signature_def = export_output.as_signature_def(input_tensors)
+
+      expected_signature_def = meta_graph_pb2.SignatureDef()
+      shape = tensor_shape_pb2.TensorShapeProto(
+          dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
+      dtype_float = types_pb2.DataType.Value('DT_FLOAT')
+      dtype_string = types_pb2.DataType.Value('DT_STRING')
+      expected_signature_def.inputs[
+          signature_constants.REGRESS_INPUTS].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='input-tensor-1:0',
+                                        dtype=dtype_string,
+                                        tensor_shape=shape))
+      expected_signature_def.outputs[
+          signature_constants.REGRESS_OUTPUTS].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='output-tensor-1:0',
+                                        dtype=dtype_float,
+                                        tensor_shape=shape))
+
+      expected_signature_def.method_name = (
+          signature_constants.REGRESS_METHOD_NAME)
+      self.assertEqual(actual_signature_def, expected_signature_def)
+
+  def test_build_standardized_signature_def_classify_classes_only(self):
+    """Tests classification with one output tensor."""
+    with context.graph_mode():
+      input_tensors = {
+          'input-1':
+              array_ops.placeholder(
+                  dtypes.string, 1, name='input-tensor-1')
+      }
+      classes = array_ops.placeholder(dtypes.string, 1, name='output-tensor-1')
+
+      export_output = export_output_lib.ClassificationOutput(classes=classes)
+      actual_signature_def = export_output.as_signature_def(input_tensors)
+
+      expected_signature_def = meta_graph_pb2.SignatureDef()
+      shape = tensor_shape_pb2.TensorShapeProto(
+          dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
+      dtype_string = types_pb2.DataType.Value('DT_STRING')
+      expected_signature_def.inputs[
+          signature_constants.CLASSIFY_INPUTS].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='input-tensor-1:0',
+                                        dtype=dtype_string,
+                                        tensor_shape=shape))
+      expected_signature_def.outputs[
+          signature_constants.CLASSIFY_OUTPUT_CLASSES].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='output-tensor-1:0',
+                                        dtype=dtype_string,
+                                        tensor_shape=shape))
+
+      expected_signature_def.method_name = (
+          signature_constants.CLASSIFY_METHOD_NAME)
+      self.assertEqual(actual_signature_def, expected_signature_def)
+
+  def test_build_standardized_signature_def_classify_both(self):
+    """Tests multiple output tensors that include classes and scores."""
+    with context.graph_mode():
+      input_tensors = {
+          'input-1':
+              array_ops.placeholder(
+                  dtypes.string, 1, name='input-tensor-1')
+      }
+      classes = array_ops.placeholder(dtypes.string, 1,
+                                      name='output-tensor-classes')
+      scores = array_ops.placeholder(dtypes.float32, 1,
+                                     name='output-tensor-scores')
+
+      export_output = export_output_lib.ClassificationOutput(
+          scores=scores, classes=classes)
+      actual_signature_def = export_output.as_signature_def(input_tensors)
+
+      expected_signature_def = meta_graph_pb2.SignatureDef()
+      shape = tensor_shape_pb2.TensorShapeProto(
+          dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
+      dtype_float = types_pb2.DataType.Value('DT_FLOAT')
+      dtype_string = types_pb2.DataType.Value('DT_STRING')
+      expected_signature_def.inputs[
+          signature_constants.CLASSIFY_INPUTS].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='input-tensor-1:0',
+                                        dtype=dtype_string,
+                                        tensor_shape=shape))
+      expected_signature_def.outputs[
+          signature_constants.CLASSIFY_OUTPUT_CLASSES].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='output-tensor-classes:0',
+                                        dtype=dtype_string,
+                                        tensor_shape=shape))
+      expected_signature_def.outputs[
+          signature_constants.CLASSIFY_OUTPUT_SCORES].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='output-tensor-scores:0',
+                                        dtype=dtype_float,
+                                        tensor_shape=shape))
+
+      expected_signature_def.method_name = (
+          signature_constants.CLASSIFY_METHOD_NAME)
+      self.assertEqual(actual_signature_def, expected_signature_def)
+
+  def test_build_standardized_signature_def_classify_scores_only(self):
+    """Tests classification without classes tensor."""
+    with context.graph_mode():
+      input_tensors = {
+          'input-1':
+              array_ops.placeholder(
+                  dtypes.string, 1, name='input-tensor-1')
+      }
+
+      scores = array_ops.placeholder(dtypes.float32, 1,
+                                     name='output-tensor-scores')
+
+      export_output = export_output_lib.ClassificationOutput(
+          scores=scores)
+      actual_signature_def = export_output.as_signature_def(input_tensors)
+
+      expected_signature_def = meta_graph_pb2.SignatureDef()
+      shape = tensor_shape_pb2.TensorShapeProto(
+          dim=[tensor_shape_pb2.TensorShapeProto.Dim(size=1)])
+      dtype_float = types_pb2.DataType.Value('DT_FLOAT')
+      dtype_string = types_pb2.DataType.Value('DT_STRING')
+      expected_signature_def.inputs[
+          signature_constants.CLASSIFY_INPUTS].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='input-tensor-1:0',
+                                        dtype=dtype_string,
+                                        tensor_shape=shape))
+      expected_signature_def.outputs[
+          signature_constants.CLASSIFY_OUTPUT_SCORES].CopyFrom(
+              meta_graph_pb2.TensorInfo(name='output-tensor-scores:0',
+                                        dtype=dtype_float,
+                                        tensor_shape=shape))
+
+      expected_signature_def.method_name = (
+          signature_constants.CLASSIFY_METHOD_NAME)
+      self.assertEqual(actual_signature_def, expected_signature_def)
+
+  def test_predict_outputs_valid(self):
+    """Tests that no errors are raised when provided outputs are valid."""
+    outputs = {
+        'output0': constant_op.constant([0]),
+        u'output1': constant_op.constant(['foo']),
+    }
+    export_output_lib.PredictOutput(outputs)
+
+    # Single Tensor is OK too
+    export_output_lib.PredictOutput(constant_op.constant([0]))
+
+  def test_predict_outputs_invalid(self):
+    with self.assertRaisesRegex(ValueError,
+                                'Prediction output key must be a string'):
+      export_output_lib.PredictOutput({1: constant_op.constant([0])})
+
+    with self.assertRaisesRegex(ValueError,
+                                'Prediction output value must be a Tensor'):
+      export_output_lib.PredictOutput({
+          'prediction1': sparse_tensor.SparseTensor(
+              indices=[[0, 0]], values=[1], dense_shape=[1, 1]),
+      })
+
+
+class MockSupervisedOutput(export_output_lib._SupervisedOutput):
+  """So that we can test the abstract class methods directly."""
+
+  def _get_signature_def_fn(self):
+    pass
+
+
+class SupervisedOutputTest(test.TestCase):
+
+  def test_supervised_outputs_valid(self):
+    """Tests that no errors are raised when provided outputs are valid."""
+    with context.graph_mode():
+      loss = {'my_loss': constant_op.constant([0])}
+      predictions = {u'output1': constant_op.constant(['foo'])}
+      mean, update_op = metrics_module.mean_tensor(constant_op.constant([0]))
+      metrics = {
+          'metrics': (mean, update_op),
+          'metrics2': (constant_op.constant([0]), constant_op.constant([10]))
+      }
+
+      outputter = MockSupervisedOutput(loss, predictions, metrics)
+      self.assertEqual(outputter.loss['loss/my_loss'], loss['my_loss'])
+      self.assertEqual(
+          outputter.predictions['predictions/output1'], predictions['output1'])
+      self.assertEqual(outputter.metrics['metrics/update_op'].name,
+                       'mean/update_op:0')
+      self.assertEqual(
+          outputter.metrics['metrics2/update_op'], metrics['metrics2'][1])
+
+      # Single Tensor is OK too
+      outputter = MockSupervisedOutput(
+          loss['my_loss'], predictions['output1'], metrics['metrics'])
+      self.assertEqual(outputter.loss, {'loss': loss['my_loss']})
+      self.assertEqual(
+          outputter.predictions, {'predictions': predictions['output1']})
+      self.assertEqual(outputter.metrics['metrics/update_op'].name,
+                       'mean/update_op:0')
+
+  def test_supervised_outputs_none(self):
+    outputter = MockSupervisedOutput(
+        constant_op.constant([0]), None, None)
+    self.assertLen(outputter.loss, 1)
+    self.assertIsNone(outputter.predictions)
+    self.assertIsNone(outputter.metrics)
+
+  def test_supervised_outputs_invalid(self):
+    with self.assertRaisesRegex(ValueError, 'predictions output value must'):
+      MockSupervisedOutput(constant_op.constant([0]), [3], None)
+    with self.assertRaisesRegex(ValueError, 'loss output value must'):
+      MockSupervisedOutput('str', None, None)
+    with self.assertRaisesRegex(ValueError, 'metrics output value must'):
+      MockSupervisedOutput(None, None, (15.3, 4))
+    with self.assertRaisesRegex(ValueError, 'loss output key must'):
+      MockSupervisedOutput({25: 'Tensor'}, None, None)
+
+  def test_supervised_outputs_tuples(self):
+    """Tests that no errors are raised when provided outputs are valid."""
+    with context.graph_mode():
+      loss = {('my', 'loss'): constant_op.constant([0])}
+      predictions = {(u'output1', '2'): constant_op.constant(['foo'])}
+      mean, update_op = metrics_module.mean_tensor(constant_op.constant([0]))
+      metrics = {
+          ('metrics', '1'): (mean, update_op),
+          ('metrics', '2'): (constant_op.constant([0]),
+                             constant_op.constant([10]))
+      }
+
+      outputter = MockSupervisedOutput(loss, predictions, metrics)
+      self.assertEqual(set(outputter.loss.keys()), set(['loss/my/loss']))
+      self.assertEqual(set(outputter.predictions.keys()),
+                       set(['predictions/output1/2']))
+      self.assertEqual(
+          set(outputter.metrics.keys()),
+          set([
+              'metrics/1/value', 'metrics/1/update_op', 'metrics/2/value',
+              'metrics/2/update_op'
+          ]))
+
+  def test_supervised_outputs_no_prepend(self):
+    """Tests that no errors are raised when provided outputs are valid."""
+    with context.graph_mode():
+      loss = {'loss': constant_op.constant([0])}
+      predictions = {u'predictions': constant_op.constant(['foo'])}
+      mean, update_op = metrics_module.mean_tensor(constant_op.constant([0]))
+      metrics = {
+          'metrics_1': (mean, update_op),
+          'metrics_2': (constant_op.constant([0]), constant_op.constant([10]))
+      }
+
+      outputter = MockSupervisedOutput(loss, predictions, metrics)
+      self.assertEqual(set(outputter.loss.keys()), set(['loss']))
+      self.assertEqual(set(outputter.predictions.keys()), set(['predictions']))
+      self.assertEqual(
+          set(outputter.metrics.keys()),
+          set([
+              'metrics_1/value', 'metrics_1/update_op', 'metrics_2/update_op',
+              'metrics_2/value'
+          ]))
+
+  def test_train_signature_def(self):
+    with context.graph_mode():
+      loss = {'my_loss': constant_op.constant([0])}
+      predictions = {u'output1': constant_op.constant(['foo'])}
+      mean, update_op = metrics_module.mean_tensor(constant_op.constant([0]))
+      metrics = {
+          'metrics_1': (mean, update_op),
+          'metrics_2': (constant_op.constant([0]), constant_op.constant([10]))
+      }
+
+      outputter = export_output_lib.TrainOutput(loss, predictions, metrics)
+
+      receiver = {u'features': constant_op.constant(100, shape=(100, 2)),
+                  'labels': constant_op.constant(100, shape=(100, 1))}
+      sig_def = outputter.as_signature_def(receiver)
+
+      self.assertIn('loss/my_loss', sig_def.outputs)
+      self.assertIn('metrics_1/value', sig_def.outputs)
+      self.assertIn('metrics_2/value', sig_def.outputs)
+      self.assertIn('predictions/output1', sig_def.outputs)
+      self.assertIn('features', sig_def.inputs)
+
+  def test_eval_signature_def(self):
+    with context.graph_mode():
+      loss = {'my_loss': constant_op.constant([0])}
+      predictions = {u'output1': constant_op.constant(['foo'])}
+
+      outputter = export_output_lib.EvalOutput(loss, predictions, None)
+
+      receiver = {u'features': constant_op.constant(100, shape=(100, 2)),
+                  'labels': constant_op.constant(100, shape=(100, 1))}
+      sig_def = outputter.as_signature_def(receiver)
+
+      self.assertIn('loss/my_loss', sig_def.outputs)
+      self.assertNotIn('metrics/value', sig_def.outputs)
+      self.assertIn('predictions/output1', sig_def.outputs)
+      self.assertIn('features', sig_def.inputs)
+
+  def test_metric_op_is_tensor(self):
+    """Tests that ops.Operation is wrapped by a tensor for metric_ops."""
+    with context.graph_mode():
+      loss = {'my_loss': constant_op.constant([0])}
+      predictions = {u'output1': constant_op.constant(['foo'])}
+      mean, update_op = metrics_module.mean_tensor(constant_op.constant([0]))
+      metrics = {
+          'metrics_1': (mean, update_op),
+          'metrics_2': (constant_op.constant([0]), control_flow_ops.no_op()),
+          # Keras metric's update_state() could return a Variable, rather than
+          # an Operation or Tensor.
+          'keras_1': (constant_op.constant([0.5]),
+                      variables.Variable(1.0, name='AssignAddVariableOp_3'))
+      }
+
+      outputter = MockSupervisedOutput(loss, predictions, metrics)
+      # If we get there, it means constructor succeeded; which is sufficient
+      # for testing the constructor.
+
+      self.assertTrue(outputter.metrics['metrics_1/update_op'].name.startswith(
+          'mean/update_op'))
+      self.assertIsInstance(
+          outputter.metrics['metrics_1/update_op'], tensor.Tensor)
+      self.assertIsInstance(outputter.metrics['metrics_1/value'], tensor.Tensor)
+
+      self.assertEqual(outputter.metrics['metrics_2/value'],
+                       metrics['metrics_2'][0])
+      self.assertTrue(outputter.metrics['metrics_2/update_op'].name.startswith(
+          'metric_op_wrapper'))
+      self.assertIsInstance(
+          outputter.metrics['metrics_2/update_op'], tensor.Tensor)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/saved_model/model_utils/export_test.py

@@ -0,0 +1,335 @@
+# 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.
+# ==============================================================================
+"""Tests for export utils."""
+
+import os
+import tempfile
+import time
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model.model_utils import export_output
+from tensorflow.python.saved_model.model_utils import export_utils
+from tensorflow.python.saved_model.model_utils.mode_keys import KerasModeKeys
+
+
+class ExportTest(test_util.TensorFlowTestCase):
+
+  def test_build_all_signature_defs_without_receiver_alternatives(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that depends on graph-only functions such
+    # as build_tensor_info.
+    with ops.Graph().as_default():
+      receiver_tensor = array_ops.placeholder(dtypes.string)
+      output_1 = constant_op.constant([1.])
+      output_2 = constant_op.constant(["2"])
+      output_3 = constant_op.constant(["3"])
+      export_outputs = {
+          signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+              export_output.RegressionOutput(value=output_1),
+          "head-2":
+              export_output.ClassificationOutput(classes=output_2),
+          "head-3":
+              export_output.PredictOutput(outputs={"some_output_3": output_3}),
+      }
+
+      signature_defs = export_utils.build_all_signature_defs(
+          receiver_tensor, export_outputs)
+
+      expected_signature_defs = {
+          "serving_default":
+              signature_def_utils.regression_signature_def(
+                  receiver_tensor, output_1),
+          "head-2":
+              signature_def_utils.classification_signature_def(
+                  receiver_tensor, output_2, None),
+          "head-3":
+              signature_def_utils.predict_signature_def(
+                  {"input": receiver_tensor}, {"some_output_3": output_3})
+      }
+
+      self.assertDictEqual(expected_signature_defs, signature_defs)
+
+  def test_build_all_signature_defs_with_dict_alternatives(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that depends on graph-only functions such
+    # as build_tensor_info.
+    with ops.Graph().as_default():
+      receiver_tensor = array_ops.placeholder(dtypes.string)
+
+      receiver_tensors_alternative_1 = {
+          "foo": array_ops.placeholder(dtypes.int64),
+          "bar": array_ops.sparse_placeholder(dtypes.float32)
+      }
+
+      unfed_input = array_ops.placeholder(dtypes.bool)
+      receiver_tensors_alternative_2 = {"unfed": unfed_input}
+
+      receiver_tensors_alternatives = {
+          "other": receiver_tensors_alternative_1,
+          "with_unfed_input": receiver_tensors_alternative_2
+      }
+
+      output_1 = constant_op.constant([1.])
+      output_2 = constant_op.constant(["2"])
+      output_3 = constant_op.constant(["3"])
+      output_4 = unfed_input
+      output_5 = math_ops.logical_not(unfed_input)
+      export_outputs = {
+          signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+              export_output.RegressionOutput(value=output_1),
+          "head-2":
+              export_output.ClassificationOutput(classes=output_2),
+          "head-3":
+              export_output.PredictOutput(outputs={"some_output_3": output_3}),
+          "head-4":
+              export_output.PredictOutput(outputs={"some_output_4": output_4}),
+          "head-5":
+              export_output.PredictOutput(outputs={"some_output_5": output_5}),
+      }
+
+      signature_defs = export_utils.build_all_signature_defs(
+          receiver_tensor, export_outputs, receiver_tensors_alternatives)
+
+      expected_signature_defs = {
+          "serving_default":
+              signature_def_utils.regression_signature_def(
+                  receiver_tensor, output_1),
+          "head-2":
+              signature_def_utils.classification_signature_def(
+                  receiver_tensor, output_2, None),
+          "head-3":
+              signature_def_utils.predict_signature_def(
+                  {"input": receiver_tensor}, {"some_output_3": output_3}),
+          "other:head-3":
+              signature_def_utils.predict_signature_def(
+                  receiver_tensors_alternative_1, {"some_output_3": output_3}),
+
+          # Note that the alternatives 'other:serving_default' and
+          # 'other:head-2' are invalid, because regression and classification
+          # signatures must take a single string input.  Here we verify that
+          # these invalid signatures are not included in the export_utils.
+
+          # Similarly, we verify that 'head-4' and 'head-5', which depend on an
+          # input that is not being fed as a receiver tensor, are also omitted.
+
+          # All the three heads are present when that input is fed, however:
+          "with_unfed_input:head-3":
+              signature_def_utils.predict_signature_def(
+                  receiver_tensors_alternative_2, {"some_output_3": output_3}),
+          "with_unfed_input:head-4":
+              signature_def_utils.predict_signature_def(
+                  receiver_tensors_alternative_2, {"some_output_4": output_4}),
+          "with_unfed_input:head-5":
+              signature_def_utils.predict_signature_def(
+                  receiver_tensors_alternative_2, {"some_output_5": output_5})
+      }
+
+      self.assertDictEqual(expected_signature_defs, signature_defs)
+
+  def test_build_all_signature_defs_with_single_alternatives(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that depends on graph-only functions such
+    # as build_tensor_info.
+    with ops.Graph().as_default():
+      receiver_tensor = array_ops.placeholder(dtypes.string)
+      receiver_tensors_alternative_1 = array_ops.placeholder(dtypes.int64)
+      receiver_tensors_alternative_2 = array_ops.sparse_placeholder(
+          dtypes.float32)
+      # Note we are passing single Tensors as values of
+      # receiver_tensors_alternatives, where normally that is a dict.
+      # In this case a dict will be created using the default receiver tensor
+      # name "input".
+      receiver_tensors_alternatives = {
+          "other1": receiver_tensors_alternative_1,
+          "other2": receiver_tensors_alternative_2
+      }
+      output_1 = constant_op.constant([1.])
+      output_2 = constant_op.constant(["2"])
+      output_3 = constant_op.constant(["3"])
+      export_outputs = {
+          signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+              export_output.RegressionOutput(value=output_1),
+          "head-2":
+              export_output.ClassificationOutput(classes=output_2),
+          "head-3":
+              export_output.PredictOutput(outputs={"some_output_3": output_3}),
+      }
+
+      signature_defs = export_utils.build_all_signature_defs(
+          receiver_tensor, export_outputs, receiver_tensors_alternatives)
+
+      expected_signature_defs = {
+          "serving_default":
+              signature_def_utils.regression_signature_def(
+                  receiver_tensor, output_1),
+          "head-2":
+              signature_def_utils.classification_signature_def(
+                  receiver_tensor, output_2, None),
+          "head-3":
+              signature_def_utils.predict_signature_def(
+                  {"input": receiver_tensor}, {"some_output_3": output_3}),
+          "other1:head-3":
+              signature_def_utils.predict_signature_def(
+                  {"input": receiver_tensors_alternative_1},
+                  {"some_output_3": output_3}),
+          "other2:head-3":
+              signature_def_utils.predict_signature_def(
+                  {"input": receiver_tensors_alternative_2},
+                  {"some_output_3": output_3})
+
+          # Note that the alternatives 'other:serving_default' and
+          # 'other:head-2' are invalid, because regression and classification
+          # signatures must take a single string input.  Here we verify that
+          # these invalid signatures are not included in the export_utils.
+      }
+
+      self.assertDictEqual(expected_signature_defs, signature_defs)
+
+  def test_build_all_signature_defs_export_outputs_required(self):
+    receiver_tensor = constant_op.constant(["11"])
+
+    with self.assertRaises(ValueError) as e:
+      export_utils.build_all_signature_defs(receiver_tensor, None)
+
+    self.assertTrue(
+        str(e.exception).startswith("`export_outputs` must be a dict"))
+
+  def test_get_timestamped_export_dir(self):
+    export_dir_base = tempfile.mkdtemp() + "export/"
+    export_dir_1 = export_utils.get_timestamped_export_dir(
+        export_dir_base)
+    time.sleep(2)
+    export_dir_2 = export_utils.get_timestamped_export_dir(
+        export_dir_base)
+    time.sleep(2)
+    export_dir_3 = export_utils.get_timestamped_export_dir(
+        export_dir_base)
+
+    # Export directories should be named using a timestamp that is seconds
+    # since epoch.  Such a timestamp is 10 digits long.
+    time_1 = os.path.basename(export_dir_1)
+    self.assertEqual(10, len(time_1))
+    time_2 = os.path.basename(export_dir_2)
+    self.assertEqual(10, len(time_2))
+    time_3 = os.path.basename(export_dir_3)
+    self.assertEqual(10, len(time_3))
+
+    self.assertLess(int(time_1), int(time_2))
+    self.assertLess(int(time_2), int(time_3))
+
+  def test_get_temp_export_dir(self):
+    export_dir = os.path.join("tmp", "export", "1576013284")
+    tmp_export_dir = export_utils.get_temp_export_dir(export_dir)
+    self.assertEqual(tmp_export_dir,
+                     os.path.join(b"tmp", b"export", b"temp-1576013284"))
+
+    export_dir = os.path.join(b"tmp", b"export", b"1576013284")
+    tmp_export_dir = export_utils.get_temp_export_dir(export_dir)
+    self.assertEqual(tmp_export_dir,
+                     os.path.join(b"tmp", b"export", b"temp-1576013284"))
+
+  def test_build_all_signature_defs_serving_only(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that depends on graph-only functions such
+    # as build_tensor_info.
+    with ops.Graph().as_default():
+      receiver_tensor = {"input": array_ops.placeholder(dtypes.string)}
+      output_1 = constant_op.constant([1.])
+      export_outputs = {
+          signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+              export_output.PredictOutput(outputs=output_1),
+          "train":
+              export_output.TrainOutput(loss=output_1),
+      }
+
+      signature_defs = export_utils.build_all_signature_defs(
+          receiver_tensor, export_outputs)
+
+      expected_signature_defs = {
+          "serving_default":
+              signature_def_utils.predict_signature_def(receiver_tensor,
+                                                        {"output": output_1})
+      }
+
+      self.assertDictEqual(expected_signature_defs, signature_defs)
+
+      signature_defs = export_utils.build_all_signature_defs(
+          receiver_tensor, export_outputs, serving_only=False)
+
+      expected_signature_defs.update({
+          "train":
+              signature_def_utils.supervised_train_signature_def(
+                  receiver_tensor, loss={"loss": output_1})
+      })
+
+      self.assertDictEqual(expected_signature_defs, signature_defs)
+
+  def test_export_outputs_for_mode(self):
+    predictions = {"predictions": constant_op.constant([1.])}
+    loss = {"loss": constant_op.constant([2.])}
+    metrics = {
+        "metrics": (constant_op.constant([3.]), constant_op.constant([4.]))}
+    expected_metrics = {
+        "metrics/value": metrics["metrics"][0],
+        "metrics/update_op": metrics["metrics"][1]
+    }
+
+    def _build_export_output(mode):
+      return export_utils.export_outputs_for_mode(
+          mode, None, predictions, loss, metrics)
+
+    ret = _build_export_output(KerasModeKeys.TRAIN)
+    self.assertIn(signature_constants.DEFAULT_TRAIN_SIGNATURE_DEF_KEY, ret)
+    export_out = ret[signature_constants.DEFAULT_TRAIN_SIGNATURE_DEF_KEY]
+    self.assertIsInstance(export_out, export_output.TrainOutput)
+    self.assertEqual(export_out.predictions, predictions)
+    self.assertEqual(export_out.loss, loss)
+    self.assertEqual(export_out.metrics, expected_metrics)
+
+    ret = _build_export_output(KerasModeKeys.TEST)
+    self.assertIn(signature_constants.DEFAULT_EVAL_SIGNATURE_DEF_KEY, ret)
+    export_out = ret[signature_constants.DEFAULT_EVAL_SIGNATURE_DEF_KEY]
+    self.assertIsInstance(export_out, export_output.EvalOutput)
+    self.assertEqual(export_out.predictions, predictions)
+    self.assertEqual(export_out.loss, loss)
+    self.assertEqual(export_out.metrics, expected_metrics)
+
+    ret = _build_export_output(KerasModeKeys.PREDICT)
+    self.assertIn(signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY, ret)
+    export_out = ret[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+    self.assertIsInstance(export_out, export_output.PredictOutput)
+    self.assertEqual(export_out.outputs, predictions)
+
+    classes = constant_op.constant(["class5"])
+    ret = export_utils.export_outputs_for_mode(
+        KerasModeKeys.PREDICT,
+        {"classify": export_output.ClassificationOutput(
+            classes=classes)})
+    self.assertIn("classify", ret)
+    export_out = ret["classify"]
+    self.assertIsInstance(export_out, export_output.ClassificationOutput)
+    self.assertEqual(export_out.classes, classes)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/model_utils/export_utils.py

@@ -0,0 +1,409 @@
+# 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.
+# ==============================================================================
+"""Utilities for creating SavedModels."""
+
+import collections
+import os
+import time
+
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.ops import op_selector
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import tag_constants
+from tensorflow.python.saved_model import utils
+from tensorflow.python.saved_model.model_utils import export_output as export_output_lib
+from tensorflow.python.saved_model.model_utils import mode_keys
+from tensorflow.python.saved_model.model_utils.mode_keys import KerasModeKeys as ModeKeys
+from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+from tensorflow.python.util import object_identity
+
+
+# Mapping of the modes to appropriate MetaGraph tags in the SavedModel.
+EXPORT_TAG_MAP = mode_keys.ModeKeyMap(**{
+    ModeKeys.PREDICT: [tag_constants.SERVING],
+    ModeKeys.TRAIN: [tag_constants.TRAINING],
+    ModeKeys.TEST: [tag_constants.EVAL]})
+
+# For every exported mode, a SignatureDef map should be created using the
+# functions `export_outputs_for_mode` and `build_all_signature_defs`. By
+# default, this map will contain a single Signature that defines the input
+# tensors and output predictions, losses, and/or metrics (depending on the mode)
+# The default keys used in the SignatureDef map are defined below.
+SIGNATURE_KEY_MAP = mode_keys.ModeKeyMap(**{
+    ModeKeys.PREDICT: signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
+    ModeKeys.TRAIN: signature_constants.DEFAULT_TRAIN_SIGNATURE_DEF_KEY,
+    ModeKeys.TEST: signature_constants.DEFAULT_EVAL_SIGNATURE_DEF_KEY})
+
+# Default names used in the SignatureDef input map, which maps strings to
+# TensorInfo protos.
+SINGLE_FEATURE_DEFAULT_NAME = 'feature'
+SINGLE_RECEIVER_DEFAULT_NAME = 'input'
+SINGLE_LABEL_DEFAULT_NAME = 'label'
+
+### Below utilities are specific to SavedModel exports.
+
+
+def _must_be_fed(op):
+  return op.type == 'Placeholder'
+
+
+def _ensure_servable(input_tensors, names_to_output_tensor_infos):
+  """Check that the signature outputs don't depend on unreachable placeholders.
+
+  Args:
+    input_tensors: An iterable of `Tensor`s specified as the signature's inputs.
+    names_to_output_tensor_infos: An mapping from output names to respective
+      `TensorInfo`s corresponding to the signature's output tensors.
+
+  Raises:
+    ValueError: If any of the signature's outputs depend on placeholders not
+      provided as signature's inputs.
+  """
+  plain_input_tensors = nest.flatten(input_tensors, expand_composites=True)
+
+  graph = op_selector.get_unique_graph(plain_input_tensors)
+
+  output_tensors = [
+      utils.get_tensor_from_tensor_info(tensor, graph=graph)
+      for tensor in names_to_output_tensor_infos.values()
+  ]
+  plain_output_tensors = nest.flatten(output_tensors, expand_composites=True)
+
+  dependency_ops = op_selector.get_backward_walk_ops(
+      plain_output_tensors, stop_at_ts=plain_input_tensors)
+
+  fed_tensors = object_identity.ObjectIdentitySet(plain_input_tensors)
+  for dependency_op in dependency_ops:
+    if _must_be_fed(dependency_op) and (not all(
+        output in fed_tensors for output in dependency_op.outputs)):
+      input_tensor_names = [tensor.name for tensor in plain_input_tensors]
+      output_tensor_keys = list(names_to_output_tensor_infos.keys())
+      output_tensor_names = [tensor.name for tensor in plain_output_tensors]
+      dependency_path = op_selector.show_path(dependency_op,
+                                              plain_output_tensors,
+                                              plain_input_tensors)
+      raise ValueError(
+          f'The signature\'s input tensors {input_tensor_names} are '
+          f'insufficient to compute its output keys {output_tensor_keys} '
+          f'(respectively, tensors {output_tensor_names}) because of the '
+          f'dependency on `{dependency_op.name}` which is not given as '
+          'a signature input, as illustrated by the following dependency path: '
+          f'{dependency_path}')
+
+
+def build_all_signature_defs(receiver_tensors,
+                             export_outputs,
+                             receiver_tensors_alternatives=None,
+                             serving_only=True):
+  """Build `SignatureDef`s for all export outputs.
+
+  Args:
+    receiver_tensors: a `Tensor`, or a dict of string to `Tensor`, specifying
+      input nodes where this receiver expects to be fed by default.  Typically,
+      this is a single placeholder expecting serialized `tf.Example` protos.
+    export_outputs: a dict of ExportOutput instances, each of which has
+      an as_signature_def instance method that will be called to retrieve
+      the signature_def for all export output tensors.
+    receiver_tensors_alternatives: a dict of string to additional
+      groups of receiver tensors, each of which may be a `Tensor` or a dict of
+      string to `Tensor`.  These named receiver tensor alternatives generate
+      additional serving signatures, which may be used to feed inputs at
+      different points within the input receiver subgraph.  A typical usage is
+      to allow feeding raw feature `Tensor`s *downstream* of the
+      tf.io.parse_example() op.  Defaults to None.
+    serving_only: boolean; if true, resulting signature defs will only include
+      valid serving signatures. If false, all requested signatures will be
+      returned.
+
+  Returns:
+    signature_def representing all passed args.
+
+  Raises:
+    ValueError: if export_outputs is not a dict
+  """
+  if not isinstance(receiver_tensors, dict):
+    receiver_tensors = {SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors}
+  if export_outputs is None or not isinstance(export_outputs, dict):
+    raise ValueError('`export_outputs` must be a dict. Received '
+                     f'{export_outputs} with type '
+                     f'{type(export_outputs).__name__}.')
+
+  signature_def_map = {}
+  excluded_signatures = {}
+  input_tensors = receiver_tensors.values()
+  for output_key, export_output in export_outputs.items():
+    signature_name = '{}'.format(output_key or 'None')
+    try:
+      signature = export_output.as_signature_def(receiver_tensors)
+      _ensure_servable(input_tensors, signature.outputs)
+      signature_def_map[signature_name] = signature
+    except ValueError as e:
+      excluded_signatures[signature_name] = str(e)
+
+  if receiver_tensors_alternatives:
+    for receiver_name, receiver_tensors_alt in (
+        receiver_tensors_alternatives.items()):
+      if not isinstance(receiver_tensors_alt, dict):
+        receiver_tensors_alt = {
+            SINGLE_RECEIVER_DEFAULT_NAME: receiver_tensors_alt
+        }
+      alt_input_tensors = receiver_tensors_alt.values()
+      for output_key, export_output in export_outputs.items():
+        signature_name = '{}:{}'.format(receiver_name or 'None', output_key or
+                                        'None')
+        try:
+          signature = export_output.as_signature_def(receiver_tensors_alt)
+          _ensure_servable(alt_input_tensors, signature.outputs)
+          signature_def_map[signature_name] = signature
+        except ValueError as e:
+          excluded_signatures[signature_name] = str(e)
+
+  _log_signature_report(signature_def_map, excluded_signatures)
+
+  # The above calls to export_output_lib.as_signature_def should return only
+  # valid signatures; if there is a validity problem, they raise a ValueError,
+  # in which case we exclude that signature from signature_def_map above.
+  # The is_valid_signature check ensures that the signatures produced are
+  # valid for serving, and acts as an additional sanity check for export
+  # signatures produced for serving. We skip this check for training and eval
+  # signatures, which are not intended for serving.
+  if serving_only:
+    signature_def_map = {
+        k: v
+        for k, v in signature_def_map.items()
+        if signature_def_utils.is_valid_signature(v)
+    }
+  return signature_def_map
+
+
+_FRIENDLY_METHOD_NAMES = {
+    signature_constants.CLASSIFY_METHOD_NAME: 'Classify',
+    signature_constants.REGRESS_METHOD_NAME: 'Regress',
+    signature_constants.PREDICT_METHOD_NAME: 'Predict',
+    signature_constants.SUPERVISED_TRAIN_METHOD_NAME: 'Train',
+    signature_constants.SUPERVISED_EVAL_METHOD_NAME: 'Eval',
+}
+
+
+def _log_signature_report(signature_def_map, excluded_signatures):
+  """Log a report of which signatures were produced."""
+  sig_names_by_method_name = collections.defaultdict(list)
+
+  # We'll collect whatever method_names are present, but also we want to make
+  # sure to output a line for each of the three standard methods even if they
+  # have no signatures.
+  for method_name in _FRIENDLY_METHOD_NAMES:
+    sig_names_by_method_name[method_name] = []
+
+  for signature_name, sig in signature_def_map.items():
+    sig_names_by_method_name[sig.method_name].append(signature_name)
+
+  # TODO(b/67733540): consider printing the full signatures, not just names
+  for method_name, sig_names in sig_names_by_method_name.items():
+    if method_name in _FRIENDLY_METHOD_NAMES:
+      method_name = _FRIENDLY_METHOD_NAMES[method_name]
+    logging.info('Signatures INCLUDED in export for {}: {}'.format(
+        method_name, sig_names if sig_names else 'None'))
+
+  if excluded_signatures:
+    logging.info('Signatures EXCLUDED from export because they cannot be '
+                 'be served via TensorFlow Serving APIs:')
+    for signature_name, message in excluded_signatures.items():
+      logging.info('\'{}\' : {}'.format(signature_name, message))
+
+  if not signature_def_map:
+    logging.warn('Export includes no signatures!')
+  elif (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY not in
+        signature_def_map):
+    logging.warn('Export includes no default signature!')
+
+
+# When we create a timestamped directory, there is a small chance that the
+# directory already exists because another process is also creating these
+# directories. In this case we just wait one second to get a new timestamp and
+# try again. If this fails several times in a row, then something is seriously
+# wrong.
+MAX_DIRECTORY_CREATION_ATTEMPTS = 10
+
+
+def get_timestamped_export_dir(export_dir_base):
+  """Builds a path to a new subdirectory within the base directory.
+
+  Each export is written into a new subdirectory named using the
+  current time.  This guarantees monotonically increasing version
+  numbers even across multiple runs of the pipeline.
+  The timestamp used is the number of seconds since epoch UTC.
+
+  Args:
+    export_dir_base: A string containing a directory to write the exported
+        graph and checkpoints.
+  Returns:
+    The full path of the new subdirectory (which is not actually created yet).
+
+  Raises:
+    RuntimeError: if repeated attempts fail to obtain a unique timestamped
+      directory name.
+  """
+  attempts = 0
+  while attempts < MAX_DIRECTORY_CREATION_ATTEMPTS:
+    timestamp = int(time.time())
+
+    result_dir = file_io.join(
+        compat.as_bytes(export_dir_base), compat.as_bytes(str(timestamp)))
+    if not gfile.Exists(result_dir):
+      # Collisions are still possible (though extremely unlikely): this
+      # directory is not actually created yet, but it will be almost
+      # instantly on return from this function.
+      return result_dir
+    time.sleep(1)
+    attempts += 1
+    logging.warn('Directory {} already exists; retrying (attempt {}/{})'.format(
+        compat.as_str(result_dir), attempts, MAX_DIRECTORY_CREATION_ATTEMPTS))
+  raise RuntimeError('Failed to obtain a unique export directory name after '
+                     f'{MAX_DIRECTORY_CREATION_ATTEMPTS} attempts.')
+
+
+def get_temp_export_dir(timestamped_export_dir):
+  """Builds a directory name based on the argument but starting with 'temp-'.
+
+  This relies on the fact that TensorFlow Serving ignores subdirectories of
+  the base directory that can't be parsed as integers.
+
+  Args:
+    timestamped_export_dir: the name of the eventual export directory, e.g.
+      /foo/bar/<timestamp>
+
+  Returns:
+    A sister directory prefixed with 'temp-', e.g. /foo/bar/temp-<timestamp>.
+  """
+  (dirname, basename) = os.path.split(timestamped_export_dir)
+  if isinstance(basename, bytes):
+    str_name = basename.decode('utf-8')
+  else:
+    str_name = str(basename)
+  temp_export_dir = file_io.join(
+      compat.as_bytes(dirname), compat.as_bytes('temp-{}'.format(str_name)))
+  return temp_export_dir
+
+
+def export_outputs_for_mode(
+    mode, serving_export_outputs=None, predictions=None, loss=None,
+    metrics=None):
+  """Util function for constructing a `ExportOutput` dict given a mode.
+
+  The returned dict can be directly passed to `build_all_signature_defs` helper
+  function as the `export_outputs` argument, used for generating a SignatureDef
+  map.
+
+  Args:
+    mode: A `ModeKeys` specifying the mode.
+    serving_export_outputs: Describes the output signatures to be exported to
+      `SavedModel` and used during serving. Should be a dict or None.
+    predictions: A dict of Tensors or single Tensor representing model
+        predictions. This argument is only used if serving_export_outputs is not
+        set.
+    loss: A dict of Tensors or single Tensor representing calculated loss.
+    metrics: A dict of (metric_value, update_op) tuples, or a single tuple.
+      metric_value must be a Tensor, and update_op must be a Tensor or Op
+
+  Returns:
+    Dictionary mapping the key to an `ExportOutput` object.
+    The key is the expected SignatureDef key for the mode.
+
+  Raises:
+    ValueError: if an appropriate ExportOutput cannot be found for the mode.
+  """
+  if mode not in SIGNATURE_KEY_MAP:
+    raise ValueError(
+        f'Export output type not found for `mode`: {mode}. Expected one of: '
+        f'{list(SIGNATURE_KEY_MAP.keys())}.')
+  signature_key = SIGNATURE_KEY_MAP[mode]
+  if mode_keys.is_predict(mode):
+    return get_export_outputs(serving_export_outputs, predictions)
+  elif mode_keys.is_train(mode):
+    return {signature_key: export_output_lib.TrainOutput(
+        loss=loss, predictions=predictions, metrics=metrics)}
+  else:
+    return {signature_key: export_output_lib.EvalOutput(
+        loss=loss, predictions=predictions, metrics=metrics)}
+
+
+def get_export_outputs(export_outputs, predictions):
+  """Validate export_outputs or create default export_outputs.
+
+  Args:
+    export_outputs: Describes the output signatures to be exported to
+      `SavedModel` and used during serving. Should be a dict or None.
+    predictions:  Predictions `Tensor` or dict of `Tensor`.
+
+  Returns:
+    Valid export_outputs dict
+
+  Raises:
+    TypeError: if export_outputs is not a dict or its values are not
+      ExportOutput instances.
+  """
+  if export_outputs is None:
+    default_output = export_output_lib.PredictOutput(predictions)
+    export_outputs = {
+        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: default_output}
+
+  if not isinstance(export_outputs, dict):
+    raise TypeError(
+        f'`export_outputs` must be dict, received: {export_outputs}.')
+  for v in export_outputs.values():
+    if not isinstance(v, export_output_lib.ExportOutput):
+      raise TypeError(
+          'Values in `export_outputs` must be ExportOutput objects, '
+          f'received: {export_outputs}.')
+
+  _maybe_add_default_serving_output(export_outputs)
+
+  return export_outputs
+
+
+def _maybe_add_default_serving_output(export_outputs):
+  """Add a default serving output to the export_outputs if not present.
+
+  Args:
+    export_outputs: Describes the output signatures to be exported to
+      `SavedModel` and used during serving. Should be a dict.
+
+  Returns:
+    export_outputs dict with default serving signature added if necessary
+
+  Raises:
+    ValueError: if multiple export_outputs were provided without a default
+      serving key.
+  """
+  if len(export_outputs) == 1:
+    (key, value), = export_outputs.items()
+    if key != signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+      export_outputs[
+          signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = value
+  if len(export_outputs) > 1:
+    if (signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
+        not in export_outputs):
+      raise ValueError(
+          'Multiple `export_outputs` were provided, but none of them are '
+          'specified as the default. Use'
+          '`tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY` to '
+          'specify a default.')
+
+  return export_outputs

tensorflow/python/saved_model/model_utils/mode_keys.py

@@ -0,0 +1,107 @@
+# Copyright 2016 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.
+# ==============================================================================
+# LINT.IfChange
+"""Utils for managing different mode strings used by Keras and Estimator models.
+"""
+
+from tensorflow.python.util.compat import collections_abc
+
+
+class KerasModeKeys(object):
+  """Standard names for model modes.
+
+  The following standard keys are defined:
+
+  * `TRAIN`: training/fitting mode.
+  * `TEST`: testing/evaluation mode.
+  * `PREDICT`: prediction/inference mode.
+  """
+
+  TRAIN = 'train'
+  TEST = 'test'
+  PREDICT = 'predict'
+
+
+# TODO(kathywu): Remove copy in Estimator after nightlies
+class EstimatorModeKeys(object):
+  """Standard names for Estimator model modes.
+
+  The following standard keys are defined:
+
+  * `TRAIN`: training/fitting mode.
+  * `EVAL`: testing/evaluation mode.
+  * `PREDICT`: predication/inference mode.
+  """
+
+  TRAIN = 'train'
+  EVAL = 'eval'
+  PREDICT = 'infer'
+
+
+def is_predict(mode):
+  return mode in [KerasModeKeys.PREDICT, EstimatorModeKeys.PREDICT]
+
+
+def is_eval(mode):
+  return mode in [KerasModeKeys.TEST, EstimatorModeKeys.EVAL]
+
+
+def is_train(mode):
+  return mode in [KerasModeKeys.TRAIN, EstimatorModeKeys.TRAIN]
+
+
+class ModeKeyMap(collections_abc.Mapping):
+  """Map using ModeKeys as keys.
+
+  This class creates an immutable mapping from modes to values. For example,
+  SavedModel export of Keras and Estimator models use this to map modes to their
+  corresponding MetaGraph tags/SignatureDef keys.
+
+  Since this class uses modes, rather than strings, as keys, both "predict"
+  (Keras's PREDICT ModeKey) and "infer" (Estimator's PREDICT ModeKey) map to the
+  same value.
+  """
+
+  def __init__(self, **kwargs):
+    self._internal_dict = {}
+    self._keys = []
+    for key in kwargs:
+      self._keys.append(key)
+      dict_key = self._get_internal_key(key)
+      if dict_key in self._internal_dict:
+        raise ValueError(
+            'Error creating ModeKeyMap. Multiple keys/values found for {} mode.'
+            .format(dict_key))
+      self._internal_dict[dict_key] = kwargs[key]
+
+  def _get_internal_key(self, key):
+    """Return keys used for the internal dictionary."""
+    if is_train(key):
+      return KerasModeKeys.TRAIN
+    if is_eval(key):
+      return KerasModeKeys.TEST
+    if is_predict(key):
+      return KerasModeKeys.PREDICT
+    raise ValueError('Invalid mode key: {}.'.format(key))
+
+  def __getitem__(self, key):
+    return self._internal_dict[self._get_internal_key(key)]
+
+  def __iter__(self):
+    return iter(self._keys)
+
+  def __len__(self):
+    return len(self._keys)
+# LINT.ThenChange(//tensorflow/python/keras/saving/utils_v1/mode_keys.py)

tensorflow/python/saved_model/model_utils/mode_keys_test.py

@@ -0,0 +1,61 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""ModeKey Tests."""
+
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model.model_utils import mode_keys
+
+
+class ModeKeyMapTest(test.TestCase):
+
+  def test_map(self):
+    mode_map = mode_keys.ModeKeyMap(**{
+        mode_keys.KerasModeKeys.PREDICT: 3,
+        mode_keys.KerasModeKeys.TEST: 1
+    })
+
+    # Test dictionary __getitem__
+    self.assertEqual(3, mode_map[mode_keys.KerasModeKeys.PREDICT])
+    self.assertEqual(3, mode_map[mode_keys.EstimatorModeKeys.PREDICT])
+    self.assertEqual(1, mode_map[mode_keys.KerasModeKeys.TEST])
+    self.assertEqual(1, mode_map[mode_keys.EstimatorModeKeys.EVAL])
+    with self.assertRaises(KeyError):
+      _ = mode_map[mode_keys.KerasModeKeys.TRAIN]
+    with self.assertRaises(KeyError):
+      _ = mode_map[mode_keys.EstimatorModeKeys.TRAIN]
+    with self.assertRaisesRegex(ValueError, 'Invalid mode'):
+      _ = mode_map['serve']
+
+    # Test common dictionary methods
+    self.assertLen(mode_map, 2)
+    self.assertEqual({1, 3}, set(mode_map.values()))
+    self.assertEqual(
+        {mode_keys.KerasModeKeys.TEST, mode_keys.KerasModeKeys.PREDICT},
+        set(mode_map.keys()))
+
+    # Map is immutable
+    with self.assertRaises(TypeError):
+      mode_map[mode_keys.KerasModeKeys.TEST] = 1  # pylint: disable=unsupported-assignment-operation
+
+  def test_invalid_init(self):
+    with self.assertRaisesRegex(ValueError, 'Multiple keys/values found'):
+      _ = mode_keys.ModeKeyMap(**{
+          mode_keys.KerasModeKeys.PREDICT: 3,
+          mode_keys.EstimatorModeKeys.PREDICT: 1
+      })
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/saved_model/nested_structure_coder.py

@@ -0,0 +1,514 @@
+# 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.
+# ==============================================================================
+"""Module that encodes (decodes) nested structures into (from) protos.
+
+The intended use is to serialize everything needed to restore a `Function` that
+was saved into a SavedModel. This may include concrete function inputs and
+outputs, signatures, function specs, etc.
+
+Example use:
+# Encode into proto.
+signature_proto = nested_structure_coder.encode_structure(
+    function.input_signature)
+# Decode into a Python object.
+restored_signature = nested_structure_coder.decode_proto(signature_proto)
+"""
+
+import collections
+import functools
+import warnings
+
+from tensorflow.core.protobuf import struct_pb2
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import type_spec_registry
+from tensorflow.python.types import internal
+from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+from tensorflow.python.util.compat import collections_abc
+from tensorflow.python.util.tf_export import tf_export
+
+
+class NotEncodableError(Exception):
+  """Error raised when a coder cannot encode an object."""
+
+
+def register_codec(x):
+  """Registers a codec to use for encoding/decoding.
+
+  Args:
+    x: The codec object to register. The object must implement can_encode,
+      do_encode, can_decode, and do_decode. See the various _*Codec classes for
+      examples.
+  """
+  _codecs.append(x)
+
+
+def _get_encoders():
+  return [(c.can_encode, c.do_encode) for c in _codecs]
+
+
+def _get_decoders():
+  return [(c.can_decode, c.do_decode) for c in _codecs]
+
+
+def _map_structure(pyobj, coders):
+  # Iterate through the codecs in the reverse order they were registered in,
+  #   as the most specific codec should be checked first.
+  for can, do in reversed(coders):
+    if can(pyobj):
+      recursion_fn = functools.partial(_map_structure, coders=coders)
+      return do(pyobj, recursion_fn)
+  raise NotEncodableError(
+      f"No encoder for object {str(pyobj)} of type {type(pyobj)}.")
+
+
+@tf_export("__internal__.saved_model.encode_structure", v1=[])
+def encode_structure(nested_structure):
+  """Encodes nested structures composed of encodable types into a proto.
+
+  Args:
+    nested_structure: Structure to encode.
+
+  Returns:
+    Encoded proto.
+
+  Raises:
+    NotEncodableError: For values for which there are no encoders.
+  """
+  return _map_structure(nested_structure, _get_encoders())
+
+
+def can_encode(nested_structure):
+  """Determines whether a nested structure can be encoded into a proto.
+
+  Args:
+    nested_structure: Structure to encode.
+
+  Returns:
+    True if the nested structured can be encoded.
+  """
+  try:
+    encode_structure(nested_structure)
+  except NotEncodableError:
+    return False
+  return True
+
+
+@tf_export("__internal__.saved_model.decode_proto", v1=[])
+def decode_proto(proto):
+  """Decodes proto representing a nested structure.
+
+  Args:
+    proto: Proto to decode.
+
+  Returns:
+    Decoded structure.
+
+  Raises:
+    NotEncodableError: For values for which there are no encoders.
+  """
+  return _map_structure(proto, _get_decoders())
+
+
+class _ListCodec:
+  """Codec for lists."""
+
+  def can_encode(self, pyobj):
+    return isinstance(pyobj, list)
+
+  def do_encode(self, list_value, encode_fn):
+    encoded_list = struct_pb2.StructuredValue()
+    encoded_list.list_value.CopyFrom(struct_pb2.ListValue())
+    for element in list_value:
+      encoded_list.list_value.values.add().CopyFrom(encode_fn(element))
+    return encoded_list
+
+  def can_decode(self, value):
+    return value.HasField("list_value")
+
+  def do_decode(self, value, decode_fn):
+    return [decode_fn(element) for element in value.list_value.values]
+
+
+def _is_tuple(obj):
+  return not _is_named_tuple(obj) and isinstance(obj, tuple)
+
+
+def _is_named_tuple(instance):
+  """Returns True iff `instance` is a `namedtuple`.
+
+  Args:
+    instance: An instance of a Python object.
+
+  Returns:
+    True if `instance` is a `namedtuple`.
+  """
+  if not isinstance(instance, tuple):
+    return False
+  return (hasattr(instance, "_fields") and
+          isinstance(instance._fields, collections_abc.Sequence) and
+          all(isinstance(f, str) for f in instance._fields))
+
+
+class _TupleCodec:
+  """Codec for tuples."""
+
+  def can_encode(self, pyobj):
+    return _is_tuple(pyobj)
+
+  def do_encode(self, tuple_value, encode_fn):
+    encoded_tuple = struct_pb2.StructuredValue()
+    encoded_tuple.tuple_value.CopyFrom(struct_pb2.TupleValue())
+    for element in tuple_value:
+      encoded_tuple.tuple_value.values.add().CopyFrom(encode_fn(element))
+    return encoded_tuple
+
+  def can_decode(self, value):
+    return value.HasField("tuple_value")
+
+  def do_decode(self, value, decode_fn):
+    return tuple(decode_fn(element) for element in value.tuple_value.values)
+
+
+class _DictCodec:
+  """Codec for dicts."""
+
+  def can_encode(self, pyobj):
+    return isinstance(pyobj, collections_abc.Mapping)
+
+  def do_encode(self, dict_value, encode_fn):
+    encoded_dict = struct_pb2.StructuredValue()
+    encoded_dict.dict_value.CopyFrom(struct_pb2.DictValue())
+    for key, value in dict_value.items():
+      encoded_dict.dict_value.fields[key].CopyFrom(encode_fn(value))
+    return encoded_dict
+
+  def can_decode(self, value):
+    return value.HasField("dict_value")
+
+  def do_decode(self, value, decode_fn):
+    return {key: decode_fn(val) for key, val in value.dict_value.fields.items()}
+
+
+class _NamedTupleCodec:
+  """Codec for namedtuples.
+
+  Encoding and decoding a namedtuple reconstructs a namedtuple with a different
+  actual Python type, but with the same `typename` and `fields`.
+  """
+
+  def can_encode(self, pyobj):
+    return _is_named_tuple(pyobj)
+
+  def do_encode(self, named_tuple_value, encode_fn):
+    encoded_named_tuple = struct_pb2.StructuredValue()
+    encoded_named_tuple.named_tuple_value.CopyFrom(struct_pb2.NamedTupleValue())
+    encoded_named_tuple.named_tuple_value.name = \
+      named_tuple_value.__class__.__name__
+    for key in named_tuple_value._fields:
+      pair = encoded_named_tuple.named_tuple_value.values.add()
+      pair.key = key
+      pair.value.CopyFrom(encode_fn(named_tuple_value._asdict()[key]))
+    return encoded_named_tuple
+
+  def can_decode(self, value):
+    return value.HasField("named_tuple_value")
+
+  def do_decode(self, value, decode_fn):
+    key_value_pairs = value.named_tuple_value.values
+    items = [(pair.key, decode_fn(pair.value)) for pair in key_value_pairs]
+    named_tuple_type = collections.namedtuple(value.named_tuple_value.name,
+                                              [item[0] for item in items])
+    return named_tuple_type(**dict(items))
+
+
+class _Float64Codec:
+  """Codec for floats."""
+
+  def can_encode(self, pyobj):
+    return isinstance(pyobj, float)
+
+  def do_encode(self, float64_value, encode_fn):
+    del encode_fn
+    value = struct_pb2.StructuredValue()
+    value.float64_value = float64_value
+    return value
+
+  def can_decode(self, value):
+    return value.HasField("float64_value")
+
+  def do_decode(self, value, decode_fn):
+    del decode_fn
+    return value.float64_value
+
+
+class _Int64Codec:
+  """Codec for Python integers (limited to 64 bit values)."""
+
+  def can_encode(self, pyobj):
+    return not isinstance(pyobj, bool) and isinstance(pyobj, int)
+
+  def do_encode(self, int_value, encode_fn):
+    del encode_fn
+    value = struct_pb2.StructuredValue()
+    value.int64_value = int_value
+    return value
+
+  def can_decode(self, value):
+    return value.HasField("int64_value")
+
+  def do_decode(self, value, decode_fn):
+    del decode_fn
+    return int(value.int64_value)
+
+
+class _StringCodec:
+  """Codec for strings.
+
+  See StructuredValue.string_value in proto/struct.proto for more detailed
+  explanation.
+  """
+
+  def can_encode(self, pyobj):
+    return isinstance(pyobj, str)
+
+  def do_encode(self, string_value, encode_fn):
+    del encode_fn
+    value = struct_pb2.StructuredValue()
+    value.string_value = string_value
+    return value
+
+  def can_decode(self, value):
+    return value.HasField("string_value")
+
+  def do_decode(self, value, decode_fn):
+    del decode_fn
+    return compat.as_str(value.string_value)
+
+
+class _NoneCodec:
+  """Codec for None."""
+
+  def can_encode(self, pyobj):
+    return pyobj is None
+
+  def do_encode(self, none_value, encode_fn):
+    del encode_fn, none_value
+    value = struct_pb2.StructuredValue()
+    value.none_value.CopyFrom(struct_pb2.NoneValue())
+    return value
+
+  def can_decode(self, value):
+    return value.HasField("none_value")
+
+  def do_decode(self, value, decode_fn):
+    del decode_fn, value
+    return None
+
+
+class _BoolCodec:
+  """Codec for booleans."""
+
+  def can_encode(self, pyobj):
+    return isinstance(pyobj, bool)
+
+  def do_encode(self, bool_value, encode_fn):
+    del encode_fn
+    value = struct_pb2.StructuredValue()
+    value.bool_value = bool_value
+    return value
+
+  def can_decode(self, value):
+    return value.HasField("bool_value")
+
+  def do_decode(self, value, decode_fn):
+    del decode_fn
+    return value.bool_value
+
+
+class _TensorTypeCodec:
+  """Codec for `TensorType`."""
+
+  def can_encode(self, pyobj):
+    return isinstance(pyobj, dtypes.DType)
+
+  def do_encode(self, tensor_dtype_value, encode_fn):
+    del encode_fn
+    encoded_tensor_type = struct_pb2.StructuredValue()
+    encoded_tensor_type.tensor_dtype_value = tensor_dtype_value.as_datatype_enum
+    return encoded_tensor_type
+
+  def can_decode(self, value):
+    return value.HasField("tensor_dtype_value")
+
+  def do_decode(self, value, decode_fn):
+    del decode_fn
+    return dtypes.DType(value.tensor_dtype_value)
+
+
+class BuiltInTypeSpecCodec:
+  """Codec for built-in `TypeSpec` classes.
+
+  Built-in TypeSpec's that do not require a custom codec implementation
+  register themselves by instantiating this class and passing it to
+  register_codec.
+
+  Attributes:
+    type_spec_class: The built-in TypeSpec class that the
+      codec is instantiated for.
+    type_spec_proto_enum: The proto enum value for the built-in TypeSpec class.
+  """
+
+  _BUILT_IN_TYPE_SPEC_PROTOS = []
+  _BUILT_IN_TYPE_SPECS = []
+
+  def __init__(self, type_spec_class, type_spec_proto_enum):
+    if not issubclass(type_spec_class, internal.TypeSpec):
+      raise ValueError(
+          f"The type '{type_spec_class}' does not subclass tf.TypeSpec.")
+
+    if type_spec_class in self._BUILT_IN_TYPE_SPECS:
+      raise ValueError(
+          f"The type '{type_spec_class}' already has an instantiated codec.")
+
+    if type_spec_proto_enum in self._BUILT_IN_TYPE_SPEC_PROTOS:
+      raise ValueError(
+          f"The proto value '{type_spec_proto_enum}' is already registered."
+      )
+
+    if (not isinstance(type_spec_proto_enum, int)
+        or type_spec_proto_enum <= 0
+        or type_spec_proto_enum > 10):
+      raise ValueError(f"The proto value '{type_spec_proto_enum}' is invalid.")
+
+    self.type_spec_class = type_spec_class
+    self.type_spec_proto_enum = type_spec_proto_enum
+
+    self._BUILT_IN_TYPE_SPECS.append(type_spec_class)
+    self._BUILT_IN_TYPE_SPEC_PROTOS.append(type_spec_proto_enum)
+
+  def can_encode(self, pyobj):
+    """Returns true if `pyobj` can be encoded as the built-in TypeSpec."""
+    return isinstance(pyobj, self.type_spec_class)
+
+  def do_encode(self, type_spec_value, encode_fn):
+    """Returns an encoded proto for the given built-in TypeSpec."""
+    type_state = type_spec_value._serialize()  # pylint: disable=protected-access
+    num_flat_components = len(nest.flatten(
+        type_spec_value._component_specs, expand_composites=True))  # pylint: disable=protected-access
+    encoded_type_spec = struct_pb2.StructuredValue()
+    encoded_type_spec.type_spec_value.CopyFrom(
+        struct_pb2.TypeSpecProto(
+            type_spec_class=self.type_spec_proto_enum,
+            type_state=encode_fn(type_state),
+            type_spec_class_name=self.type_spec_class.__name__,
+            num_flat_components=num_flat_components))
+    return encoded_type_spec
+
+  def can_decode(self, value):
+    """Returns true if `value` can be decoded into its built-in TypeSpec."""
+    if value.HasField("type_spec_value"):
+      type_spec_class_enum = value.type_spec_value.type_spec_class
+      return type_spec_class_enum == self.type_spec_proto_enum
+    return False
+
+  def do_decode(self, value, decode_fn):
+    """Returns the built in `TypeSpec` encoded by the proto `value`."""
+    type_spec_proto = value.type_spec_value
+    # pylint: disable=protected-access
+    return self.type_spec_class._deserialize(
+        decode_fn(type_spec_proto.type_state)
+    )
+
+
+# TODO(b/238903802): Use TraceType serialization and specific protos.
+class _TypeSpecCodec:
+  """Codec for `tf.TypeSpec`."""
+
+  def can_encode(self, pyobj):
+    """Returns true if `pyobj` can be encoded as a TypeSpec."""
+    # Check if it's a registered type.
+    if isinstance(pyobj, internal.TypeSpec):
+      try:
+        type_spec_registry.get_name(type(pyobj))
+        return True
+      except ValueError:
+        return False
+
+    return False
+
+  def do_encode(self, type_spec_value, encode_fn):
+    """Returns an encoded proto for the given `tf.TypeSpec`."""
+    type_spec_class_name = type_spec_registry.get_name(type(type_spec_value))
+    type_spec_class = struct_pb2.TypeSpecProto.REGISTERED_TYPE_SPEC
+    # Support for saving registered TypeSpecs is currently experimental.
+    # Issue a warning to indicate the limitations.
+    warnings.warn("Encoding a StructuredValue with type %s; loading this "
+                  "StructuredValue will require that this type be "
+                  "imported and registered." % type_spec_class_name)
+
+    type_state = type_spec_value._serialize()  # pylint: disable=protected-access
+    num_flat_components = len(
+        nest.flatten(type_spec_value._component_specs, expand_composites=True))  # pylint: disable=protected-access
+    encoded_type_spec = struct_pb2.StructuredValue()
+    encoded_type_spec.type_spec_value.CopyFrom(
+        struct_pb2.TypeSpecProto(
+            type_spec_class=type_spec_class,
+            type_state=encode_fn(type_state),
+            type_spec_class_name=type_spec_class_name,
+            num_flat_components=num_flat_components))
+    return encoded_type_spec
+
+  def can_decode(self, value):
+    """Returns true if `value` can be decoded into a `tf.TypeSpec`."""
+    return value.HasField("type_spec_value")
+
+  def do_decode(self, value, decode_fn):
+    """Returns the `tf.TypeSpec` encoded by the proto `value`."""
+    type_spec_proto = value.type_spec_value
+    type_spec_class_enum = type_spec_proto.type_spec_class
+    class_name = type_spec_proto.type_spec_class_name
+
+    if type_spec_class_enum == struct_pb2.TypeSpecProto.REGISTERED_TYPE_SPEC:
+      try:
+        type_spec_class = type_spec_registry.lookup(class_name)
+      except ValueError as e:
+        raise ValueError(
+            f"The type '{class_name}' has not been registered.  It must be "
+            "registered before you load this object (typically by importing "
+            "its module).") from e
+    else:
+      raise ValueError(
+          f"The type '{class_name}' is not supported by this version of "
+          "TensorFlow. (The object you are loading must have been created "
+          "with a newer version of TensorFlow.)")
+
+    # pylint: disable=protected-access
+    return type_spec_class._deserialize(decode_fn(type_spec_proto.type_state))
+
+
+_codecs = [
+    _ListCodec(),
+    _TupleCodec(),
+    _NamedTupleCodec(),
+    _StringCodec(),
+    _Float64Codec(),
+    _NoneCodec(),
+    _Int64Codec(),
+    _BoolCodec(),
+    _DictCodec(),
+    _TypeSpecCodec(),
+    _TensorTypeCodec(),
+]

tensorflow/python/saved_model/nested_structure_coder_test.py

@@ -0,0 +1,546 @@
+# 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.
+# ==============================================================================
+"""Tests for nested structure coding."""
+
+import collections
+import typing
+import warnings
+
+import numpy as np
+
+from google.protobuf import text_format
+from tensorflow.core.protobuf import struct_pb2
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import extension_type
+from tensorflow.python.framework import immutable_dict
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework import type_spec
+from tensorflow.python.framework import type_spec_registry
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.types import internal
+
+
+class NestedStructureCoderTest(test.TestCase):
+
+  def testEncodeDecodeList(self):
+    structure = [1.5, 2.5, 3.0]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected.list_value.values.add().float64_value = 1.5
+    expected.list_value.values.add().float64_value = 2.5
+    expected.list_value.values.add().float64_value = 3.0
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeTuple(self):
+    structure = ("hello", [3, (2, 1)])
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected.tuple_value.values.add().string_value = "hello"
+    list_value = expected.tuple_value.values.add().list_value
+    list_value.values.add().int64_value = 3
+    tuple_value = list_value.values.add().tuple_value
+    tuple_value.values.add().int64_value = 2
+    tuple_value.values.add().int64_value = 1
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeDict(self):
+    structure = dict(a=3, b=[7, 2.5])
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected.dict_value.fields["a"].int64_value = 3
+    list_value = expected.dict_value.fields["b"].list_value
+    list_value.values.add().int64_value = 7
+    list_value.values.add().float64_value = 2.5
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertIsInstance(decoded["a"], int)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeImmutableDict(self):
+    structure = immutable_dict.ImmutableDict(dict(a=3, b=[7, 2.5]))
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected.dict_value.fields["a"].int64_value = 3
+    list_value = expected.dict_value.fields["b"].list_value
+    list_value.values.add().int64_value = 7
+    list_value.values.add().float64_value = 2.5
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertIsInstance(decoded["a"], int)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeTensorShape(self):
+    structure = [tensor_shape.TensorShape([1, 2, 3]), "hello"]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected_list = expected.list_value
+    expected_tensor_shape = expected_list.values.add().tensor_shape_value
+    expected_tensor_shape.dim.add().size = 1
+    expected_tensor_shape.dim.add().size = 2
+    expected_tensor_shape.dim.add().size = 3
+    expected_tensor_shape = expected_list.values.add().string_value = "hello"
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeNamedTuple(self):
+    named_tuple_type = collections.namedtuple("NamedTuple", ["x", "y"])
+    named_tuple = named_tuple_type(x=[1, 2], y="hello")
+    self.assertTrue(nested_structure_coder.can_encode(named_tuple))
+    encoded = nested_structure_coder.encode_structure(named_tuple)
+    expected = struct_pb2.StructuredValue()
+    expected_named_tuple = expected.named_tuple_value
+    expected_named_tuple.name = "NamedTuple"
+    key_value_pair = expected_named_tuple.values.add()
+    key_value_pair.key = "x"
+    list_value = key_value_pair.value.list_value
+    list_value.values.add().int64_value = 1
+    list_value.values.add().int64_value = 2
+    key_value_pair = expected_named_tuple.values.add()
+    key_value_pair.key = "y"
+    key_value_pair.value.string_value = "hello"
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(named_tuple._asdict(), decoded._asdict())
+    self.assertEqual(named_tuple.__class__.__name__, decoded.__class__.__name__)
+
+  def testNone(self):
+    structure = [1.0, None]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected.list_value.values.add().float64_value = 1.0
+    expected.list_value.values.add().none_value.CopyFrom(struct_pb2.NoneValue())
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testBool(self):
+    structure = [False]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected.list_value.values.add().bool_value = False
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEmptyStructures(self):
+    structure = [list(), dict(), tuple()]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected.list_value.values.add().list_value.CopyFrom(struct_pb2.ListValue())
+    expected.list_value.values.add().dict_value.CopyFrom(struct_pb2.DictValue())
+    expected.list_value.values.add().tuple_value.CopyFrom(
+        struct_pb2.TupleValue())
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testDtype(self):
+    structure = [dtypes.int64]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    list_value = expected.list_value.values.add()
+    list_value.tensor_dtype_value = dtypes.int64.as_datatype_enum
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeTensorSpec(self):
+    structure = [tensor.TensorSpec([1, 2, 3], dtypes.int64, "hello")]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected_list = expected.list_value
+    expected_tensor_spec = expected_list.values.add().tensor_spec_value
+    expected_tensor_spec.shape.dim.add().size = 1
+    expected_tensor_spec.shape.dim.add().size = 2
+    expected_tensor_spec.shape.dim.add().size = 3
+    expected_tensor_spec.name = "hello"
+    expected_tensor_spec.dtype = dtypes.int64.as_datatype_enum
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeTensorSpecWithNoName(self):
+    structure = [tensor.TensorSpec([1, 2, 3], dtypes.int64)]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected_list = expected.list_value
+    expected_tensor_spec = expected_list.values.add().tensor_spec_value
+    expected_tensor_spec.shape.dim.add().size = 1
+    expected_tensor_spec.shape.dim.add().size = 2
+    expected_tensor_spec.shape.dim.add().size = 3
+    expected_tensor_spec.name = ""
+    expected_tensor_spec.dtype = dtypes.int64.as_datatype_enum
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeRaggedTensorSpec(self):
+    structure = [
+        ragged_tensor.RaggedTensorSpec([1, 2, 3], dtypes.int64, 2, dtypes.int32)
+    ]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected_pbtxt = r"""
+      list_value {
+        values {
+          type_spec_value {
+            type_spec_class: RAGGED_TENSOR_SPEC
+            type_spec_class_name: 'RaggedTensorSpec'
+            num_flat_components: 3
+            type_state {
+              tuple_value {
+                # spec._shape
+                values {
+                  tensor_shape_value {
+                    dim { size: 1 }
+                    dim { size: 2 }
+                    dim { size: 3 }
+                  }
+                }
+                # spec._dtype
+                values { tensor_dtype_value: DT_INT64 }
+                # spec._ragged_rank
+                values { int64_value: 2 }
+                # spec._row_splits_dtype
+                values { tensor_dtype_value: DT_INT32 }
+              }
+            }
+          }
+        }
+      }
+    """
+    expected = struct_pb2.StructuredValue()
+    text_format.Parse(expected_pbtxt, expected)
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeSparseTensorSpec(self):
+    structure = [sparse_tensor.SparseTensorSpec([10, 20], dtypes.float32)]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected_pbtxt = r"""
+      list_value {
+        values {
+          type_spec_value {
+            type_spec_class: SPARSE_TENSOR_SPEC
+            type_spec_class_name: 'SparseTensorSpec'
+            num_flat_components: 3
+            type_state {
+              tuple_value {
+                # spec._shape
+                values {
+                  tensor_shape_value {
+                    dim { size: 10 }
+                    dim { size: 20 }
+                  }
+                }
+                # spec._dtype
+                values { tensor_dtype_value: DT_FLOAT }
+              }
+            }
+          }
+        }
+      }
+    """
+    expected = struct_pb2.StructuredValue()
+    text_format.Parse(expected_pbtxt, expected)
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeExtensionTypeSpec(self):
+
+    class Zoo(extension_type.ExtensionType):
+      __name__ = "tf.nested_structure_coder_test.Zoo"
+      zookeepers: typing.Tuple[str, ...]
+      animals: typing.Mapping[str, tensor.Tensor]
+
+    structure = [
+        Zoo.Spec(
+            zookeepers=["Zoey", "Zack"],
+            animals={"tiger": tensor.TensorSpec([16])})
+    ]
+
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected_pbtxt = r"""
+    list_value {
+      values {
+        type_spec_value {
+          type_spec_class: EXTENSION_TYPE_SPEC
+          type_spec_class_name: "tf.nested_structure_coder_test.Zoo.Spec"
+          num_flat_components: 1
+          type_state {
+            tuple_value {
+              values {
+                tuple_value {
+                  values { string_value: "zookeepers" }
+                  values { tuple_value {
+                    values { string_value: "Zoey" }
+                    values { string_value: "Zack" } } } } }
+              values {
+                tuple_value {
+                  values { string_value: "animals" }
+                  values { dict_value {
+                    fields {
+                      key: "tiger"
+                      value { tensor_spec_value {
+                        shape { dim { size: 16 } }
+                        dtype: DT_FLOAT } } } } } } } } } } } }
+    """
+    expected = struct_pb2.StructuredValue()
+    text_format.Parse(expected_pbtxt, expected)
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testDecodeUnknownTensorSpec(self):
+    encoded = struct_pb2.StructuredValue()
+    encoded.type_spec_value.type_spec_class = 0
+    encoded.type_spec_value.type_spec_class_name = "FutureTensorSpec"
+    with self.assertRaisesRegex(ValueError,
+                                "The type 'FutureTensorSpec' is not supported"):
+      nested_structure_coder.decode_proto(encoded)
+
+  def testEncodeDecodeBoundedTensorSpec(self):
+    structure = [
+        tensor.BoundedTensorSpec([1, 2, 3], dtypes.int64, 0, 10, "hello_0_10")
+    ]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected_list = expected.list_value
+    expected_tensor_spec = expected_list.values.add().bounded_tensor_spec_value
+    expected_tensor_spec.shape.dim.add().size = 1
+    expected_tensor_spec.shape.dim.add().size = 2
+    expected_tensor_spec.shape.dim.add().size = 3
+    expected_tensor_spec.name = "hello_0_10"
+    expected_tensor_spec.dtype = dtypes.int64.as_datatype_enum
+    expected_tensor_spec.minimum.CopyFrom(
+        tensor_util.make_tensor_proto([0], dtype=dtypes.int64, shape=[]))
+    expected_tensor_spec.maximum.CopyFrom(
+        tensor_util.make_tensor_proto([10], dtype=dtypes.int64, shape=[]))
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDecodeBoundedTensorSpecNoName(self):
+    structure = [
+        tensor.BoundedTensorSpec((28, 28, 3), dtypes.float64, -2, (1, 1, 20))
+    ]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected = struct_pb2.StructuredValue()
+    expected_list = expected.list_value
+    expected_tensor_spec = expected_list.values.add().bounded_tensor_spec_value
+    expected_tensor_spec.shape.dim.add().size = 28
+    expected_tensor_spec.shape.dim.add().size = 28
+    expected_tensor_spec.shape.dim.add().size = 3
+    expected_tensor_spec.name = ""
+    expected_tensor_spec.dtype = dtypes.float64.as_datatype_enum
+    expected_tensor_spec.minimum.CopyFrom(
+        tensor_util.make_tensor_proto([-2], dtype=dtypes.float64, shape=[]))
+    expected_tensor_spec.maximum.CopyFrom(
+        tensor_util.make_tensor_proto([1, 1, 20],
+                                      dtype=dtypes.float64,
+                                      shape=[3]))
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testEncodeDataSetSpec(self):
+    structure = [
+        dataset_ops.DatasetSpec({
+            "rt": ragged_tensor.RaggedTensorSpec([10, None], dtypes.int32),
+            "st": sparse_tensor.SparseTensorSpec([10, 20], dtypes.float32),
+            "t": tensor.TensorSpec([10, 8], dtypes.string)
+        })
+    ]
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testEncodeDecodeTensor(self):
+    structure = constant_op.constant(1)
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected_pbtxt = r"""
+      tensor_value {
+        dtype: DT_INT32
+        tensor_shape {
+        }
+        int_val: 1
+      }
+    """
+    expected = struct_pb2.StructuredValue()
+    text_format.Parse(expected_pbtxt, expected)
+    self.assertEqual(expected, encoded)
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertAllEqual(structure, decoded)
+
+  def testEncodeDecodeNumpy(self):
+    structure = np.array(1.0)
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    encoded = nested_structure_coder.encode_structure(structure)
+    expected_pbtxt = r"""
+      numpy_value {
+        dtype: DT_DOUBLE
+        tensor_shape {
+        }
+        double_val: 1.0
+      }
+    """
+    expected = struct_pb2.StructuredValue()
+    text_format.Parse(expected_pbtxt, expected)
+    self.assertEqual(expected, encoded)
+
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertIsInstance(decoded, np.ndarray)
+    self.assertAllEqual(structure, decoded)
+
+  def testNotEncodable(self):
+
+    class NotEncodable(object):
+      pass
+
+    self.assertFalse(nested_structure_coder.can_encode([NotEncodable()]))
+
+  def testRegisterCustomCodec(self):
+
+    class MyObject(object):
+      pass
+
+    class MyObjectCodec(object):
+      """Codec for MyObject."""
+
+      def can_encode(self, pyobj):
+        return isinstance(pyobj, MyObject)
+
+      def do_encode(self, array, encode_fn):
+        del array, encode_fn
+        return struct_pb2.StructuredValue()
+
+      def can_decode(self, value):
+        del value
+        return False
+
+      def do_decode(self, value, decode_fn):
+        raise NotImplementedError("Test only.")
+
+    nested_structure_coder.register_codec(MyObjectCodec())
+    my_object = MyObject()
+    self.assertTrue(nested_structure_coder.can_encode(my_object))
+
+  def testRegisteredTypeSpec(self):
+    expected_warning = ("Encoding a StructuredValue with type "
+                        "NestedStructureTest.RegisteredTypeSpec; loading "
+                        "this StructuredValue will require that this type "
+                        "be imported and registered")
+    structure = {"x": RegisteredTypeSpec()}
+
+    self.assertTrue(nested_structure_coder.can_encode(structure))
+    with warnings.catch_warnings(record=True) as w:
+      encoded = nested_structure_coder.encode_structure(structure)
+      self.assertLen(w, 1)
+      self.assertIn(expected_warning, str(w[0].message))
+    decoded = nested_structure_coder.decode_proto(encoded)
+    self.assertEqual(structure, decoded)
+
+  def testUnregisteredTypeSpec(self):
+    structure = {"x": UnregisteredTypeSpec()}
+    self.assertFalse(nested_structure_coder.can_encode(structure))
+    with self.assertRaises(nested_structure_coder.NotEncodableError):
+      nested_structure_coder.encode_structure(structure)
+
+  def testBuiltInTypeSpecCodecInvalidInputs(self):
+    class Foo:
+      pass
+
+    class Bar(internal.TypeSpec):
+      pass
+
+    with self.assertRaisesRegex(
+        ValueError, "The type '(.*?)' does not subclass tf.TypeSpec."):
+      nested_structure_coder.BuiltInTypeSpecCodec(Foo, 0)
+    with self.assertRaisesRegex(
+        ValueError, "The type '(.*?)' already has an instantiated codec."):
+      nested_structure_coder.BuiltInTypeSpecCodec(
+          dataset_ops.DatasetSpec, struct_pb2.TypeSpecProto.DATA_DATASET_SPEC)
+    with self.assertRaisesRegex(
+        ValueError, "The proto value '(.*?)' is already registered."):
+      nested_structure_coder.BuiltInTypeSpecCodec(
+          Bar, struct_pb2.TypeSpecProto.DATA_DATASET_SPEC)
+    with self.assertRaisesRegex(
+        ValueError, "The proto value '(.*?)' is invalid."):
+      nested_structure_coder.BuiltInTypeSpecCodec(Bar, 0)
+    with self.assertRaisesRegex(
+        ValueError, "The proto value '(.*?)' is invalid."):
+      nested_structure_coder.BuiltInTypeSpecCodec(Bar, 11)
+    with self.assertRaisesRegex(
+        ValueError, "The proto value '(.*?)' is invalid."):
+      nested_structure_coder.BuiltInTypeSpecCodec(Bar, 12)
+    with self.assertRaisesRegex(
+        ValueError, "The proto value '(.*?)' is invalid."):
+      nested_structure_coder.BuiltInTypeSpecCodec(Bar, 13)
+
+
+# Trivial TypeSpec class for testing.
+class UnregisteredTypeSpec(type_spec.TypeSpec):
+  value_type = property(lambda self: None)
+  _component_specs = property(lambda self: ())
+  _to_components = lambda self, v: ()
+  _from_components = classmethod(lambda cls, c: cls())
+  _serialize = lambda self: ()
+
+
+# Trivial TypeSpec class for testing.
+@type_spec_registry.register("NestedStructureTest.RegisteredTypeSpec")
+class RegisteredTypeSpec(type_spec.TypeSpec):
+  value_type = property(lambda self: None)
+  _component_specs = property(lambda self: ())
+  _to_components = lambda self, v: ()
+  _from_components = classmethod(lambda cls, c: cls())
+  _serialize = lambda self: ()
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/path_helpers.py

@@ -0,0 +1,82 @@
+# 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.
+# ==============================================================================
+
+"""Path helpers utility functions."""
+
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.saved_model import constants
+from tensorflow.python.util import compat
+
+
+def get_or_create_variables_dir(export_dir):
+  """Return variables sub-directory, or create one if it doesn't exist."""
+  variables_dir = get_variables_dir(export_dir)
+  file_io.recursive_create_dir(variables_dir)
+  return variables_dir
+
+
+def get_variables_dir(export_dir):
+  """Return variables sub-directory in the SavedModel."""
+  return file_io.join(
+      compat.as_text(export_dir), compat.as_text(constants.VARIABLES_DIRECTORY))
+
+
+def get_variables_path(export_dir):
+  """Return the variables path, used as the prefix for checkpoint files."""
+  return file_io.join(
+      compat.as_text(get_variables_dir(export_dir)),
+      compat.as_text(constants.VARIABLES_FILENAME))
+
+
+def get_or_create_assets_dir(export_dir):
+  """Return assets sub-directory, or create one if it doesn't exist."""
+  assets_destination_dir = get_assets_dir(export_dir)
+
+  file_io.recursive_create_dir(assets_destination_dir)
+
+  return assets_destination_dir
+
+
+def get_assets_dir(export_dir):
+  """Return path to asset directory in the SavedModel."""
+  return file_io.join(
+      compat.as_text(export_dir), compat.as_text(constants.ASSETS_DIRECTORY))
+
+
+def get_or_create_debug_dir(export_dir):
+  """Returns path to the debug sub-directory, creating if it does not exist."""
+  debug_dir = get_debug_dir(export_dir)
+
+  file_io.recursive_create_dir(debug_dir)
+
+  return debug_dir
+
+
+def get_saved_model_pbtxt_path(export_dir):
+  return file_io.join(
+      compat.as_bytes(compat.path_to_str(export_dir)),
+      compat.as_bytes(constants.SAVED_MODEL_FILENAME_PBTXT))
+
+
+def get_saved_model_pb_path(export_dir):
+  return file_io.join(
+      compat.as_bytes(compat.path_to_str(export_dir)),
+      compat.as_bytes(constants.SAVED_MODEL_FILENAME_PB))
+
+
+def get_debug_dir(export_dir):
+  """Returns path to the debug sub-directory in the SavedModel."""
+  return file_io.join(
+      compat.as_text(export_dir), compat.as_text(constants.DEBUG_DIRECTORY))

tensorflow/python/saved_model/proto_splitter_save_test.py

@@ -0,0 +1,66 @@
+# Copyright 2023 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 SavedModel save using experimental_image_format."""
+
+import os
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.module import module
+from tensorflow.python.saved_model import save
+from tensorflow.python.saved_model import save_options
+from tensorflow.tools.proto_splitter import constants
+
+
+class ProtoSplitterSaveTest(test.TestCase, parameterized.TestCase):
+
+  def test_save_experimental_image_format(self):
+    root = module.Module()
+    root.c = constant_op.constant(np.random.random_sample([150, 150]))
+    root.get_c = def_function.function(lambda: root.c)
+    save_dir = os.path.join(self.get_temp_dir(), "chunked_model")
+    constants.debug_set_max_size(80000)
+    options = save_options.SaveOptions(experimental_image_format=True)
+    save.save(
+        root,
+        save_dir,
+        signatures=root.get_c.get_concrete_function(),
+        options=options,
+    )
+    self.assertTrue(os.path.exists(save_dir + "/saved_model.cpb"))
+
+  def test_save_experimental_image_format_not_chunked(self):
+    root = module.Module()
+    root.c = constant_op.constant(np.random.random_sample([150, 150]))
+    root.get_c = def_function.function(lambda: root.c)
+    save_dir = os.path.join(self.get_temp_dir(), "not_chunked_model")
+    constants.debug_set_max_size(1 << 31)  # 2GB
+    options = save_options.SaveOptions(experimental_image_format=True)
+    save.save(
+        root,
+        save_dir,
+        signatures=root.get_c.get_concrete_function(),
+        options=options,
+    )
+    # Should save an unchunked proto (.pb) and not .cpb
+    self.assertTrue(os.path.exists(save_dir + "/saved_model.pb"))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/pywrap_saved_model.cc

@@ -0,0 +1,46 @@
+/* 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.
+==============================================================================*/
+
+// Defines the pywrap_saved_model module. In order to have only one dynamically-
+// linked shared object, all SavedModel python bindings should be added here.
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/cc/experimental/libexport/save.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+#include "tensorflow/python/saved_model/pywrap_saved_model_constants.h"
+#include "tensorflow/python/saved_model/pywrap_saved_model_fingerprinting.h"
+// Placeholder for protosplitter merger include.
+#include "tensorflow/python/saved_model/pywrap_saved_model_metrics.h"
+
+namespace tensorflow {
+namespace saved_model {
+namespace python {
+
+PYBIND11_MODULE(pywrap_saved_model, m) {
+  m.doc() = "TensorFlow SavedModel Python bindings";
+
+  m.def("Save", [](const char* export_dir) {
+    MaybeRaiseFromStatus(libexport::Save(export_dir));
+  });
+
+  DefineConstantsModule(m);
+  DefineMetricsModule(m);
+  DefineFingerprintingModule(m);
+  // Placeholder for protosplitter merger module definition.
+}
+
+}  // namespace python
+}  // namespace saved_model
+}  // namespace tensorflow

tensorflow/python/saved_model/pywrap_saved_model/__init__.pyi

@@ -0,0 +1,18 @@
+# Copyright 2023 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 . import constants as constants, fingerprinting as fingerprinting, merger as merger, metrics as metrics
+
+def Save(arg0: str) -> None: ...

tensorflow/python/saved_model/pywrap_saved_model/constants.pyi

@@ -0,0 +1,33 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+ASSETS_DIRECTORY: str
+ASSETS_KEY: str
+DEBUG_DIRECTORY: str
+DEBUG_INFO_FILENAME_PB: str
+EXTRA_ASSETS_DIRECTORY: str
+FINGERPRINT_FILENAME: str
+INIT_OP_SIGNATURE_KEY: str
+LEGACY_INIT_OP_KEY: str
+MAIN_OP_KEY: str
+SAVED_MODEL_FILENAME_CPB: str
+SAVED_MODEL_FILENAME_PB: str
+SAVED_MODEL_FILENAME_PBTXT: str
+SAVED_MODEL_FILENAME_PREFIX: str
+SAVED_MODEL_SCHEMA_VERSION: int
+TRAIN_OP_KEY: str
+TRAIN_OP_SIGNATURE_KEY: str
+VARIABLES_DIRECTORY: str
+VARIABLES_FILENAME: str

tensorflow/python/saved_model/pywrap_saved_model/fingerprinting.pyi

@@ -0,0 +1,24 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+class FileNotFoundException(Exception): ...
+
+class FingerprintException(Exception): ...
+
+def CreateFingerprintDef(export_dir: str) -> bytes: ...
+def ReadSavedModelFingerprint(export_dir: str) -> bytes: ...
+def Singleprint(graph_def_program_hash: int, signature_def_hash: int, saved_object_graph_hash: int, checkpoint_hash: int) -> str: ...
+def SingleprintFromFP(export_dir: str) -> str: ...
+def SingleprintFromSM(export_dir: str) -> str: ...

tensorflow/python/saved_model/pywrap_saved_model/merger.pyi

@@ -0,0 +1,18 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+class MergerException(Exception): ...
+
+def MergerRead(*args, **kwargs): ...

tensorflow/python/saved_model/pywrap_saved_model/metrics.pyi

@@ -0,0 +1,60 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+kFingerprintError: str
+kFingerprintFound: str
+kFingerprintNotFound: str
+
+class MetricException(Exception): ...
+
+def AddAsyncCheckpointWriteDuration(api_label: str, microseconds: float) -> None: ...
+def AddCheckpointReadDuration(api_label: str, microseconds: float) -> None: ...
+def AddCheckpointWriteDuration(api_label: str, microseconds: float) -> None: ...
+def AddNumCheckpointShardsWritten(num_shards: int) -> None: ...
+def AddShardingCallbackDuration(callback_duration: int) -> None: ...
+def AddTrainingTimeSaved(api_label: str, microseconds: float) -> None: ...
+def CalculateFileSize(arg0: str) -> int: ...
+def GetAsyncCheckpointWriteDurations(api_label: str) -> bytes: ...
+def GetCheckpointReadDurations(api_label: str) -> bytes: ...
+def GetCheckpointSize(api_label: str, filesize: int) -> int: ...
+def GetCheckpointWriteDurations(api_label: str) -> bytes: ...
+def GetFoundFingerprintOnLoad() -> str: ...
+def GetNumCheckpointShardsWritten() -> int: ...
+def GetRead(write_version: str) -> int: ...
+def GetReadApi(arg0: str) -> int: ...
+def GetReadFingerprint() -> str: ...
+def GetReadPath() -> str: ...
+def GetReadPathAndSingleprint() -> tuple[str, str]: ...
+def GetShardingCallbackDescription() -> str: ...
+def GetShardingCallbackDuration() -> int: ...
+def GetTrainingTimeSaved(api_label: str) -> int: ...
+def GetWrite(write_version: str) -> int: ...
+def GetWriteApi(arg0: str) -> int: ...
+def GetWriteFingerprint() -> str: ...
+def GetWritePath() -> str: ...
+def GetWritePathAndSingleprint() -> tuple[str, str]: ...
+def IncrementRead(write_version: str) -> None: ...
+def IncrementReadApi(arg0: str) -> None: ...
+def IncrementWrite(write_version: str) -> None: ...
+def IncrementWriteApi(arg0: str) -> None: ...
+def RecordCheckpointSize(api_label: str, filesize: int) -> None: ...
+def SetFoundFingerprintOnLoad(found_status: str) -> None: ...
+def SetReadFingerprint(fingerprint: bytes) -> None: ...
+def SetReadPath(saved_model_path: str) -> None: ...
+def SetReadPathAndSingleprint(path: str, singleprint: str) -> None: ...
+def SetShardingCallbackDescription(description: str) -> None: ...
+def SetWriteFingerprint(fingerprint: bytes) -> None: ...
+def SetWritePath(saved_model_path: str) -> None: ...
+def SetWritePathAndSingleprint(path: str, singleprint: str) -> None: ...

tensorflow/python/saved_model/pywrap_saved_model_constants.cc

@@ -0,0 +1,82 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/python/saved_model/pywrap_saved_model_constants.h"
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "tensorflow/cc/saved_model/constants.h"
+
+namespace tensorflow {
+namespace saved_model {
+namespace python {
+
+namespace py = pybind11;
+
+void DefineConstantsModule(py::module main_module) {
+  auto m = main_module.def_submodule("constants");
+
+  m.doc() = "Python bindings for TensorFlow SavedModel Constants";
+
+  m.attr("ASSETS_DIRECTORY") = py::str(tensorflow::kSavedModelAssetsDirectory);
+
+  m.attr("EXTRA_ASSETS_DIRECTORY") =
+      py::str(tensorflow::kSavedModelAssetsExtraDirectory);
+
+  m.attr("ASSETS_KEY") = py::str(tensorflow::kSavedModelAssetsKey);
+
+  m.attr("DEBUG_DIRECTORY") = py::str(tensorflow::kSavedModelDebugDirectory);
+
+  m.attr("DEBUG_INFO_FILENAME_PB") =
+      py::str(tensorflow::kSavedModelDebugInfoFilenamePb);
+
+  m.attr("INIT_OP_SIGNATURE_KEY") =
+      py::str(tensorflow::kSavedModelInitOpSignatureKey);
+
+  m.attr("LEGACY_INIT_OP_KEY") =
+      py::str(tensorflow::kSavedModelLegacyInitOpKey);
+
+  m.attr("MAIN_OP_KEY") = py::str(tensorflow::kSavedModelMainOpKey);
+
+  m.attr("TRAIN_OP_KEY") = py::str(tensorflow::kSavedModelTrainOpKey);
+
+  m.attr("TRAIN_OP_SIGNATURE_KEY") =
+      py::str(tensorflow::kSavedModelTrainOpSignatureKey);
+
+  m.attr("SAVED_MODEL_FILENAME_PREFIX") =
+      py::str(tensorflow::kSavedModelFilenamePrefix);
+
+  m.attr("SAVED_MODEL_FILENAME_PB") =
+      py::str(tensorflow::kSavedModelFilenamePb);
+
+  m.attr("SAVED_MODEL_FILENAME_CPB") =
+      py::str(tensorflow::kSavedModelFilenameCpb);
+
+  m.attr("SAVED_MODEL_FILENAME_PBTXT") =
+      py::str(tensorflow::kSavedModelFilenamePbTxt);
+
+  m.attr("SAVED_MODEL_SCHEMA_VERSION") = tensorflow::kSavedModelSchemaVersion;
+
+  m.attr("VARIABLES_DIRECTORY") =
+      py::str(tensorflow::kSavedModelVariablesDirectory);
+
+  m.attr("VARIABLES_FILENAME") =
+      py::str(tensorflow::kSavedModelVariablesFilename);
+
+  m.attr("FINGERPRINT_FILENAME") = py::str(tensorflow::kFingerprintFilenamePb);
+}
+
+}  // namespace python
+}  // namespace saved_model
+}  // namespace tensorflow

tensorflow/python/saved_model/pywrap_saved_model_constants.h

@@ -0,0 +1,35 @@
+/* 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.
+==============================================================================*/
+
+// Defines the 'constants' wrapper submodule.
+
+#ifndef TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_CONSTANTS_H_
+#define TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_CONSTANTS_H_
+
+// placeholder for index annotation headers
+#include "pybind11/pybind11.h"  // from @pybind11
+
+namespace tensorflow {
+namespace saved_model {
+namespace python {
+
+// Wraps the SavedModel constants for exporting to Python.
+void DefineConstantsModule(pybind11::module main_module);
+
+}  // namespace python
+}  // namespace saved_model
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_CONSTANTS_H_

tensorflow/python/saved_model/pywrap_saved_model_fingerprinting.cc

@@ -0,0 +1,209 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/python/saved_model/pywrap_saved_model_fingerprinting.h"
+
+#include <cstdint>
+#include <exception>
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "pybind11_abseil/absl_casters.h"  // from @pybind11_abseil
+#include "tensorflow/cc/saved_model/constants.h"
+#include "tensorflow/cc/saved_model/fingerprinting.h"
+#include "tensorflow/cc/saved_model/reader.h"
+#include "tensorflow/core/common_runtime/graph_runner.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/protobuf/fingerprint.pb.h"
+#include "tensorflow/core/protobuf/saved_model.pb.h"
+#include "tensorflow/python/lib/core/pybind11_status.h"
+
+namespace tensorflow {
+namespace saved_model {
+namespace python {
+
+namespace py = pybind11;
+
+class FingerprintException : public std::exception {
+ public:
+  explicit FingerprintException(const char *m) : message_{m} {}
+  const char *what() const noexcept override { return message_.c_str(); }
+
+ private:
+  std::string message_ = "";
+};
+
+class FileNotFoundException : public std::exception {
+ public:
+  explicit FileNotFoundException(const char *m) : message_{m} {}
+  const char *what() const noexcept override { return message_.c_str(); }
+
+ private:
+  std::string message_ = "";
+};
+
+void DefineFingerprintingModule(py::module main_module) {
+  auto m = main_module.def_submodule("fingerprinting");
+
+  m.doc() = "Python bindings for TensorFlow SavedModel Fingerprinting.";
+
+  static py::exception<FingerprintException> fp_ex(m, "FingerprintException");
+  py::register_exception_translator([](std::exception_ptr p) {
+    try {
+      if (p) {
+        std::rethrow_exception(p);
+      }
+    } catch (const FingerprintException &e) {
+      fp_ex(e.what());
+    }
+  });
+
+  static py::exception<FileNotFoundException> fnf_ex(m,
+                                                     "FileNotFoundException");
+  py::register_exception_translator([](std::exception_ptr p) {
+    try {
+      if (p) {
+        std::rethrow_exception(p);
+      }
+    } catch (const FileNotFoundException &e) {
+      fnf_ex(e.what());
+    }
+  });
+
+  m.def(
+      "CreateFingerprintDef",
+      [](std::string export_dir) -> absl::StatusOr<py::bytes> {
+        absl::StatusOr<FingerprintDef> fingerprint =
+            fingerprinting::CreateFingerprintDef(export_dir);
+        if (fingerprint.ok()) {
+          return py::bytes(fingerprint.value().SerializeAsString());
+        }
+        throw FingerprintException(
+            absl::StrCat(
+                std::string("Could not create fingerprint in directory: " +
+                            export_dir),
+                "\n", fingerprint.status().ToString())
+                .c_str());
+      },
+      py::arg("export_dir"),
+      py::doc(
+          "Returns the serialized FingerprintDef of a SavedModel on disk."));
+
+  m.def(
+      "ReadSavedModelFingerprint",
+      [](std::string export_dir) {
+        absl::StatusOr<FingerprintDef> fingerprint =
+            fingerprinting::ReadSavedModelFingerprint(export_dir);
+        if (fingerprint.ok()) {
+          return py::bytes(fingerprint.value().SerializeAsString());
+        } else if (fingerprint.status().code() == absl::StatusCode::kNotFound) {
+          throw FileNotFoundException(
+              absl::StrCat(
+                  std::string("Could not find fingerprint in directory: " +
+                              export_dir),
+                  "\n", fingerprint.status().ToString())
+                  .c_str());
+        } else {
+          throw FingerprintException(
+              absl::StrCat(
+                  std::string(
+                      "Could not read fingerprint from fingerprint.pb file "
+                      "in directory: " +
+                      export_dir),
+                  "\n", fingerprint.status().ToString())
+                  .c_str());
+        }
+      },
+      py::arg("export_dir"),
+      py::doc(
+          "Loads the `fingerprint.pb` from `export_dir`, returns an error if "
+          "there is none."));
+
+  m.def(
+      "SingleprintFromFP",
+      [](std::string export_dir) {
+        absl::StatusOr<std::string> singleprint =
+            fingerprinting::Singleprint(export_dir);
+        if (singleprint.ok()) {
+          return py::str(singleprint.value());
+        }
+        throw FingerprintException(
+            absl::StrCat(
+                std::string("Could not create singleprint from the fingerprint "
+                            "specified by the export_dir."),
+                "\n", singleprint.status().ToString())
+                .c_str());
+      },
+      py::arg("export_dir"),
+      py::doc("Canonical fingerprinting ID for a SavedModel."));
+
+  m.def(
+      "SingleprintFromSM",
+      [](std::string export_dir) {
+        absl::StatusOr<FingerprintDef> fingerprint_def =
+            fingerprinting::CreateFingerprintDef(export_dir);
+        if (!fingerprint_def.ok()) {
+          throw FingerprintException(
+              absl::StrCat(
+                  std::string(
+                      "Could not create singleprint from the saved_model."),
+                  "\n", fingerprint_def.status().ToString())
+                  .c_str());
+        }
+
+        absl::StatusOr<std::string> singleprint =
+            fingerprinting::Singleprint(fingerprint_def.value());
+        if (!singleprint.ok()) {
+          throw FingerprintException(
+              absl::StrCat(
+                  std::string(
+                      "Could not create singleprint from the saved_model."),
+                  "\n", singleprint.status().ToString())
+                  .c_str());
+        }
+
+        return py::str(singleprint.value());
+      },
+      py::arg("export_dir"),
+      py::doc("Canonical fingerprinting ID for a SavedModel."));
+
+  m.def(
+      "Singleprint",
+      [](uint64_t graph_def_program_hash, uint64_t signature_def_hash,
+         uint64_t saved_object_graph_hash, uint64_t checkpoint_hash) {
+        absl::StatusOr<std::string> singleprint = fingerprinting::Singleprint(
+            graph_def_program_hash, signature_def_hash, saved_object_graph_hash,
+            checkpoint_hash);
+        if (singleprint.ok()) {
+          return py::str(singleprint.value());
+        }
+        throw FingerprintException(
+            absl::StrCat(
+                std::string("Could not create singleprint from given values."),
+                "\n", singleprint.status().ToString())
+                .c_str());
+      },
+      py::arg("graph_def_program_hash"), py::arg("signature_def_hash"),
+      py::arg("saved_object_graph_hash"), py::arg("checkpoint_hash"),
+      py::doc("Canonical fingerprinting ID for a SavedModel."));
+}
+
+}  // namespace python
+}  // namespace saved_model
+}  // namespace tensorflow

tensorflow/python/saved_model/pywrap_saved_model_fingerprinting.h

@@ -0,0 +1,35 @@
+/* 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.
+==============================================================================*/
+
+// Defines the 'fingerprinting' wrapper submodule.
+
+#ifndef TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_FINGERPRINTING_H_
+#define TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_FINGERPRINTING_H_
+
+// placeholder for index annotation headers
+#include "pybind11/pybind11.h"  // from @pybind11
+
+namespace tensorflow {
+namespace saved_model {
+namespace python {
+
+// Wraps the SavedModel fingerprinting methods for exporting to Python.
+void DefineFingerprintingModule(pybind11::module main_module);
+
+}  // namespace python
+}  // namespace saved_model
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_FINGERPRINTING_H_

tensorflow/python/saved_model/pywrap_saved_model_fingerprinting_test.py

@@ -0,0 +1,123 @@
+# 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 pywrap_saved_model_fingerprinting."""
+
+from tensorflow.core.protobuf import fingerprint_pb2
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model.pywrap_saved_model import fingerprinting as pywrap_fingerprinting
+
+is_oss = True  # Updated by copybara.
+
+
+class FingerprintingTest(test.TestCase):
+  def test_create_fingerprint_def(self):
+    export_dir = test.test_src_dir_path(
+        "cc/saved_model/testdata/VarsAndArithmeticObjectGraph")
+
+    fingerprint = fingerprint_pb2.FingerprintDef().FromString(
+        pywrap_fingerprinting.CreateFingerprintDef(export_dir))
+
+    # We cannot check the value of the saved_model_checksum due to
+    # non-determinism in saving.
+    self.assertGreater(fingerprint.saved_model_checksum, 0)
+    self.assertEqual(fingerprint.graph_def_program_hash, 10127142238652115842)
+    self.assertEqual(fingerprint.signature_def_hash, 15570736222402453744)
+    self.assertEqual(fingerprint.saved_object_graph_hash, 3678101440349108924)
+    # TODO(b/242348400): The checkpoint hash is non-deterministic, so we cannot
+    # check its value here.
+    self.assertGreater(fingerprint.checkpoint_hash, 0)
+
+  def test_read_saved_model_fingerprint(self):
+    export_dir = test.test_src_dir_path(
+        "cc/saved_model/testdata/VarsAndArithmeticObjectGraph")
+    fingerprint = fingerprint_pb2.FingerprintDef().FromString(
+        pywrap_fingerprinting.ReadSavedModelFingerprint(export_dir))
+
+    self.assertGreater(fingerprint.saved_model_checksum, 0)
+    self.assertEqual(fingerprint.graph_def_program_hash, 706963557435316516)
+    self.assertEqual(fingerprint.signature_def_hash, 5693392539583495303)
+    self.assertEqual(fingerprint.saved_object_graph_hash, 12074714563970609759)
+    self.assertGreater(fingerprint.checkpoint_hash, 0)
+    self.assertEqual(fingerprint.version.producer, 1)
+
+  def test_read_nonexistent_fingerprint(self):
+    export_dir = test.test_src_dir_path("cc/saved_model/testdata/AssetModule")
+    with self.assertRaises(
+        pywrap_fingerprinting.FileNotFoundException) as excinfo:
+      pywrap_fingerprinting.ReadSavedModelFingerprint(export_dir)
+    self.assertRegex(str(excinfo.exception), "Could not find fingerprint.")
+
+  def test_read_saved_model_singleprint(self):
+    export_dir = test.test_src_dir_path(
+        "cc/saved_model/testdata/VarsAndArithmeticObjectGraph")
+    fingerprint = fingerprint_pb2.FingerprintDef().FromString(
+        pywrap_fingerprinting.ReadSavedModelFingerprint(export_dir))
+    singleprint = pywrap_fingerprinting.Singleprint(
+        fingerprint.graph_def_program_hash,
+        fingerprint.signature_def_hash,
+        fingerprint.saved_object_graph_hash,
+        fingerprint.checkpoint_hash)
+    # checkpoint_hash is non-deterministic and not included
+    self.assertRegex(singleprint,
+                     "/".join([
+                         "706963557435316516",  # graph_def_program_hash
+                         "5693392539583495303",  # signature_def_hash
+                         "12074714563970609759",  # saved_object_graph_hash
+                         ]))
+
+  def test_read_saved_model_singleprint_from_fp(self):
+    export_dir = test.test_src_dir_path(
+        "cc/saved_model/testdata/VarsAndArithmeticObjectGraph")
+    singleprint = pywrap_fingerprinting.SingleprintFromFP(export_dir)
+    # checkpoint_hash is non-deterministic and not included
+    self.assertRegex(singleprint,
+                     "/".join([
+                         "706963557435316516",  # graph_def_program_hash
+                         "5693392539583495303",  # signature_def_hash
+                         "12074714563970609759",  # saved_object_graph_hash
+                         ]))
+
+  def test_read_saved_model_singleprint_from_sm(self):
+    export_dir = test.test_src_dir_path(
+        "cc/saved_model/testdata/AssetModule")
+    singleprint = pywrap_fingerprinting.SingleprintFromSM(export_dir)
+    # checkpoint_hash is non-deterministic and not included
+    self.assertRegex(singleprint,
+                     "/".join([
+                         "14732473038199296573",  # graph_def_program_hash
+                         "11983586671997178523",  # signature_def_hash
+                         "14640180866165615446",  # saved_object_graph_hash
+                         ]))
+
+  def test_read_chunked_saved_model_fingerprint(self):
+    if is_oss:
+      self.skipTest("Experimental image format disabled in OSS.")
+    export_dir = test.test_src_dir_path(
+        "cc/saved_model/testdata/chunked_saved_model/chunked_model")
+    fingerprint = fingerprint_pb2.FingerprintDef().FromString(
+        pywrap_fingerprinting.CreateFingerprintDef(export_dir))
+    self.assertGreater(fingerprint.saved_model_checksum, 0)
+    # We test for multiple fingerprints due to non-determinism when building
+    # with different compilation_mode flag options.
+    self.assertIn(fingerprint.graph_def_program_hash,
+                  [906548630859202535, 9562420523583756263])
+    self.assertEqual(fingerprint.signature_def_hash, 1043582354059066488)
+    self.assertIn(fingerprint.saved_object_graph_hash,
+                  [11894619660760763927, 2766043449526180728])
+    self.assertEqual(fingerprint.checkpoint_hash, 0)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/pywrap_saved_model_metrics.cc

@@ -0,0 +1,478 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/python/saved_model/pywrap_saved_model_metrics.h"
+
+#include <cstdint>
+#include <exception>
+#include <string>
+#include <utility>
+
+// Placeholder for lineage logging import.
+// Placeholder for lineage logging additional import.
+#include "absl/base/call_once.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
+#include "pybind11/attr.h"  // from @pybind11
+#include "pybind11/cast.h"  // from @pybind11
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "tensorflow/cc/saved_model/metrics.h"
+#include "tensorflow/core/protobuf/fingerprint.pb.h"
+
+namespace tensorflow {
+namespace saved_model {
+namespace python {
+
+// Helper variable for logging.
+
+namespace py = pybind11;
+
+class MetricException : public std::exception {
+ public:
+  explicit MetricException(const char* m) : message_{m} {}
+  const char* what() const noexcept override { return message_.c_str(); }
+
+ private:
+  std::string message_ = "";
+};
+
+// Placeholder for a helper function for logging.
+
+void DefineMetricsModule(py::module main_module) {
+  auto m = main_module.def_submodule("metrics");
+
+  m.doc() = "Python bindings for TensorFlow SavedModel and Checkpoint Metrics.";
+
+  static py::exception<MetricException> fp_ex(m, "MetricException");
+  py::register_exception_translator([](std::exception_ptr p) {
+    try {
+      if (p) {
+        std::rethrow_exception(p);
+      }
+    } catch (const MetricException& e) {
+      fp_ex(e.what());
+    }
+  });
+
+  m.def(
+      "IncrementWrite",
+      [](const char* write_version) {
+        metrics::SavedModelWriteCount(write_version).IncrementBy(1);
+      },
+      py::kw_only(), py::arg("write_version"),
+      py::doc("Increment the '/tensorflow/core/saved_model/write/count' "
+              "counter."));
+
+  m.def(
+      "GetWrite",
+      [](const char* write_version) {
+        return metrics::SavedModelWriteCount(write_version).value();
+      },
+      py::kw_only(), py::arg("write_version"),
+      py::doc("Get value of '/tensorflow/core/saved_model/write/count' "
+              "counter."));
+
+  m.def(
+      "IncrementWriteApi",
+      [](const char* api_label) {
+        metrics::SavedModelWriteApi(api_label).IncrementBy(1);
+      },
+      py::doc("Increment the '/tensorflow/core/saved_model/write/api' "
+              "counter for API with `api_label`"));
+
+  m.def(
+      "GetWriteApi",
+      [](const char* api_label) {
+        return metrics::SavedModelWriteApi(api_label).value();
+      },
+      py::doc("Get value of '/tensorflow/core/saved_model/write/api' "
+              "counter for `api_label` cell."));
+
+  m.def(
+      "IncrementRead",
+      [](const char* write_version) {
+        metrics::SavedModelReadCount(write_version).IncrementBy(1);
+      },
+      py::kw_only(), py::arg("write_version"),
+      py::doc("Increment the '/tensorflow/core/saved_model/read/count' "
+              "counter after reading a SavedModel with the specifed "
+              "`write_version`."));
+
+  m.def(
+      "GetRead",
+      [](const char* write_version) {
+        return metrics::SavedModelReadCount(write_version).value();
+      },
+      py::kw_only(), py::arg("write_version"),
+      py::doc("Get value of '/tensorflow/core/saved_model/read/count' "
+              "counter for SavedModels with the specified `write_version`."));
+
+  m.def(
+      "IncrementReadApi",
+      [](const char* api_label) {
+        metrics::SavedModelReadApi(api_label).IncrementBy(1);
+      },
+      py::doc("Increment the '/tensorflow/core/saved_model/read/api' "
+              "counter for API with `api_label`."));
+
+  m.def(
+      "GetReadApi",
+      [](const char* api_label) {
+        return metrics::SavedModelReadApi(api_label).value();
+      },
+      py::doc("Get value of '/tensorflow/core/saved_model/read/api' "
+              "counter for `api_label` cell."));
+
+  m.def(
+      "SetReadFingerprint",
+      [](const py::bytes fingerprint) {
+        FingerprintDef fingerprint_def;
+        fingerprint_def.ParseFromString(std::string(fingerprint));
+        metrics::SavedModelReadFingerprint().Set(
+            metrics::MakeFingerprintJson(fingerprint_def).c_str());
+      },
+      py::kw_only(), py::arg("fingerprint"),
+      py::doc("Set the '/tensorflow/core/saved_model/read/fingerprint' gauge "
+              "with `fingerprint`."));
+
+  m.def(
+      "GetReadFingerprint",
+      []() { return metrics::SavedModelReadFingerprint().value(); },
+      py::doc("Get value of '/tensorflow/core/saved_model/read/fingerprint' "
+              "gauge."));
+
+  m.def(
+      "SetWriteFingerprint",
+      [](const py::bytes fingerprint) {
+        FingerprintDef fingerprint_def;
+        fingerprint_def.ParseFromString(std::string(fingerprint));
+        metrics::SavedModelWriteFingerprint().Set(
+            metrics::MakeFingerprintJson(fingerprint_def).c_str());
+      },
+      py::kw_only(), py::arg("fingerprint"),
+      py::doc("Set the '/tensorflow/core/saved_model/write/fingerprint' gauge "
+              "with `fingerprint`."));
+
+  m.def(
+      "GetWriteFingerprint",
+      []() { return metrics::SavedModelWriteFingerprint().value(); },
+      py::doc("Get value of '/tensorflow/core/saved_model/write/fingerprint' "
+              "gauge."));
+
+  m.def(
+      "SetReadPath",
+      [](const char* saved_model_path) {
+        metrics::SavedModelReadPath().Set(saved_model_path);
+      },
+      py::kw_only(), py::arg("saved_model_path"),
+      py::doc("Set the '/tensorflow/core/saved_model/read/path' gauge "
+              "with `saved_model_path`."));
+
+  m.def(
+      "GetReadPath", []() { return metrics::SavedModelReadPath().value(); },
+      py::doc("Get value of '/tensorflow/core/saved_model/read/path' gauge."));
+
+  m.def(
+      "SetWritePath",
+      [](const char* saved_model_path) {
+        metrics::SavedModelWritePath().Set(saved_model_path);
+      },
+      py::kw_only(), py::arg("saved_model_path"),
+      py::doc("Set the '/tensorflow/core/saved_model/write/path' gauge "
+              "with `saved_model_path`."));
+
+  m.def(
+      "GetWritePath", []() { return metrics::SavedModelWritePath().value(); },
+      py::doc("Get value of '/tensorflow/core/saved_model/write/path' gauge."));
+
+  m.def(
+      "SetReadPathAndSingleprint",
+      [](const char* path, const char* singleprint) {
+        auto path_and_singleprint =
+            metrics::MakeSavedModelPathAndSingleprint(path, singleprint);
+        if (!path_and_singleprint.ok()) {
+          throw MetricException(path_and_singleprint.status().message().data());
+        }
+        metrics::SavedModelReadPathAndSingleprint().Set(
+            path_and_singleprint.value());
+        // Placeholder for lineage logging dedup setup.
+        // Placeholder for lineage logging input call.
+      },
+      py::kw_only(), py::arg("path"), py::arg("singleprint"),
+      py::doc(
+          "Set the '/tensorflow/core/saved_model/read/path_and_singleprint' "
+          "gauge with `path` and `singleprint`."));
+
+  m.def(
+      "GetReadPathAndSingleprint",
+      []() {
+        std::string path_and_singleprint =
+            metrics::SavedModelReadPathAndSingleprint().value();
+        if (path_and_singleprint.empty())
+          return std::pair<std::string, std::string>();
+        auto parsed_path_and_singleprint =
+            metrics::ParseSavedModelPathAndSingleprint(path_and_singleprint);
+        if (!parsed_path_and_singleprint.ok()) {
+          throw MetricException(
+              parsed_path_and_singleprint.status().message().data());
+        }
+        return parsed_path_and_singleprint.value();
+      },
+      py::doc(
+          "Get tuple of `path` and `singleprint` values of "
+          "'/tensorflow/core/saved_model/read/path_and_singleprint' gauge."));
+
+  m.def(
+      "SetWritePathAndSingleprint",
+      [](const char* path, const char* singleprint) {
+        auto path_and_singleprint =
+            metrics::MakeSavedModelPathAndSingleprint(path, singleprint);
+        if (!path_and_singleprint.ok()) {
+          throw MetricException(path_and_singleprint.status().message().data());
+        }
+        metrics::SavedModelWritePathAndSingleprint().Set(
+            path_and_singleprint.value());
+        // Placeholder for lineage logging dedup setup.
+        // Placeholder for lineage logging output call.
+      },
+      py::kw_only(), py::arg("path"), py::arg("singleprint"),
+      py::doc("Set the "
+              "'/tensorflow/core/saved_model/write/path_and_singleprint' gauge "
+              "with `path` and `singleprint`."));
+
+  m.def(
+      "GetWritePathAndSingleprint",
+      []() {
+        std::string path_and_singleprint =
+            metrics::SavedModelWritePathAndSingleprint().value();
+        if (path_and_singleprint.empty())
+          return std::pair<std::string, std::string>();
+        auto parsed_path_and_singleprint =
+            metrics::ParseSavedModelPathAndSingleprint(path_and_singleprint);
+        if (!parsed_path_and_singleprint.ok()) {
+          throw MetricException(
+              parsed_path_and_singleprint.status().message().data());
+        }
+        return parsed_path_and_singleprint.value();
+      },
+      py::doc(
+          "Get tuple of `path` and `singleprint` values of "
+          "'/tensorflow/core/saved_model/write/path_and_singleprint' gauge."));
+
+  m.attr("kFingerprintFound") = metrics::kFingerprintFound;
+  m.attr("kFingerprintNotFound") = metrics::kFingerprintNotFound;
+  m.attr("kFingerprintError") = metrics::kFingerprintError;
+
+  m.def(
+      "GetFoundFingerprintOnLoad",
+      []() { return metrics::SavedModelFoundFingerprintOnLoad().value(); },
+      py::doc(
+          "Get value of "
+          "'/tensorflow/core/saved_model/found_fingerprint_on_load' gauge."));
+
+  m.def(
+      "SetFoundFingerprintOnLoad",
+      [](std::string found_status) {
+        if (found_status == metrics::kFingerprintFound ||
+            found_status == metrics::kFingerprintNotFound ||
+            found_status == metrics::kFingerprintError) {
+          metrics::SavedModelFoundFingerprintOnLoad().Set(found_status);
+        } else {
+          metrics::SavedModelFoundFingerprintOnLoad().Set("");
+        }
+      },
+      py::kw_only(), py::arg("found_status"),
+      py::doc("Set value of "
+              "'/tensorflow/core/saved_model/found_fingerprint_on_load' gauge "
+              "with 'found_status' if status is one of { \"FOUND\", "
+              "\"NOT_FOUND\", \"ERROR\" }."));
+
+  m.def(
+      "AddCheckpointReadDuration",
+      [](const char* api_label, double microseconds) {
+        metrics::CheckpointReadDuration(api_label).Add(microseconds);
+      },
+      py::kw_only(), py::arg("api_label"), py::arg("microseconds"),
+      py::doc("Add `microseconds` to the cell `api_label`for "
+              "'/tensorflow/core/checkpoint/read/read_durations'."));
+
+  m.def(
+      "GetCheckpointReadDurations",
+      [](const char* api_label) {
+        // This function is called sparingly in unit tests, so protobuf
+        // (de)-serialization round trip is not an issue.
+        return py::bytes(metrics::CheckpointReadDuration(api_label)
+                             .value()
+                             .SerializeAsString());
+      },
+      py::kw_only(), py::arg("api_label"),
+      py::doc("Get serialized HistogramProto of `api_label` cell for "
+              "'/tensorflow/core/checkpoint/read/read_durations'."));
+
+  m.def(
+      "AddCheckpointWriteDuration",
+      [](const char* api_label, double microseconds) {
+        metrics::CheckpointWriteDuration(api_label).Add(microseconds);
+      },
+      py::kw_only(), py::arg("api_label"), py::arg("microseconds"),
+      py::doc("Add `microseconds` to the cell `api_label` for "
+              "'/tensorflow/core/checkpoint/write/write_durations'."));
+
+  m.def(
+      "GetCheckpointWriteDurations",
+      [](const char* api_label) {
+        // This function is called sparingly, so protobuf (de)-serialization
+        // round trip is not an issue.
+        return py::bytes(metrics::CheckpointWriteDuration(api_label)
+                             .value()
+                             .SerializeAsString());
+      },
+      py::kw_only(), py::arg("api_label"),
+      py::doc("Get serialized HistogramProto of `api_label` cell for "
+              "'/tensorflow/core/checkpoint/write/write_durations'."));
+
+  m.def(
+      "AddAsyncCheckpointWriteDuration",
+      [](const char* api_label, double microseconds) {
+        metrics::AsyncCheckpointWriteDuration(api_label).Add(microseconds);
+      },
+      py::kw_only(), py::arg("api_label"), py::arg("microseconds"),
+      py::doc("Add `microseconds` to the cell `api_label` for "
+              "'/tensorflow/core/checkpoint/write/async_write_durations'."));
+
+  m.def(
+      "GetAsyncCheckpointWriteDurations",
+      [](const char* api_label) {
+        // This function is called sparingly, so protobuf (de)-serialization
+        // round trip is not an issue.
+        return py::bytes(metrics::AsyncCheckpointWriteDuration(api_label)
+                             .value()
+                             .SerializeAsString());
+      },
+      py::kw_only(), py::arg("api_label"),
+      py::doc("Get serialized HistogramProto of `api_label` cell for "
+              "'/tensorflow/core/checkpoint/write/async_write_durations'."));
+
+  m.def(
+      "AddTrainingTimeSaved",
+      [](const char* api_label, double microseconds) {
+        metrics::TrainingTimeSaved(api_label).IncrementBy(microseconds);
+      },
+      py::kw_only(), py::arg("api_label"), py::arg("microseconds"),
+      py::doc("Add `microseconds` to the cell `api_label` for "
+              "'/tensorflow/core/checkpoint/write/training_time_saved'."));
+
+  m.def(
+      "GetTrainingTimeSaved",
+      [](const char* api_label) {
+        return metrics::TrainingTimeSaved(api_label).value();
+      },
+      py::kw_only(), py::arg("api_label"),
+      py::doc("Get cell `api_label` for "
+              "'/tensorflow/core/checkpoint/write/training_time_saved'."));
+
+  m.def(
+      "CalculateFileSize",
+      [](const char* filename) {
+        Env* env = Env::Default();
+        uint64_t filesize = 0;
+        if (!env->GetFileSize(filename, &filesize).ok()) {
+          return (int64_t)-1;
+        }
+        // Convert to MB.
+        int64_t filesize_mb = filesize / 1000;
+        // Round to the nearest 100 MB.
+        // Smaller multiple.
+        int64_t a = (filesize_mb / 100) * 100;
+        // Larger multiple.
+        int64_t b = a + 100;
+        // Return closest of two.
+        return (filesize_mb - a > b - filesize_mb) ? b : a;
+      },
+      py::doc("Calculate filesize (MB) for `filename`, rounding to the nearest "
+              "100MB. Returns -1 if `filename` is invalid."));
+
+  m.def(
+      "RecordCheckpointSize",
+      [](const char* api_label, int64_t filesize) {
+        metrics::CheckpointSize(api_label, filesize).IncrementBy(1);
+      },
+      py::kw_only(), py::arg("api_label"), py::arg("filesize"),
+      py::doc("Increment the "
+              "'/tensorflow/core/checkpoint/write/checkpoint_size' counter for "
+              "cell (api_label, filesize) after writing a checkpoint."));
+
+  m.def(
+      "GetCheckpointSize",
+      [](const char* api_label, uint64_t filesize) {
+        return metrics::CheckpointSize(api_label, filesize).value();
+      },
+      py::kw_only(), py::arg("api_label"), py::arg("filesize"),
+      py::doc("Get cell (api_label, filesize) for "
+              "'/tensorflow/core/checkpoint/write/checkpoint_size'."));
+
+  m.def(
+      "GetShardingCallbackDuration",
+      []() { return metrics::ShardingCallbackDuration().value(); },
+      py::doc("Get value of "
+              "'/tensorflow/core/checkpoint/sharding/callback_duration'."));
+
+  m.def(
+      "AddShardingCallbackDuration",
+      [](int64_t callback_duration) {
+        metrics::ShardingCallbackDuration().IncrementBy(callback_duration);
+      },
+      py::kw_only(), py::arg("callback_duration"),
+      py::doc("Set value of "
+              "'/tensorflow/core/checkpoint/sharding/callback_duration'."));
+
+  m.def(
+      "GetNumCheckpointShardsWritten",
+      []() { return metrics::NumCheckpointShardsWritten().value(); },
+      py::doc("Get value of "
+              "'/tensorflow/core/checkpoint/sharding/"
+              "num_checkpoint_shards_written'."));
+
+  m.def(
+      "AddNumCheckpointShardsWritten",
+      [](int64_t num_shards) {
+        metrics::NumCheckpointShardsWritten().IncrementBy(num_shards);
+      },
+      py::kw_only(), py::arg("num_shards"),
+      py::doc("Set value of "
+              "'/tensorflow/core/checkpoint/sharding/"
+              "num_checkpoint_shards_written'."));
+
+  m.def(
+      "GetShardingCallbackDescription",
+      []() { return metrics::ShardingCallbackDescription().value(); },
+      py::doc("Get value of "
+              "'/tensorflow/core/checkpoint/sharding/callback_description'."));
+
+  m.def(
+      "SetShardingCallbackDescription",
+      [](std::string description) {
+        metrics::ShardingCallbackDescription().Set(description);
+      },
+      py::kw_only(), py::arg("description"),
+      py::doc("Set value of "
+              "'/tensorflow/core/checkpoint/sharding/callback_description'."));
+}
+
+}  // namespace python
+}  // namespace saved_model
+}  // namespace tensorflow

tensorflow/python/saved_model/pywrap_saved_model_metrics.h

@@ -0,0 +1,35 @@
+/* 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.
+==============================================================================*/
+
+// Defines the 'metrics' wrapper submodule.
+
+#ifndef TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_METRICS_H_
+#define TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_METRICS_H_
+
+// placeholder for index annotation headers
+#include "pybind11/pybind11.h"  // from @pybind11
+
+namespace tensorflow {
+namespace saved_model {
+namespace python {
+
+// Wraps the SM and Checkpoint Metrics API methods for exporting to Python.
+void DefineMetricsModule(pybind11::module main_module);
+
+}  // namespace python
+}  // namespace saved_model
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_PYTHON_SAVED_MODEL_PYWRAP_SAVED_MODEL_METRICS_H_

tensorflow/python/saved_model/pywrap_saved_model_metrics_test.py

@@ -0,0 +1,232 @@
+# 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 SavedModel and Checkpoint metrics Python bindings."""
+
+import os
+
+from tensorflow.core.framework import summary_pb2
+from tensorflow.core.framework import versions_pb2
+from tensorflow.core.protobuf import fingerprint_pb2
+from tensorflow.python.eager import test
+from tensorflow.python.saved_model.pywrap_saved_model import metrics
+
+
+class MetricsTest(test.TestCase):
+
+  def _get_histogram_proto(self, proto_bytes):
+    histogram_proto = summary_pb2.HistogramProto()
+    histogram_proto.ParseFromString(proto_bytes)
+    return histogram_proto
+
+  def _get_serialized_fingerprint_def(self):
+    return fingerprint_pb2.FingerprintDef(
+        saved_model_checksum=1,
+        graph_def_program_hash=2,
+        signature_def_hash=3,
+        saved_object_graph_hash=4,
+        checkpoint_hash=5,
+        version=versions_pb2.VersionDef(producer=6)).SerializeToString()
+
+  def test_SM_increment_write(self):
+    self.assertEqual(metrics.GetWrite(write_version="1"), 0)
+    metrics.IncrementWriteApi("foo")
+    self.assertEqual(metrics.GetWriteApi("foo"), 1)
+    metrics.IncrementWrite(write_version="1")
+    self.assertEqual(metrics.GetWrite(write_version="1"), 1)
+
+  def test_SM_increment_read(self):
+    self.assertEqual(metrics.GetRead(write_version="2"), 0)
+    metrics.IncrementReadApi("bar")
+    self.assertEqual(metrics.GetReadApi("bar"), 1)
+    metrics.IncrementRead(write_version="2")
+    self.assertEqual(metrics.GetRead(write_version="2"), 1)
+
+  def test_checkpoint_add_write_duration(self):
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetCheckpointWriteDurations(api_label="foo")).num, 0)
+
+    metrics.AddCheckpointWriteDuration(api_label="foo", microseconds=100)
+    metrics.AddCheckpointWriteDuration(api_label="foo", microseconds=200)
+
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetCheckpointWriteDurations(api_label="foo")).num, 2)
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetCheckpointWriteDurations(api_label="foo")).min, 100)
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetCheckpointWriteDurations(api_label="foo")).max, 200)
+
+  def test_async_checkpoint_add_write_duration(self):
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetAsyncCheckpointWriteDurations(api_label="foo")).num, 0)
+
+    metrics.AddAsyncCheckpointWriteDuration(api_label="foo", microseconds=20)
+    metrics.AddAsyncCheckpointWriteDuration(api_label="foo", microseconds=50)
+
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetAsyncCheckpointWriteDurations(api_label="foo")).num, 2)
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetAsyncCheckpointWriteDurations(api_label="foo")).min, 20)
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetAsyncCheckpointWriteDurations(api_label="foo")).max, 50)
+
+  def test_checkpoint_add_read_duration(self):
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetCheckpointReadDurations(api_label="bar")).num, 0)
+
+    metrics.AddCheckpointReadDuration(api_label="bar", microseconds=200)
+    metrics.AddCheckpointReadDuration(api_label="bar", microseconds=20000)
+
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetCheckpointReadDurations(api_label="bar")).num, 2)
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetCheckpointReadDurations(api_label="bar")).min, 200)
+    self.assertEqual(
+        self._get_histogram_proto(
+            metrics.GetCheckpointReadDurations(api_label="bar")).max, 20000)
+
+  def test_training_time_saved(self):
+    self.assertEqual(metrics.GetTrainingTimeSaved(api_label="baz"), 0)
+    metrics.AddTrainingTimeSaved(api_label="baz", microseconds=1000)
+    self.assertEqual(metrics.GetTrainingTimeSaved(api_label="baz"), 1000)
+
+  def test_checkpoint_size(self):
+    self.assertEqual(
+        metrics.GetCheckpointSize(api_label="baz", filesize=100), 0)
+    metrics.RecordCheckpointSize(api_label="baz", filesize=100)
+    metrics.RecordCheckpointSize(api_label="baz", filesize=100)
+    self.assertEqual(
+        metrics.GetCheckpointSize(api_label="baz", filesize=100), 2)
+
+  def test_filesize(self):
+    filename = os.path.join(self.get_temp_dir(), "test.txt")
+    with open(filename, "w") as file:
+      file.write("Hello! \n")
+    self.assertEqual(metrics.CalculateFileSize(filename), 0)
+
+  def test_invalid_file(self):
+    self.assertEqual(metrics.CalculateFileSize("not_a_file.txt"), -1)
+
+  def test_SM_read_fingerprint(self):
+    self.assertEqual(metrics.GetReadFingerprint(), "")
+    metrics.SetReadFingerprint(
+        fingerprint=self._get_serialized_fingerprint_def())
+    read_fingerprint = metrics.GetReadFingerprint()
+    self.assertIn('"saved_model_checksum" : 1', read_fingerprint)
+    self.assertIn('"graph_def_program_hash" : 2', read_fingerprint)
+    self.assertIn('"signature_def_hash" : 3', read_fingerprint)
+    self.assertIn('"saved_object_graph_hash" : 4', read_fingerprint)
+    self.assertIn('"checkpoint_hash" : 5', read_fingerprint)
+
+  def test_SM_write_fingerprint(self):
+    self.assertEqual(metrics.GetWriteFingerprint(), "")
+    metrics.SetWriteFingerprint(
+        fingerprint=self._get_serialized_fingerprint_def())
+    write_fingerprint = metrics.GetWriteFingerprint()
+    self.assertIn('"saved_model_checksum" : 1', write_fingerprint)
+    self.assertIn('"graph_def_program_hash" : 2', write_fingerprint)
+    self.assertIn('"signature_def_hash" : 3', write_fingerprint)
+    self.assertIn('"saved_object_graph_hash" : 4', write_fingerprint)
+    self.assertIn('"checkpoint_hash" : 5', write_fingerprint)
+
+  def test_SM_read_path(self):
+    self.assertEqual(metrics.GetReadPath(), "")
+    metrics.SetReadPath(saved_model_path="foo")
+    self.assertEqual(metrics.GetReadPath(), "foo")
+
+  def test_SM_write_path(self):
+    self.assertEqual(metrics.GetWritePath(), "")
+    metrics.SetWritePath(saved_model_path="foo")
+    self.assertEqual(metrics.GetWritePath(), "foo")
+
+  def test_SM_read_path_and_singleprint(self):
+    self.assertEqual(metrics.GetReadPathAndSingleprint(), ("", ""))
+    metrics.SetReadPathAndSingleprint(path="foo", singleprint="bar")
+    self.assertEqual(metrics.GetReadPathAndSingleprint(), ("foo", "bar"))
+
+  def test_SM_read_invalid_path_and_singleprint(self):
+    with self.assertRaises(metrics.MetricException) as excinfo:
+      metrics.SetReadPathAndSingleprint(path="", singleprint="bar")
+    self.assertRegex(str(excinfo.exception),
+                     "Invalid path_and_singleprint argument. Empty path.")
+
+    with self.assertRaises(metrics.MetricException) as excinfo:
+      metrics.SetReadPathAndSingleprint(path="foo", singleprint="")
+    self.assertRegex(
+        str(excinfo.exception),
+        "Invalid path_and_singleprint argument. Empty singleprint.")
+
+  def test_SM_write_path_and_singleprint(self):
+    self.assertEqual(metrics.GetWritePathAndSingleprint(), ("", ""))
+    metrics.SetWritePathAndSingleprint(path="foo", singleprint="bar")
+    self.assertEqual(metrics.GetWritePathAndSingleprint(), ("foo", "bar"))
+
+  def test_SM_write_invalid_path_and_singleprint(self):
+    with self.assertRaises(metrics.MetricException) as excinfo:
+      metrics.SetWritePathAndSingleprint(path="", singleprint="bar")
+    self.assertRegex(str(excinfo.exception),
+                     "Invalid path_and_singleprint argument. Empty path.")
+
+    with self.assertRaises(metrics.MetricException) as excinfo:
+      metrics.SetWritePathAndSingleprint(path="foo", singleprint="")
+    self.assertRegex(
+        str(excinfo.exception),
+        "Invalid path_and_singleprint argument. Empty singleprint.")
+
+  def test_SM_found_fingerprint_on_load(self):
+    metrics.SetFoundFingerprintOnLoad(found_status=metrics.kFingerprintFound)
+    self.assertEqual(metrics.GetFoundFingerprintOnLoad(), "FOUND")
+
+    metrics.SetFoundFingerprintOnLoad(found_status=metrics.kFingerprintNotFound)
+    self.assertEqual(metrics.GetFoundFingerprintOnLoad(), "NOT_FOUND")
+
+    metrics.SetFoundFingerprintOnLoad(found_status=metrics.kFingerprintError)
+    self.assertEqual(metrics.GetFoundFingerprintOnLoad(), "ERROR")
+
+  def test_invalid_SM_found_fingerprint_on_load(self):
+    metrics.SetFoundFingerprintOnLoad(found_status="absolute nonsense")
+    self.assertEqual(metrics.GetFoundFingerprintOnLoad(), "")
+
+    metrics.SetFoundFingerprintOnLoad(found_status="found")
+    self.assertEqual(metrics.GetFoundFingerprintOnLoad(), "")
+
+  def test_checkpoint_sharding_callback_duration(self):
+    self.assertEqual(metrics.GetShardingCallbackDuration(), 0)
+    metrics.AddShardingCallbackDuration(callback_duration=100)
+    self.assertEqual(metrics.GetShardingCallbackDuration(), 100)
+
+  def test_num_checkpoint_shards_written(self):
+    self.assertEqual(metrics.GetNumCheckpointShardsWritten(), 0)
+    metrics.AddNumCheckpointShardsWritten(num_shards=10)
+    self.assertEqual(metrics.GetNumCheckpointShardsWritten(), 10)
+
+  def test_sharding_callback_description(self):
+    self.assertEqual(metrics.GetShardingCallbackDescription(), "")
+    metrics.SetShardingCallbackDescription(description="foo")
+    self.assertEqual(metrics.GetShardingCallbackDescription(), "foo")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/registration/BUILD

@@ -0,0 +1,100 @@
+# Description:
+# TensorFlow SavedModel Registration.
+
+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 = "registration",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+    deps = [
+        ":registration_lib",
+    ],
+)
+
+py_library(
+    name = "registration_lib",
+    srcs = [
+        "registration.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:tf_inspect",
+        "@absl_py//absl/logging",
+    ],
+)
+
+tf_py_strict_test(
+    name = "registration_test",
+    srcs = ["registration_test.py"],
+    deps = [
+        ":registration",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/trackable:base",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "registration_saving_test",
+    srcs = ["registration_saving_test.py"],
+    deps = [
+        ":registration",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:io_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/trackable:autotrackable",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "test_util",
+    srcs = [
+        "test_util.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":registration_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "tf_registration_test",
+    srcs = ["tf_registration_test.py"],
+    data = [
+        "tf_checkpoint_saver_allowlist.txt",
+        "tf_serializable_allowlist.txt",
+    ],
+    tags = ["no_pip"],
+    deps = [
+        ":registration",
+        ":test_util",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/lib/io:file_io",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:resource_loader",
+    ],
+)

tensorflow/python/saved_model/registration/README.md

@@ -0,0 +1,173 @@
+# Registrations
+
+To configure SaveModel or checkpointing beyond the basic saving and loading
+steps [documentation TBD], registration is required.
+
+Currently, only TensorFlow-internal
+registrations are allowed, and must be added to the allowlist.
+
+* `tensorflow.python.saved_model.registration.register_tf_serializable`
+  * Allowlist: tf_serializable_allowlist.txt
+*  `tensorflow.python.saved_model.registration.register_tf_checkpoint_saver`
+  * Allowlist: tf_checkpoint_saver_allowlist.txt
+
+[TOC]
+
+## SavedModel serializable registration
+
+Custom objects must be registered in order to get the correct deserialization
+method when loading. The registered name of the class is saved to the proto.
+
+Keras already has a similar mechanism for registering serializables:
+[`tf.keras.utils.register_keras_serializable(package, name)`](https://www.tensorflow.org/api_docs/python/tf/keras/utils/register_keras_serializable).
+This has been imported to core TensorFlow:
+
+```python
+registration.register_serializable(package, name)
+registration.register_tf_serializable(name)  # If TensorFlow-internal.
+```
+
+*   package: The package that this class belongs to.
+*   name: The name of this class. The registered name that is saved in the proto
+    is "{package}.{name}" (for TensorFlow internal registration, the package
+    name is `tf`)
+
+## Checkpoint saver registration
+
+If `Trackables` share state or require complicated coordination between multiple
+`Trackables` (e.g. `DTensor`), then users may register a save and restore
+functions for these objects.
+
+```
+tf.saved_model.register_checkpoint_saver(
+    predicate, save_fn=None, restore_fn=None):
+```
+
+*   `predicate`: A function that returns `True` if a `Trackable` object should
+    be saved using the registered `save_fn` or `restore_fn`.
+*   `save_fn`: A python function or `tf.function` or `None`. If `None`, run the
+    default saving process which calls `Trackable._serialize_to_tensors`.
+*   `restore_fn`: A `tf.function` or `None`. If `None`, run the default
+    restoring process which calls `Trackable._restore_from_tensors`.
+
+**`save_fn` details**
+
+```
+@tf.function  # optional decorator
+def save_fn(trackables, file_prefix): -> List[shard filenames]
+```
+
+*   `trackables`: A dictionary of `{object_prefix: Trackable}`. The
+    object_prefix can be used as the object names, and uniquely identify each
+    `Trackable`. `trackables` is the filtered set of trackables that pass the
+    predicate.
+*   `file_prefix`: A string or string tensor of the checkpoint prefix.
+*   `shard filenames`: A list of filenames written using `io_ops.save_v2`, which
+    will be merged into the checkpoint data files. These should be prefixed by
+    `file_prefix`.
+
+This function can be a python function, in which case shard filenames can be an
+empty list (if the values are written without the `SaveV2` op).
+
+If this function is a `tf.function`, then the shards must be written using the
+SaveV2 op. This guarantees the checkpoint format is compatible with existing
+checkpoint readers and managers.
+
+**`restore_fn` details**
+
+```
+@tf.function  # required decorator
+def restore_fn(trackables, file_prefix): -> None
+```
+
+A `tf.function` with the spec:
+
+*   `trackables`: A dictionary of `{object_prefix: Trackable}`. The
+    `object_prefix` can be used as the object name, and uniquely identifies each
+    Trackable. The Trackable objects are the filtered results of the registered
+    predicate.
+*   `file_prefix`: A string or string tensor of the checkpoint prefix.
+
+**Why are restore functions required to be a `tf.function`?** The short answer
+is, the SavedModel format must maintain the invariant that SavedModel packages
+can be used for inference on any platform and language. SavedModel inference
+needs to be able to restore checkpointed values, so the restore function must be
+directly encoded into the SavedModel in the Graph. We also have security
+measures over FunctionDef and GraphDef, so users can check that the SavedModel
+will not run arbitrary code (a feature of `saved_model_cli`).
+
+## Example
+
+Below shows a `Stack` module that contains multiple `Parts` (a subclass of
+`tf.Variable`). When a `Stack` is saved to a checkpoint, the `Parts` are stacked
+together and a single entry in the checkpoint is created. The checkpoint value
+is restored to all of the `Parts` in the `Stack`.
+
+```
+@registration.register_serializable()
+class Part(resource_variable_ops.ResourceVariable):
+
+  def __init__(self, value):
+    self._init_from_args(value)
+
+  @classmethod
+  def _deserialize_from_proto(cls, **kwargs):
+    return cls([0, 0])
+
+
+@registration.register_serializable()
+class Stack(tracking.AutoTrackable):
+
+  def __init__(self, parts=None):
+    self.parts = parts
+
+  @def_function.function(input_signature=[])
+  def value(self):
+    return array_ops_stack.stack(self.parts)
+
+
+def get_tensor_slices(trackables):
+  tensor_names = []
+  shapes_and_slices = []
+  tensors = []
+  restored_trackables = []
+  for obj_prefix, obj in trackables.items():
+    if isinstance(obj, Part):
+      continue  # only save stacks
+    tensor_names.append(obj_prefix + "/value")
+    shapes_and_slices.append("")
+    x = obj.value()
+    with ops.device("/device:CPU:0"):
+      tensors.append(array_ops.identity(x))
+    restored_trackables.append(obj)
+
+  return tensor_names, shapes_and_slices, tensors, restored_trackables
+
+
+def save_stacks_and_parts(trackables, file_prefix):
+  """Save stack and part objects to a checkpoint shard."""
+  tensor_names, shapes_and_slices, tensors, _ = get_tensor_slices(trackables)
+  io_ops.save_v2(file_prefix, tensor_names, shapes_and_slices, tensors)
+  return file_prefix
+
+
+def restore_stacks_and_parts(trackables, merged_prefix):
+  tensor_names, shapes_and_slices, tensors, restored_trackables = (
+      get_tensor_slices(trackables))
+  dtypes = [t.dtype for t in tensors]
+  restored_tensors = io_ops.restore_v2(merged_prefix, tensor_names,
+                                       shapes_and_slices, dtypes)
+  for trackable, restored_tensor in zip(restored_trackables, restored_tensors):
+    expected_shape = trackable.value().get_shape()
+    restored_tensor = array_ops.reshape(restored_tensor, expected_shape)
+    parts = array_ops_stack.unstack(restored_tensor)
+    for part, restored_part in zip(trackable.parts, parts):
+      part.assign(restored_part)
+
+
+registration.register_checkpoint_saver(
+    name="stacks",
+    predicate=lambda x: isinstance(x, (Stack, Part)),
+    save_fn=save_stacks_and_parts,
+    restore_fn=restore_stacks_and_parts)
+```

tensorflow/python/saved_model/registration/__init__.py

@@ -0,0 +1,49 @@
+# 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.
+# ==============================================================================
+"""Utilities for registering the saving/loading steps for advanced objects."""
+
+from tensorflow.python.saved_model.registration.registration import get_registered_class
+from tensorflow.python.saved_model.registration.registration import get_registered_class_name
+from tensorflow.python.saved_model.registration.registration import get_registered_saver_name
+from tensorflow.python.saved_model.registration.registration import get_restore_function
+from tensorflow.python.saved_model.registration.registration import get_save_function
+from tensorflow.python.saved_model.registration.registration import get_strict_predicate_restore
+
+# These are currently an evolving feature. Use with care.
+from tensorflow.python.saved_model.registration.registration import register_checkpoint_saver
+from tensorflow.python.saved_model.registration.registration import register_serializable
+
+from tensorflow.python.saved_model.registration.registration import RegisteredSaver
+from tensorflow.python.saved_model.registration.registration import validate_restore_function
+
+
+def register_tf_serializable(name=None, predicate=None):
+  """See the docstring for `register_serializable`."""
+  return register_serializable(package="tf", name=name, predicate=predicate)
+
+
+def register_tf_checkpoint_saver(name=None,
+                                 predicate=None,
+                                 save_fn=None,
+                                 restore_fn=None,
+                                 strict_predicate_restore=True):
+  """See the docstring for `register_checkpoint_saver`."""
+  return register_checkpoint_saver(
+      package="tf",
+      name=name,
+      predicate=predicate,
+      save_fn=save_fn,
+      restore_fn=restore_fn,
+      strict_predicate_restore=strict_predicate_restore)

tensorflow/python/saved_model/registration/registration.py

@@ -0,0 +1,391 @@
+# 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.
+# ==============================================================================
+"""Serialization Registration for SavedModel.
+
+revived_types registration will be migrated to this infrastructure.
+
+See the Advanced saving section in go/savedmodel-configurability.
+This API is approved for TF internal use only.
+"""
+import collections
+import re
+
+from absl import logging
+
+from tensorflow.python.util import tf_inspect
+
+
+# Only allow valid file/directory characters
+_VALID_REGISTERED_NAME = re.compile(r"^[a-zA-Z0-9._-]+$")
+
+
+class _PredicateRegistry(object):
+  """Registry with predicate-based lookup.
+
+  See the documentation for `register_checkpoint_saver` and
+  `register_serializable` for reasons why predicates are required over a
+  class-based registry.
+
+  Since this class is used for global registries, each object must be registered
+  to unique names (an error is raised if there are naming conflicts). The lookup
+  searches the predicates in reverse order, so that later-registered predicates
+  are executed first.
+  """
+  __slots__ = ("_registry_name", "_registered_map", "_registered_predicates",
+               "_registered_names")
+
+  def __init__(self, name):
+    self._registry_name = name
+    # Maps registered name -> object
+    self._registered_map = {}
+    # Maps registered name -> predicate
+    self._registered_predicates = {}
+    # Stores names in the order of registration
+    self._registered_names = []
+
+  @property
+  def name(self):
+    return self._registry_name
+
+  def register(self, package, name, predicate, candidate):
+    """Registers a candidate object under the package, name and predicate."""
+    if not isinstance(package, str) or not isinstance(name, str):
+      raise TypeError(
+          f"The package and name registered to a {self.name} must be strings, "
+          f"got: package={type(package)}, name={type(name)}")
+    if not callable(predicate):
+      raise TypeError(
+          f"The predicate registered to a {self.name} must be callable, "
+          f"got: {type(predicate)}")
+    registered_name = package + "." + name
+    if not _VALID_REGISTERED_NAME.match(registered_name):
+      raise ValueError(
+          f"Invalid registered {self.name}. Please check that the package and "
+          f"name follow the regex '{_VALID_REGISTERED_NAME.pattern}': "
+          f"(package='{package}', name='{name}')")
+    if registered_name in self._registered_map:
+      raise ValueError(
+          f"The name '{registered_name}' has already been registered to a "
+          f"{self.name}. Found: {self._registered_map[registered_name]}")
+
+    self._registered_map[registered_name] = candidate
+    self._registered_predicates[registered_name] = predicate
+    self._registered_names.append(registered_name)
+
+  def lookup(self, obj):
+    """Looks up the registered object using the predicate.
+
+    Args:
+      obj: Object to pass to each of the registered predicates to look up the
+        registered object.
+    Returns:
+      The object registered with the first passing predicate.
+    Raises:
+      LookupError if the object does not match any of the predicate functions.
+    """
+    return self._registered_map[self.get_registered_name(obj)]
+
+  def name_lookup(self, registered_name):
+    """Looks up the registered object using the registered name."""
+    try:
+      return self._registered_map[registered_name]
+    except KeyError:
+      raise LookupError(f"The {self.name} registry does not have name "
+                        f"'{registered_name}' registered.")
+
+  def get_registered_name(self, obj):
+    for registered_name in reversed(self._registered_names):
+      predicate = self._registered_predicates[registered_name]
+      if predicate(obj):
+        return registered_name
+    raise LookupError(f"Could not find matching {self.name} for {type(obj)}.")
+
+  def get_predicate(self, registered_name):
+    try:
+      return self._registered_predicates[registered_name]
+    except KeyError:
+      raise LookupError(f"The {self.name} registry does not have name "
+                        f"'{registered_name}' registered.")
+
+  def get_registrations(self):
+    return self._registered_predicates
+
+_class_registry = _PredicateRegistry("serializable class")
+_saver_registry = _PredicateRegistry("checkpoint saver")
+
+
+def get_registered_class_name(obj):
+  try:
+    return _class_registry.get_registered_name(obj)
+  except LookupError:
+    return None
+
+
+def get_registered_class(registered_name):
+  try:
+    return _class_registry.name_lookup(registered_name)
+  except LookupError:
+    return None
+
+
+def register_serializable(package="Custom", name=None, predicate=None):  # pylint: disable=unused-argument
+  """Decorator for registering a serializable class.
+
+  THIS METHOD IS STILL EXPERIMENTAL AND MAY CHANGE AT ANY TIME.
+
+  Registered classes will be saved with a name generated by combining the
+  `package` and `name` arguments. When loading a SavedModel, modules saved with
+  this registered name will be created using the `_deserialize_from_proto`
+  method.
+
+  By default, only direct instances of the registered class will be saved/
+  restored with the `serialize_from_proto`/`deserialize_from_proto` methods. To
+  extend the registration to subclasses, use the `predicate argument`:
+
+  ```python
+  class A(tf.Module):
+    pass
+
+  register_serializable(
+      package="Example", predicate=lambda obj: isinstance(obj, A))(A)
+  ```
+
+  Args:
+    package: The package that this class belongs to.
+    name: The name to serialize this class under in this package. If None, the
+      class's name will be used.
+    predicate: An optional function that takes a single Trackable argument, and
+      determines whether that object should be serialized with this `package`
+      and `name`. The default predicate checks whether the object's type exactly
+      matches the registered class. Predicates are executed in the reverse order
+      that they are added (later registrations are checked first).
+
+  Returns:
+    A decorator that registers the decorated class with the passed names and
+    predicate.
+  """
+  def decorator(arg):
+    """Registers a class with the serialization framework."""
+    nonlocal predicate
+    if not tf_inspect.isclass(arg):
+      raise TypeError("Registered serializable must be a class: {}".format(arg))
+
+    class_name = name if name is not None else arg.__name__
+    if predicate is None:
+      predicate = lambda x: isinstance(x, arg)
+    _class_registry.register(package, class_name, predicate, arg)
+    return arg
+
+  return decorator
+
+
+RegisteredSaver = collections.namedtuple(
+    "RegisteredSaver", ["name", "predicate", "save_fn", "restore_fn"])
+_REGISTERED_SAVERS = {}
+_REGISTERED_SAVER_NAMES = []  # Stores names in the order of registration
+
+
+def register_checkpoint_saver(package="Custom",
+                              name=None,
+                              predicate=None,
+                              save_fn=None,
+                              restore_fn=None,
+                              strict_predicate_restore=True):
+  """Registers functions which checkpoints & restores objects with custom steps.
+
+  If you have a class that requires complicated coordination between multiple
+  objects when checkpointing, then you will need to register a custom saver
+  and restore function. An example of this is a custom Variable class that
+  splits the variable across different objects and devices, and needs to write
+  checkpoints that are compatible with different configurations of devices.
+
+  The registered save and restore functions are used in checkpoints and
+  SavedModel.
+
+  Please make sure you are familiar with the concepts in the [Checkpointing
+  guide](https://www.tensorflow.org/guide/checkpoint), and ops used to save the
+  V2 checkpoint format:
+
+  * io_ops.SaveV2
+  * io_ops.MergeV2Checkpoints
+  * io_ops.RestoreV2
+
+  **Predicate**
+
+  The predicate is a filter that will run on every `Trackable` object connected
+  to the root object. This function determines whether a `Trackable` should use
+  the registered functions.
+
+  Example: `lambda x: isinstance(x, CustomClass)`
+
+  **Custom save function**
+
+  This is how checkpoint saving works normally:
+  1. Gather all of the Trackables with saveable values.
+  2. For each Trackable, gather all of the saveable tensors.
+  3. Save checkpoint shards (grouping tensors by device) with SaveV2
+  4. Merge the shards with MergeCheckpointV2. This combines all of the shard's
+     metadata, and renames them to follow the standard shard pattern.
+
+  When a saver is registered, Trackables that pass the registered `predicate`
+  are automatically marked as having saveable values. Next, the custom save
+  function replaces steps 2 and 3 of the saving process. Finally, the shards
+  returned by the custom save function are merged with the other shards.
+
+  The save function takes in a dictionary of `Trackables` and a `file_prefix`
+  string. The function should save checkpoint shards using the SaveV2 op, and
+  list of the shard prefixes. SaveV2 is currently required to work a correctly,
+  because the code merges all of the returned shards, and the `restore_fn` will
+  only be given the prefix of the merged checkpoint. If you need to be able to
+  save and restore from unmerged shards, please file a feature request.
+
+  Specification and example of the save function:
+
+  ```
+  def save_fn(trackables, file_prefix):
+    # trackables: A dictionary mapping unique string identifiers to trackables
+    # file_prefix: A unique file prefix generated using the registered name.
+    ...
+    # Gather the tensors to save.
+    ...
+    io_ops.SaveV2(file_prefix, tensor_names, shapes_and_slices, tensors)
+    return file_prefix  # Returns a tensor or a list of string tensors
+  ```
+
+  The save function is executed before the unregistered save ops.
+
+  **Custom restore function**
+
+  Normal checkpoint restore behavior:
+  1. Gather all of the Trackables that have saveable values.
+  2. For each Trackable, get the names of the desired tensors to extract from
+     the checkpoint.
+  3. Use RestoreV2 to read the saved values, and pass the restored tensors to
+     the corresponding Trackables.
+
+  The custom restore function replaces steps 2 and 3.
+
+  The restore function also takes a dictionary of `Trackables` and a
+  `merged_prefix` string. The `merged_prefix` is different from the
+  `file_prefix`, since it contains the renamed shard paths. To read from the
+  merged checkpoint, you must use `RestoreV2(merged_prefix, ...)`.
+
+  Specification:
+
+  ```
+  def restore_fn(trackables, merged_prefix):
+    # trackables: A dictionary mapping unique string identifiers to Trackables
+    # merged_prefix: File prefix of the merged shard names.
+
+    restored_tensors = io_ops.restore_v2(
+        merged_prefix, tensor_names, shapes_and_slices, dtypes)
+    ...
+    # Restore the checkpoint values for the given Trackables.
+  ```
+
+  The restore function is executed after the non-registered restore ops.
+
+  Args:
+    package: Optional, the package that this class belongs to.
+    name: (Required) The name of this saver, which is saved to the checkpoint.
+      When a checkpoint is restored, the name and package are used to find the
+      the matching restore function. The name and package are also used to
+      generate a unique file prefix that is passed to the save_fn.
+    predicate: (Required) A function that returns a boolean indicating whether a
+      `Trackable` object should be checkpointed with this function. Predicates
+      are executed in the reverse order that they are added (later registrations
+      are checked first).
+    save_fn: (Required) A function that takes a dictionary of trackables and a
+      file prefix as the arguments, writes the checkpoint shards for the given
+      Trackables, and returns the list of shard prefixes.
+    restore_fn: (Required) A function that takes a dictionary of trackables and
+      a file prefix as the arguments and restores the trackable values.
+    strict_predicate_restore: If this is `True` (default), then an error will be
+      raised if the predicate fails during checkpoint restoration. If this is
+      `True`, checkpoint restoration will skip running the restore function.
+      This value is generally set to `False` when the predicate does not pass on
+      the Trackables after being saved/loaded from SavedModel.
+
+  Raises:
+    ValueError: if the package and name are already registered.
+  """
+  if not callable(save_fn):
+    raise TypeError(f"The save_fn must be callable, got: {type(save_fn)}")
+  if not callable(restore_fn):
+    raise TypeError(f"The restore_fn must be callable, got: {type(restore_fn)}")
+
+  _saver_registry.register(package, name, predicate, (save_fn, restore_fn,
+                                                      strict_predicate_restore))
+
+
+def get_registered_saver_name(trackable):
+  """Returns the name of the registered saver to use with Trackable."""
+  try:
+    return _saver_registry.get_registered_name(trackable)
+  except LookupError:
+    return None
+
+
+def get_save_function(registered_name):
+  """Returns save function registered to name."""
+  return _saver_registry.name_lookup(registered_name)[0]
+
+
+def get_restore_function(registered_name):
+  """Returns restore function registered to name."""
+  return _saver_registry.name_lookup(registered_name)[1]
+
+
+def get_strict_predicate_restore(registered_name):
+  """Returns if the registered restore can be ignored if the predicate fails."""
+  try:
+    return _saver_registry.name_lookup(registered_name)[2]
+  except LookupError:
+    logging.warning(
+        "Registered saver %s was not found when restoring checkpoints.",
+        registered_name,
+    )
+    return False  # Return false as the default if the name isn't registered.
+
+
+def validate_restore_function(trackable, registered_name):
+  """Validates whether the trackable can be restored with the saver.
+
+  When using a checkpoint saved with a registered saver, that same saver must
+  also be also registered when loading. The name of that saver is saved to the
+  checkpoint and set in the `registered_name` arg.
+
+  Args:
+    trackable: A `Trackable` object.
+    registered_name: String name of the expected registered saver. This argument
+      should be set using the name saved in a checkpoint.
+
+  Raises:
+    ValueError if the saver could not be found, or if the predicate associated
+      with the saver does not pass.
+  """
+  try:
+    _saver_registry.name_lookup(registered_name)
+  except LookupError:
+    raise ValueError(
+        f"Error when restoring object {trackable} from checkpoint. This "
+        "object was saved using a registered saver named "
+        f"'{registered_name}', but this saver cannot be found in the "
+        "current context.")
+  if not _saver_registry.get_predicate(registered_name)(trackable):
+    raise ValueError(
+        f"Object {trackable} was saved with the registered saver named "
+        f"'{registered_name}'. However, this saver cannot be used to restore the "
+        "object because the predicate does not pass.")

tensorflow/python/saved_model/registration/registration_saving_test.py

@@ -0,0 +1,419 @@
+# 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 saving with registered Trackable classes and checkpoint functions."""
+
+import os
+import tempfile
+
+from absl.testing import parameterized
+
+from google.protobuf import wrappers_pb2
+from tensorflow.python.checkpoint import checkpoint as util
+from tensorflow.python.client import session
+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 errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import io_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import gfile
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import loader
+from tensorflow.python.saved_model import registration
+from tensorflow.python.saved_model import save
+from tensorflow.python.trackable import autotrackable
+
+
+@registration.register_serializable()
+class Part(resource_variable_ops.ResourceVariable):
+
+  def __init__(self, value):
+    self._init_from_args(value)
+
+  @classmethod
+  def _deserialize_from_proto(cls, **kwargs):
+    return cls([0, 0])
+
+  def _export_to_saved_model_graph(self, object_map, tensor_map, **kwargs):
+    p = Part(array_ops.zeros(self.shape, self.dtype))
+    object_map[self] = p
+    tensor_map[self.handle] = p.handle
+    return [self.handle]
+
+
+@registration.register_serializable()
+class Stack(autotrackable.AutoTrackable):
+
+  def __init__(self, parts=None):
+    self.parts = parts
+
+  @def_function.function(input_signature=[])
+  def value(self):
+    return array_ops_stack.stack(self.parts)
+
+
+def get_tensor_slices(trackables):
+  tensor_names = []
+  shapes_and_slices = []
+  tensors = []
+  restored_trackables = []
+  for obj_prefix, obj in trackables.items():
+    if isinstance(obj, Part):
+      continue  # only save stacks
+    tensor_names.append(obj_prefix + "/value")
+    shapes_and_slices.append("")
+    x = obj.value()
+    with ops.device("/device:CPU:0"):
+      tensors.append(array_ops.identity(x))
+    restored_trackables.append(obj)
+
+  return tensor_names, shapes_and_slices, tensors, restored_trackables
+
+
+def save_stacks_and_parts(trackables, file_prefix):
+  """Save stack and part objects to a checkpoint shard."""
+  tensor_names, shapes_and_slices, tensors, _ = get_tensor_slices(trackables)
+  io_ops.save_v2(file_prefix, tensor_names, shapes_and_slices, tensors)
+  return file_prefix
+
+
+def restore_stacks_and_parts(trackables, merged_prefix):
+  tensor_names, shapes_and_slices, tensors, restored_trackables = (
+      get_tensor_slices(trackables))
+  dtypes = [t.dtype for t in tensors]
+  restored_tensors = io_ops.restore_v2(merged_prefix, tensor_names,
+                                       shapes_and_slices, dtypes)
+  for trackable, restored_tensor in zip(restored_trackables, restored_tensors):
+    expected_shape = trackable.value().get_shape()
+    restored_tensor = array_ops.reshape(restored_tensor, expected_shape)
+    parts = array_ops_stack.unstack(restored_tensor)
+    for part, restored_part in zip(trackable.parts, parts):
+      part.assign(restored_part)
+
+
+registration.register_checkpoint_saver(
+    name="stacks",
+    predicate=lambda x: isinstance(x, (Stack, Part)),
+    save_fn=save_stacks_and_parts,
+    restore_fn=restore_stacks_and_parts)
+
+
+def cycle(obj, cycles, signatures=None, options=None):
+  to_save = obj
+  for _ in range(cycles):
+    path = tempfile.mkdtemp(prefix=test.get_temp_dir())
+    # If available, we'll run the save and restore preferring the GPU. This
+    # just makes sure we aren't throwing errors and have enough
+    # device("CPU") blocks to satisfy the placer.
+    with test_util.use_gpu():
+      save.save(to_save, path, signatures, options=options)
+      loaded = load.load(path)
+      signatures = loaded.signatures
+    to_save = loaded
+  return loaded
+
+
+@parameterized.named_parameters(
+    dict(testcase_name="ReloadOnce", cycles=1),
+    dict(testcase_name="ReloadTwice", cycles=2),
+    dict(testcase_name="ReloadThrice", cycles=3))
+class SavedModelTest(test.TestCase, parameterized.TestCase):
+
+  def test_registered_serializable(self, cycles):
+
+    @registration.register_serializable(name=f"SaveAndLoad{cycles}")
+    class Module(autotrackable.AutoTrackable):
+
+      def __init__(self, name="module"):
+        self.v = variables.Variable(1.)
+        self.name = name
+
+      def _serialize_to_proto(self, **unused_kwargs):
+        return wrappers_pb2.StringValue(value=self.name)
+
+      @classmethod
+      def _deserialize_from_proto(cls, proto, **unused_kwargs):
+        if proto.Is(wrappers_pb2.StringValue.DESCRIPTOR):
+          unpacked = wrappers_pb2.StringValue()
+          proto.Unpack(unpacked)
+          return cls(name=unpacked.value)
+        raise AssertionError(
+            "Did not receive proto of correct type during deserialization. "
+            f"Expected type {wrappers_pb2.StringValue.DESCRIPTOR.full_name}, "
+            f"got {proto.TypeName()}")
+
+    m = Module("a")
+    m.v.assign(5)
+    loaded = cycle(m, cycles)
+    self.assertIsInstance(loaded, Module)
+    self.assertEqual(5, loaded.v.numpy())
+    self.assertEqual("a", loaded.name)
+
+  def test_none_proto(self, cycles):
+
+    @registration.register_serializable(name=f"NoneProto{cycles}")
+    class Module(autotrackable.AutoTrackable):
+
+      def __init__(self, name="module"):
+        self.v = variables.Variable(1.)
+        self.name = name
+
+      # Leave _serialize_to_proto as the default (returns `None`).
+
+      @classmethod
+      def _deserialize_from_proto(cls, proto, **unused_kwargs):
+        self.assertEqual(proto.ByteSize(), 0)
+        return cls("deserialized")
+
+    m = Module("a")
+    m.v.assign(5)
+    loaded = cycle(m, cycles)
+    self.assertIsInstance(loaded, Module)
+    self.assertEqual(5, loaded.v.numpy())
+    self.assertEqual("deserialized", loaded.name)
+
+  def test_deserialization_dependencies(self, cycles):
+    @registration.register_serializable(name=f"Dependency{cycles}")
+    class Module(autotrackable.AutoTrackable):
+
+      def __init__(self, v=None):
+        self.v = v if v is not None else variables.Variable(1.)
+
+      def _deserialization_dependencies(self, children):
+        del children  # Unused.
+        return {"v": self.v}
+
+      @classmethod
+      def _deserialize_from_proto(cls, dependencies, **unused_kwargs):
+        self.assertIn("v", dependencies)
+        return cls(v=dependencies["v"])
+
+    m = Module()
+    m.v.assign(5)
+    loaded = cycle(m, cycles)
+    self.assertIsInstance(loaded, Module)
+    self.assertEqual(5, loaded.v.numpy())
+
+  def test_registered_saver(self, cycles):
+    p1 = Part([1, 4])
+    p2 = Part([2, 5])
+    p3 = Part([3, 6])
+    s = Stack([p1, p2, p3])
+    loaded = cycle(s, cycles)
+    self.assertAllEqual(s.value(), loaded.value())
+
+
+class SingleCycleTest(test.TestCase):
+
+  @test_util.deprecated_graph_mode_only
+  def test_registered_saver_fails_in_saved_model_graph_mode(self):
+    with context.eager_mode():
+      p1 = Part([1, 4])
+      p2 = Part([2, 5])
+      p3 = Part([3, 6])
+      s = Stack([p1, p2, p3])
+      save_dir = os.path.join(self.get_temp_dir(), "save_dir")
+      save.save(s, save_dir)
+
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        "registered checkpoint saver is not supported in graph mode"):
+      load.load(save_dir)
+
+  def test_registered_saver_checkpoint(self):
+    p1 = Part([1, 4])
+    p2 = Part([2, 5])
+    p3 = Part([3, 6])
+    s = Stack([p1, p2, p3])
+    s2 = Stack([p3, p1, p2])
+
+    expected_value_s = s.value()
+    expected_value_s2 = s2.value()
+
+    ckpt_path = os.path.join(self.get_temp_dir(), "ckpt")
+    util.Checkpoint(s=s, s2=s2).write(ckpt_path)
+
+    del s, s2, p1, p2, p3
+
+    restore_s = Stack([Part([0, 0]) for _ in range(3)])
+    util.Checkpoint(s=restore_s).read(ckpt_path).expect_partial()
+    self.assertAllEqual(expected_value_s, restore_s.value())
+    util.Checkpoint(s2=restore_s).read(ckpt_path).expect_partial()
+    self.assertAllEqual(expected_value_s2, restore_s.value())
+
+  def test_compatible_with_v1_savedmodel(self):
+    p1 = Part([1, 4])
+    p2 = Part([2, 5])
+    p3 = Part([3, 6])
+    s = Stack([p1, p2, p3])
+    save_path = os.path.join(self.get_temp_dir(), "savedmodel")
+
+    @def_function.function(input_signature=[])
+    def serve():
+      return {"value": s.value()}
+
+    exported_value = serve()["value"]
+
+    save.save(s, save_path, signatures=serve)
+    with ops.Graph().as_default(), session.Session() as sess:
+      metagraph = loader.load(sess, ["serve"], save_path)
+      value_output = metagraph.signature_def["serving_default"].outputs["value"]
+      self.assertAllEqual(exported_value, sess.run(value_output.name))
+
+  def test_non_strict_predicate(self):
+    class NonStrictPredicateClass(autotrackable.AutoTrackable):
+      pass
+    registration.register_checkpoint_saver(
+        name="NonStrictPredicate",
+        predicate=lambda x: isinstance(x, NonStrictPredicateClass),
+        save_fn=lambda **kwargs: [],
+        restore_fn=lambda **kwargs: None,
+        strict_predicate_restore=False)
+
+    root = NonStrictPredicateClass()
+    ckpt_path = os.path.join(self.get_temp_dir(), "ckpt")
+    util.Checkpoint(root).write(ckpt_path)
+
+    root2 = autotrackable.AutoTrackable()
+    # This should run without throwing an error.
+    util.Checkpoint(root2).read(ckpt_path)
+
+  def test_strict_predicate(self):
+    class StrictPredicateClass(autotrackable.AutoTrackable):
+      pass
+    registration.register_checkpoint_saver(
+        name="StrictPredicate",
+        predicate=lambda x: isinstance(x, StrictPredicateClass),
+        save_fn=lambda **kwargs: [],
+        restore_fn=lambda **kwargs: None,
+        strict_predicate_restore=True)
+
+    root = StrictPredicateClass()
+    ckpt_path = os.path.join(self.get_temp_dir(), "ckpt")
+    util.Checkpoint(root).write(ckpt_path)
+
+    root2 = autotrackable.AutoTrackable()
+    with self.assertRaisesRegex(ValueError, "saver cannot be used"):
+      util.Checkpoint(root2).read(ckpt_path)
+
+  def test_registered_saver_is_called_before_save_after_load(self):
+    if not context.executing_eagerly():
+      self.skipTest("This test must run under eager mode.")
+
+    class RestoreClass(autotrackable.AutoTrackable):
+      pass
+    def save_fn(trackables, file_prefix):
+      del trackables  # Unused.
+      # Check that directory is empty
+      files = gfile.ListDirectory(os.path.dirname(file_prefix.numpy()))
+      self.assertEmpty(files)
+
+    def restore_fn(trackables, merged_prefix):
+      del merged_prefix  # Unused.
+      root = next(trackables.values())
+      self.assertEqual(root.v.numpy(), 123)
+
+    registration.register_checkpoint_saver(
+        name="OptionalRestore",
+        predicate=lambda x: isinstance(x, RestoreClass),
+        save_fn=save_fn,
+        restore_fn=restore_fn)
+
+    root = RestoreClass()
+    root.v = variables.Variable(123.0)
+
+    ckpt_path = os.path.join(self.get_temp_dir(), "ckpt")
+    util.Checkpoint(root).write(ckpt_path)
+
+  def test_migration_backwards_compatibility(self):
+    # Tests that objects migrated to using the advanced saver registration can
+    # use pre-migration checkpoints.
+
+    class NoRegisteredSaver(autotrackable.AutoTrackable):
+
+      def __init__(self, name):
+        self.name = name
+
+      def _serialize_to_tensors(self):
+        return {"name": constant_op.constant(self.name)}
+
+    class RegisteredSaver(autotrackable.AutoTrackable):
+
+      def __init__(self, name):
+        self.name = name
+
+    def _get_tensors(trackables, append_name=True):
+      tensor_names = []
+      shapes_and_slices = []
+      tensors = []
+      restored_trackables = []
+      for obj_prefix, obj in trackables.items():
+        tensor_names.append(obj_prefix + "name" if append_name else obj_prefix)
+        shapes_and_slices.append("")
+        tensors.append(constant_op.constant(obj.name))
+        restored_trackables.append(obj)
+      return tensor_names, shapes_and_slices, tensors, restored_trackables
+
+    def save_fn(trackables, file_prefix):
+      tensor_names, shapes_and_slices, tensors, _ = _get_tensors(trackables)
+      io_ops.save_v2(file_prefix, tensor_names, shapes_and_slices, tensors)
+      return file_prefix
+
+    def restore_fn(trackables, merged_prefix):
+      tensor_names, shapes_and_slices, tensors, restored_trackables = (
+          _get_tensors(trackables))
+      dtypes = [t.dtype for t in tensors]
+      try:
+        restored_tensors = io_ops.restore_v2(merged_prefix, tensor_names,
+                                             shapes_and_slices, dtypes)
+      except errors_impl.NotFoundError:
+        # If a NotFoundError is caught, then it means that the checkpoint
+        # was written prior to the saver registration migration.
+        tensor_names, shapes_and_slices, tensors, restored_trackables = (
+            _get_tensors(trackables, append_name=False))
+        restored_tensors = io_ops.restore_v2(merged_prefix, tensor_names,
+                                             shapes_and_slices, dtypes)
+      for trackable, name_tensor in zip(restored_trackables, restored_tensors):
+        trackable.name = name_tensor
+
+    registration.register_checkpoint_saver(
+        name="MigratedSaver",
+        predicate=lambda x: isinstance(x, RegisteredSaver),
+        save_fn=save_fn,
+        restore_fn=restore_fn,
+    )
+
+    before = NoRegisteredSaver("before")
+    after = RegisteredSaver("after")
+    before_ckpt_path = os.path.join(self.get_temp_dir(), "before_ckpt")
+    util.Checkpoint(before).write(before_ckpt_path)
+
+    after_ckpt = util.Checkpoint(after)
+    after_ckpt_path = os.path.join(self.get_temp_dir(), "after_ckpt")
+    after_ckpt.write(after_ckpt_path)
+
+    # Try loading the pre-migrated checkpoint to the migrated object.
+    after_ckpt.read(before_ckpt_path)
+    self.assertEqual(b"before", self.evaluate(after.name))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/registration/registration_test.py

@@ -0,0 +1,220 @@
+# 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 the registration functions.
+
+For integration tests that use save and load functions, see
+registration_saving_test.py.
+"""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import test
+from tensorflow.python.saved_model import registration
+from tensorflow.python.trackable import base
+
+
+@registration.register_serializable()
+class RegisteredClass(base.Trackable):
+  pass
+
+
+@registration.register_serializable(name="Subclass")
+class RegisteredSubclass(RegisteredClass):
+  pass
+
+
+@registration.register_serializable(package="testing")
+class CustomPackage(base.Trackable):
+  pass
+
+
+@registration.register_serializable(package="testing", name="name")
+class CustomPackageAndName(base.Trackable):
+  pass
+
+
+class SerializableRegistrationTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      (RegisteredClass, "Custom.RegisteredClass"),
+      (RegisteredSubclass, "Custom.Subclass"),
+      (CustomPackage, "testing.CustomPackage"),
+      (CustomPackageAndName, "testing.name"),
+  ])
+  def test_registration(self, expected_cls, expected_name):
+    obj = expected_cls()
+    self.assertEqual(registration.get_registered_class_name(obj), expected_name)
+    self.assertIs(
+        registration.get_registered_class(expected_name), expected_cls)
+
+  def test_get_invalid_name(self):
+    self.assertIsNone(registration.get_registered_class("invalid name"))
+
+  def test_get_unregistered_class(self):
+
+    class NotRegistered(base.Trackable):
+      pass
+
+    no_register = NotRegistered
+    self.assertIsNone(registration.get_registered_class_name(no_register))
+
+  def test_duplicate_registration(self):
+
+    @registration.register_serializable()
+    class Duplicate(base.Trackable):
+      pass
+
+    dup = Duplicate()
+    self.assertEqual(
+        registration.get_registered_class_name(dup), "Custom.Duplicate")
+    # Registrations with different names are ok.
+    registration.register_serializable(package="duplicate")(Duplicate)
+    # Registrations are checked in reverse order.
+    self.assertEqual(
+        registration.get_registered_class_name(dup), "duplicate.Duplicate")
+    # Both names should resolve to the same class.
+    self.assertIs(
+        registration.get_registered_class("Custom.Duplicate"), Duplicate)
+    self.assertIs(
+        registration.get_registered_class("duplicate.Duplicate"), Duplicate)
+
+    # Registrations of the same name fails
+    with self.assertRaisesRegex(ValueError, "already been registered"):
+      registration.register_serializable(
+          package="testing", name="CustomPackage")(
+              Duplicate)
+
+  def test_register_non_class_fails(self):
+    obj = RegisteredClass()
+    with self.assertRaisesRegex(TypeError, "must be a class"):
+      registration.register_serializable()(obj)
+
+  def test_register_bad_predicate_fails(self):
+    with self.assertRaisesRegex(TypeError, "must be callable"):
+      registration.register_serializable(predicate=0)(RegisteredClass)
+
+  def test_predicate(self):
+
+    class Predicate(base.Trackable):
+
+      def __init__(self, register_this):
+        self.register_this = register_this
+
+    registration.register_serializable(
+        name="RegisterThisOnlyTrue",
+        predicate=lambda x: isinstance(x, Predicate) and x.register_this)(
+            Predicate)
+
+    a = Predicate(True)
+    b = Predicate(False)
+    self.assertEqual(
+        registration.get_registered_class_name(a),
+        "Custom.RegisterThisOnlyTrue")
+    self.assertIsNone(registration.get_registered_class_name(b))
+
+    registration.register_serializable(
+        name="RegisterAllPredicate",
+        predicate=lambda x: isinstance(x, Predicate))(
+            Predicate)
+
+    self.assertEqual(
+        registration.get_registered_class_name(a),
+        "Custom.RegisterAllPredicate")
+    self.assertEqual(
+        registration.get_registered_class_name(b),
+        "Custom.RegisterAllPredicate")
+
+
+class CheckpointSaverRegistrationTest(test.TestCase):
+
+  def test_invalid_registration(self):
+    with self.assertRaisesRegex(TypeError, "must be string"):
+      registration.register_checkpoint_saver(
+          package=None,
+          name="test",
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(TypeError, "must be string"):
+      registration.register_checkpoint_saver(
+          name=None,
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(ValueError,
+                                "Invalid registered checkpoint saver."):
+      registration.register_checkpoint_saver(
+          package="package",
+          name="t/est",
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(ValueError,
+                                "Invalid registered checkpoint saver."):
+      registration.register_checkpoint_saver(
+          package="package",
+          name="t/est",
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(
+        TypeError,
+        "The predicate registered to a checkpoint saver must be callable"
+    ):
+      registration.register_checkpoint_saver(
+          name="test",
+          predicate=None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(TypeError, "The save_fn must be callable"):
+      registration.register_checkpoint_saver(
+          name="test",
+          predicate=lambda: None,
+          save_fn=None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(TypeError, "The restore_fn must be callable"):
+      registration.register_checkpoint_saver(
+          name="test",
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=None)
+
+  def test_registration(self):
+    registration.register_checkpoint_saver(
+        package="Testing",
+        name="test_predicate",
+        predicate=lambda x: hasattr(x, "check_attr"),
+        save_fn=lambda: "save",
+        restore_fn=lambda: "restore")
+    x = base.Trackable()
+    self.assertIsNone(registration.get_registered_saver_name(x))
+
+    x.check_attr = 1
+    saver_name = registration.get_registered_saver_name(x)
+    self.assertEqual(saver_name, "Testing.test_predicate")
+
+    self.assertEqual(registration.get_save_function(saver_name)(), "save")
+    self.assertEqual(registration.get_restore_function(saver_name)(), "restore")
+
+    registration.validate_restore_function(x, "Testing.test_predicate")
+    with self.assertRaisesRegex(ValueError, "saver cannot be found"):
+      registration.validate_restore_function(x, "Invalid.name")
+    x2 = base.Trackable()
+    with self.assertRaisesRegex(ValueError, "saver cannot be used"):
+      registration.validate_restore_function(x2, "Testing.test_predicate")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/registration/test_util.py

@@ -0,0 +1,25 @@
+# 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.
+# ==============================================================================
+"""Utils for testing registered objects."""
+from tensorflow.python.saved_model.registration import registration as registration_lib
+# pylint: disable=protected-access
+
+
+def get_all_registered_serializables():
+  return registration_lib._class_registry.get_registrations()
+
+
+def get_all_registered_checkpoint_savers():
+  return registration_lib._saver_registry.get_registrations()

tensorflow/python/saved_model/registration/tf_checkpoint_saver_allowlist.txt

@@ -0,0 +1,4 @@
+# Please enter your registered saver names, separated by new lines.
+
+tf.TestDummySaver
+tf.TPUEmbeddingCallback

tensorflow/python/saved_model/registration/tf_registration_test.py

@@ -0,0 +1,110 @@
+# 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 TensorFlow registered objects.
+
+New or deleted registrations must be approved by the Saved Model team.
+"""
+import os
+
+import tensorflow as tf
+
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.platform import resource_loader
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import registration
+from tensorflow.python.saved_model.registration import test_util
+
+SERIALIZABLE_ALLOWLIST = os.path.join(resource_loader.get_data_files_path(),
+                                      "tf_serializable_allowlist.txt")
+CHECKPOINT_SAVER_ALLOWLIST = os.path.join(resource_loader.get_data_files_path(),
+                                          "tf_checkpoint_saver_allowlist.txt")
+
+# call TPUEmbedding to load the file and run its registration
+_ = tf.tpu.experimental.embedding.TPUEmbedding
+# call save to load the file and run its registration
+_ = tf.saved_model.save
+
+
+@registration.register_tf_serializable()
+class NotInAllowlistExample(tf.Variable):
+  pass
+
+
+registration.register_tf_checkpoint_saver(
+    name="TestDummySaver",
+    predicate=lambda: False,
+    save_fn=lambda: None,
+    restore_fn=lambda: None)
+
+registration.register_tf_checkpoint_saver(
+    name="NotInAllowlistExample",
+    predicate=lambda: False,
+    save_fn=lambda: None,
+    restore_fn=lambda: None)
+
+
+class AllowlistTest(test.TestCase):
+
+  def _test_allowlist(self, allowlist_file, registrations):
+    allowlist = set([
+        s.strip()
+        for s in file_io.read_file_to_string(allowlist_file).splitlines()
+        if s.strip() and not s.startswith("#")
+    ])
+    registered_names = set(registrations)
+
+    missing_from_allowlist = registered_names - allowlist
+    self.assertIn("tf.NotInAllowlistExample", missing_from_allowlist)
+    missing_from_allowlist.remove("tf.NotInAllowlistExample")
+
+    if missing_from_allowlist:
+      msg = ("[NEEDS ATTENTION] Registered names found that were not added to "
+             "the allowlist. Add the following names to the list:\n\t" +
+             "\n\t".join(missing_from_allowlist))
+    else:
+      msg = "[OK] All registered names have been added to the allowlist.  ✓"
+
+    msg += "\n\n"
+
+    missing_registered_names = allowlist - registered_names
+    if missing_registered_names:
+      msg += ("[NEEDS ATTENTION] Some names were found in the allowlist that "
+              "are not registered in TensorFlow. This could mean that a "
+              "registration was removed from the codebase. If this was "
+              "intended, please remove the following from the allowlist:\n\t" +
+              "\n\t".join(missing_registered_names))
+    else:
+      msg += ("[OK] All allowlisted names are registered in the Tensorflow "
+              "library. ✓")
+
+    if missing_from_allowlist or missing_registered_names:
+      raise AssertionError(
+          "Error found in the registration allowlist.\nPlease update the "
+          "allowlist at .../tensorflow/python/saved_model/registration/"
+          f"{os.path.basename(allowlist_file)}.\n\n" + msg +
+          "\n\nAfter making changes, request approval from "
+          " tf-saved-model-owners@.")
+
+  def test_checkpoint_savers(self):
+    self._test_allowlist(CHECKPOINT_SAVER_ALLOWLIST,
+                         test_util.get_all_registered_checkpoint_savers())
+
+  def test_serializables(self):
+    self._test_allowlist(SERIALIZABLE_ALLOWLIST,
+                         test_util.get_all_registered_serializables())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/registration/tf_serializable_allowlist.txt

@@ -0,0 +1,5 @@
+# Please enter your registered saver names, separated by new lines.
+# These names are not the same as the exported tf API symbols.
+
+tf.StaticHashTable
+tf.TrackableConstant

tensorflow/python/saved_model/revived_types.py

@@ -0,0 +1,249 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Handles types registrations for tf.saved_model.load."""
+
+import operator
+
+from tensorflow.core.framework import versions_pb2
+from tensorflow.core.protobuf import saved_object_graph_pb2
+from tensorflow.python.trackable import data_structures
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("__internal__.saved_model.load.VersionedTypeRegistration", v1=[])
+class VersionedTypeRegistration(object):
+  """Holds information about one version of a revived type."""
+
+  def __init__(self, object_factory, version, min_producer_version,
+               min_consumer_version, bad_consumers=None, setter=setattr):
+    """Identify a revived type version.
+
+    Args:
+      object_factory: A callable which takes a SavedUserObject proto and returns
+        a trackable object. Dependencies are added later via `setter`.
+      version: An integer, the producer version of this wrapper type. When
+        making incompatible changes to a wrapper, add a new
+        `VersionedTypeRegistration` with an incremented `version`. The most
+        recent version will be saved, and all registrations with a matching
+        identifier will be searched for the highest compatible version to use
+        when loading.
+      min_producer_version: The minimum producer version number required to use
+        this `VersionedTypeRegistration` when loading a proto.
+      min_consumer_version: `VersionedTypeRegistration`s with a version number
+        less than `min_consumer_version` will not be used to load a proto saved
+        with this object. `min_consumer_version` should be set to the lowest
+        version number which can successfully load protos saved by this
+        object. If no matching registration is available on load, the object
+        will be revived with a generic trackable type.
+
+        `min_consumer_version` and `bad_consumers` are a blunt tool, and using
+        them will generally break forward compatibility: previous versions of
+        TensorFlow will revive newly saved objects as opaque trackable
+        objects rather than wrapped objects. When updating wrappers, prefer
+        saving new information but preserving compatibility with previous
+        wrapper versions. They are, however, useful for ensuring that
+        previously-released buggy wrapper versions degrade gracefully rather
+        than throwing exceptions when presented with newly-saved SavedModels.
+      bad_consumers: A list of consumer versions which are incompatible (in
+        addition to any version less than `min_consumer_version`).
+      setter: A callable with the same signature as `setattr` to use when adding
+        dependencies to generated objects.
+    """
+    self.setter = setter
+    self.identifier = None  # Set after registration
+    self._object_factory = object_factory
+    self.version = version
+    self._min_consumer_version = min_consumer_version
+    self._min_producer_version = min_producer_version
+    if bad_consumers is None:
+      bad_consumers = []
+    self._bad_consumers = bad_consumers
+
+  def to_proto(self):
+    """Create a SavedUserObject proto."""
+    # For now wrappers just use dependencies to save their state, so the
+    # SavedUserObject doesn't depend on the object being saved.
+    # TODO(allenl): Add a wrapper which uses its own proto.
+    return saved_object_graph_pb2.SavedUserObject(
+        identifier=self.identifier,
+        version=versions_pb2.VersionDef(
+            producer=self.version,
+            min_consumer=self._min_consumer_version,
+            bad_consumers=self._bad_consumers))
+
+  def from_proto(self, proto):
+    """Recreate a trackable object from a SavedUserObject proto."""
+    return self._object_factory(proto)
+
+  def should_load(self, proto):
+    """Checks if this object should load the SavedUserObject `proto`."""
+    if proto.identifier != self.identifier:
+      return False
+    if self.version < proto.version.min_consumer:
+      return False
+    if proto.version.producer < self._min_producer_version:
+      return False
+    for bad_version in proto.version.bad_consumers:
+      if self.version == bad_version:
+        return False
+    return True
+
+
+# string identifier -> (predicate, [VersionedTypeRegistration])
+_REVIVED_TYPE_REGISTRY = {}
+_TYPE_IDENTIFIERS = []
+
+
+@tf_export("__internal__.saved_model.load.register_revived_type", v1=[])
+def register_revived_type(identifier, predicate, versions):
+  """Register a type for revived objects.
+
+  Args:
+    identifier: A unique string identifying this class of objects.
+    predicate: A Boolean predicate for this registration. Takes a
+      trackable object as an argument. If True, `type_registration` may be
+      used to save and restore the object.
+    versions: A list of `VersionedTypeRegistration` objects.
+  """
+  # Keep registrations in order of version. We always use the highest matching
+  # version (respecting the min consumer version and bad consumers).
+  versions.sort(key=lambda reg: reg.version, reverse=True)
+  if not versions:
+    raise AssertionError("Need at least one version of a registered type.")
+  version_numbers = set()
+  for registration in versions:
+    # Copy over the identifier for use in generating protos
+    registration.identifier = identifier
+    if registration.version in version_numbers:
+      raise AssertionError(
+          f"Got multiple registrations with version {registration.version} for "
+          f"type {identifier}.")
+    version_numbers.add(registration.version)
+
+  _REVIVED_TYPE_REGISTRY[identifier] = (predicate, versions)
+  _TYPE_IDENTIFIERS.append(identifier)
+
+
+def serialize(obj):
+  """Create a SavedUserObject from a trackable object."""
+  for identifier in _TYPE_IDENTIFIERS:
+    predicate, versions = _REVIVED_TYPE_REGISTRY[identifier]
+    if predicate(obj):
+      # Always uses the most recent version to serialize.
+      return versions[0].to_proto()
+  return None
+
+
+def deserialize(proto):
+  """Create a trackable object from a SavedUserObject proto.
+
+  Args:
+    proto: A SavedUserObject to deserialize.
+
+  Returns:
+    A tuple of (trackable, assignment_fn) where assignment_fn has the same
+    signature as setattr and should be used to add dependencies to
+    `trackable` when they are available.
+  """
+  _, type_registrations = _REVIVED_TYPE_REGISTRY.get(
+      proto.identifier, (None, None))
+  if type_registrations is not None:
+    for type_registration in type_registrations:
+      if type_registration.should_load(proto):
+        return (type_registration.from_proto(proto), type_registration.setter)
+  return None
+
+
+@tf_export("__internal__.saved_model.load.registered_identifiers", v1=[])
+def registered_identifiers():
+  """Return all the current registered revived object identifiers.
+
+  Returns:
+    A set of strings.
+  """
+  return _REVIVED_TYPE_REGISTRY.keys()
+
+
+@tf_export("__internal__.saved_model.load.get_setter", v1=[])
+def get_setter(proto):
+  """Gets the registered setter function for the SavedUserObject proto.
+
+  See VersionedTypeRegistration for info about the setter function.
+
+  Args:
+    proto: SavedUserObject proto
+
+  Returns:
+    setter function
+  """
+  _, type_registrations = _REVIVED_TYPE_REGISTRY.get(
+      proto.identifier, (None, None))
+  if type_registrations is not None:
+    for type_registration in type_registrations:
+      if type_registration.should_load(proto):
+        return type_registration.setter
+  return None
+
+
+register_revived_type(
+    "trackable_dict_wrapper",
+    # pylint: disable=protected-access
+    lambda obj: isinstance(obj, data_structures._DictWrapper),
+    versions=[
+        VersionedTypeRegistration(
+            # Standard dependencies are enough to reconstruct the trackable
+            # items in dictionaries, so we don't need to save any extra
+            # information.
+            # pylint: disable=protected-access
+            object_factory=lambda proto: data_structures._DictWrapper({}),
+            version=1,
+            min_producer_version=1,
+            min_consumer_version=1,
+            setter=operator.setitem,
+        )
+    ],
+)
+
+
+register_revived_type(
+    "trackable_list_wrapper",
+    lambda obj: isinstance(obj, data_structures.ListWrapper),
+    versions=[
+        VersionedTypeRegistration(
+            object_factory=lambda proto: data_structures.ListWrapper([]),
+            version=1,
+            min_producer_version=1,
+            min_consumer_version=1,
+            setter=data_structures.set_list_item,
+        )
+    ],
+)
+
+
+# Revive tuples as lists so we can append any dependencies during loading.
+register_revived_type(
+    "trackable_tuple_wrapper",
+    # pylint: disable=protected-access
+    lambda obj: isinstance(obj, data_structures._TupleWrapper),
+    versions=[
+        VersionedTypeRegistration(
+            object_factory=lambda proto: data_structures.ListWrapper([]),
+            version=1,
+            min_producer_version=1,
+            min_consumer_version=1,
+            setter=data_structures.set_tuple_item,
+        )
+    ],
+)

tensorflow/python/saved_model/revived_types_test.py

@@ -0,0 +1,106 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for revived type matching."""
+
+from tensorflow.core.framework import versions_pb2
+from tensorflow.core.protobuf import saved_object_graph_pb2
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import revived_types
+from tensorflow.python.trackable import autotrackable
+
+
+class CustomTestClass(autotrackable.AutoTrackable):
+
+  def __init__(self, version):
+    self.version = version
+
+
+revived_types.register_revived_type(
+    "test_type",
+    lambda obj: isinstance(obj, CustomTestClass),
+    versions=[
+        revived_types.VersionedTypeRegistration(
+            object_factory=lambda _: CustomTestClass(1),
+            version=1, min_producer_version=1,
+            min_consumer_version=1),
+        revived_types.VersionedTypeRegistration(
+            object_factory=lambda _: CustomTestClass(2),
+            version=2, min_producer_version=2, min_consumer_version=1),
+        revived_types.VersionedTypeRegistration(
+            object_factory=lambda _: CustomTestClass(3),
+            version=3, min_producer_version=3, min_consumer_version=2),
+        revived_types.VersionedTypeRegistration(
+            object_factory=lambda _: CustomTestClass(4),
+            version=4, min_producer_version=4, min_consumer_version=2,
+            bad_consumers=[3]),
+    ]
+)
+
+
+class RegistrationMatchingTest(test.TestCase):
+
+  def test_save_typecheck(self):
+    self.assertIs(revived_types.serialize(autotrackable.AutoTrackable()), None)
+
+  def test_load_identifier_not_found(self):
+    nothing_matches = revived_types.deserialize(
+        saved_object_graph_pb2.SavedUserObject(
+            identifier="_unregistered_type",
+            version=versions_pb2.VersionDef(
+                producer=1,
+                min_consumer=1,
+                bad_consumers=[])))
+    self.assertIs(nothing_matches, None)
+
+  def test_most_recent_version_saved(self):
+    serialized = revived_types.serialize(CustomTestClass(None))
+    self.assertEqual([3], serialized.version.bad_consumers)
+    deserialized, _ = revived_types.deserialize(serialized)
+    self.assertIsInstance(deserialized, CustomTestClass)
+    self.assertEqual(4, deserialized.version)
+
+  def test_min_consumer_version(self):
+    nothing_matches = revived_types.deserialize(
+        saved_object_graph_pb2.SavedUserObject(
+            identifier="test_type",
+            version=versions_pb2.VersionDef(
+                producer=5,
+                min_consumer=5,
+                bad_consumers=[])))
+    self.assertIs(nothing_matches, None)
+
+  def test_bad_versions(self):
+    deserialized, _ = revived_types.deserialize(
+        saved_object_graph_pb2.SavedUserObject(
+            identifier="test_type",
+            version=versions_pb2.VersionDef(
+                producer=5,
+                min_consumer=1,
+                bad_consumers=[4, 3])))
+    self.assertEqual(2, deserialized.version)
+
+  def test_min_producer_version(self):
+    deserialized, _ = revived_types.deserialize(
+        saved_object_graph_pb2.SavedUserObject(
+            identifier="test_type",
+            version=versions_pb2.VersionDef(
+                producer=3,
+                min_consumer=0,
+                bad_consumers=[])))
+    self.assertEqual(3, deserialized.version)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/save_context.py

@@ -0,0 +1,66 @@
+# 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.
+# ==============================================================================
+"""Context for building SavedModel."""
+
+import contextlib
+import threading
+
+
+class SaveContext(threading.local):
+  """A context for building a graph of SavedModel."""
+
+  def __init__(self):
+    super(SaveContext, self).__init__()
+    self._in_save_context = False
+    self._options = None
+
+  def options(self):
+    if not self.in_save_context():
+      raise ValueError("Not in a SaveContext.")
+    return self._options
+
+  def enter_save_context(self, options):
+    self._in_save_context = True
+    self._options = options
+
+  def exit_save_context(self):
+    self._in_save_context = False
+    self._options = None
+
+  def in_save_context(self):
+    return self._in_save_context
+
+_save_context = SaveContext()
+
+
+@contextlib.contextmanager
+def save_context(options):
+  if in_save_context():
+    raise ValueError("Already in a SaveContext.")
+  _save_context.enter_save_context(options)
+  try:
+    yield
+  finally:
+    _save_context.exit_save_context()
+
+
+def in_save_context():
+  """Returns whether under a save context."""
+  return _save_context.in_save_context()
+
+
+def get_save_options():
+  """Returns the save options if under a save context."""
+  return _save_context.options()

tensorflow/python/saved_model/save_context_test.py

@@ -0,0 +1,83 @@
+# 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.
+# ==============================================================================
+"""Test for SaveContext."""
+
+import threading
+
+from tensorflow.python.eager import test
+from tensorflow.python.saved_model import save_context
+from tensorflow.python.saved_model import save_options
+
+
+class SaveContextTest(test.TestCase):
+
+  def test_multi_thread(self):
+    self.assertFalse(save_context.in_save_context())
+    with self.assertRaisesRegex(ValueError, 'Not in a SaveContext'):
+      save_context.get_save_options()
+
+    options = save_options.SaveOptions(save_debug_info=True)
+    with save_context.save_context(options):
+      self.assertTrue(save_context.in_save_context())
+      self.assertTrue(save_context.get_save_options().save_debug_info)
+
+      entered_context_in_thread = threading.Event()
+      continue_thread = threading.Event()
+
+      def thread_fn():
+        self.assertFalse(save_context.in_save_context())
+        with self.assertRaisesRegex(ValueError, 'Not in a SaveContext'):
+          save_context.get_save_options()
+
+        options = save_options.SaveOptions(save_debug_info=False)
+        with save_context.save_context(options):
+          self.assertTrue(save_context.in_save_context())
+          # save_debug_info has a different value in this thread.
+          self.assertFalse(save_context.get_save_options().save_debug_info)
+          entered_context_in_thread.set()
+          continue_thread.wait()
+
+        self.assertFalse(save_context.in_save_context())
+        with self.assertRaisesRegex(ValueError, 'Not in a SaveContext'):
+          save_context.get_save_options()
+
+      t = threading.Thread(target=thread_fn)
+      t.start()
+
+      entered_context_in_thread.wait()
+      # Another thread shouldn't affect this thread.
+      self.assertTrue(save_context.in_save_context())
+      self.assertTrue(save_context.get_save_options().save_debug_info)
+
+      continue_thread.set()
+      t.join()
+      # Another thread exiting SaveContext shouldn't affect this thread.
+      self.assertTrue(save_context.in_save_context())
+      self.assertTrue(save_context.get_save_options().save_debug_info)
+
+    self.assertFalse(save_context.in_save_context())
+    with self.assertRaisesRegex(ValueError, 'Not in a SaveContext'):
+      save_context.get_save_options()
+
+  def test_enter_multiple(self):
+    options = save_options.SaveOptions()
+    with self.assertRaisesRegex(ValueError, 'Already in a SaveContext'):
+      with save_context.save_context(options):
+        with save_context.save_context(options):
+          pass
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/saved_model/save_options.py

@@ -0,0 +1,245 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Options for saving SavedModels."""
+
+import enum
+
+from tensorflow.python.checkpoint.sharding import sharding_util
+from tensorflow.python.util import compat
+from tensorflow.python.util.tf_export import tf_export
+
+
+is_oss = True  # Updated by copybara.
+
+
+@tf_export("saved_model.experimental.VariablePolicy")
+class VariablePolicy(enum.Enum):
+  """Enum defining options for variable handling when saving.
+
+  NONE
+    No policy applied: Distributed variables are saved as one variable, with no
+    device attached.
+
+  SAVE_VARIABLE_DEVICES
+    When saving variables, also save their device assignment.
+    This is useful if one wants to hardcode devices in saved models, but it also
+    makes them non-portable if soft device placement is disabled (more details
+    in `tf.config.set_soft_device_placement`). This is currently not
+    fully supported by `saved_model.load`, and is mainly intended to be used
+    when one will be reading the saved model at a lower API level. In the
+    example below, the graph saved by the call to `saved_model.save` will have
+    the variable devices correctly specified:
+    ```python
+    exported = tf.train.Checkpoint()
+    with tf.device('/GPU:0'):
+      exported.x_gpu = tf.Variable(1.0)
+    with tf.device('/CPU:0'):
+      exported.x_cpu = tf.Variable(1.0)
+    tf.saved_model.save(exported, export_dir,
+        options = tf.saved_model.SaveOptions(
+            experimental_variable_policy=
+              tf.saved_model.experimental.VariablePolicy.SAVE_VARIABLE_DEVICES))
+    ```
+    Distributed variables are still saved as one variable under this policy.
+
+  EXPAND_DISTRIBUTED_VARIABLES
+    Distributed variables will be saved with information about their components,
+    allowing for their restoration on load. Also, the saved graph will contain
+    references to those variables. This is useful when one wants to use the
+    model for training in environments where the original distribution strategy
+    is not available.
+  """
+
+  NONE = None
+
+  SAVE_VARIABLE_DEVICES = "save_variable_devices"
+
+  EXPAND_DISTRIBUTED_VARIABLES = "expand_distributed_variables"
+
+  def _save_variable_devices(self):
+    """Checks whether variable devices should be saved."""
+    return self != VariablePolicy.NONE
+
+  def _expand_distributed_variables(self):
+    """Checks whether distributed variables should be expanded."""
+    return self == VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES
+
+  @staticmethod
+  def from_obj(obj):
+    """Tries to convert `obj` to a VariablePolicy instance."""
+    if obj is None:
+      return VariablePolicy.NONE
+    if isinstance(obj, VariablePolicy):
+      return obj
+    key = str(obj).lower()
+    for policy in VariablePolicy:
+      if key == policy.value:
+        return policy
+    raise ValueError(f"Received invalid VariablePolicy value: {obj}.")
+
+
+@tf_export("saved_model.SaveOptions")
+class SaveOptions:
+  """Options for saving to SavedModel.
+
+  This function may be used in the `options` argument in functions that
+  save a SavedModel (`tf.saved_model.save`, `tf.keras.models.save_model`).
+  """
+
+  # Define object attributes in __slots__ for improved memory and performance.
+  __slots__ = (
+      "namespace_whitelist",
+      "save_debug_info",
+      "function_aliases",
+      "experimental_debug_stripper",
+      "experimental_io_device",
+      "experimental_variable_policy",
+      "experimental_custom_gradients",
+      "experimental_image_format",
+      "experimental_skip_saver",
+      "experimental_sharding_callback",
+      "extra_tags",
+  )
+
+  def __init__(
+      self,
+      namespace_whitelist=None,
+      save_debug_info=False,
+      function_aliases=None,
+      experimental_debug_stripper=False,
+      experimental_io_device=None,
+      experimental_variable_policy=None,
+      experimental_custom_gradients=True,
+      experimental_image_format=False,
+      experimental_skip_saver=False,
+      experimental_sharding_callback=None,
+      extra_tags=None,
+  ):
+    """Creates an object that stores options for SavedModel saving.
+
+    Args:
+      namespace_whitelist: List of strings containing op namespaces to whitelist
+        when saving a model. Saving an object that uses namespaced ops must
+        explicitly add all namespaces to the whitelist. The namespaced ops must
+        be registered into the framework when loading the SavedModel. If no
+        whitelist is provided, all namespaced ops will be allowed.
+      save_debug_info: Boolean indicating whether debug information is saved. If
+        True, then a debug/saved_model_debug_info.pb file will be written with
+        the contents of a GraphDebugInfo binary protocol buffer containing stack
+        trace information for all ops and functions that are saved.
+      function_aliases: Python dict. Mapping from string to object returned by
+        @tf.function. A single tf.function can generate many ConcreteFunctions.
+        If a downstream tool wants to refer to all concrete functions generated
+        by a single tf.function you can use the `function_aliases` argument to
+        store a map from the alias name to all concrete function names. E.g. >>>
+        class Adder(tf.Module): ...   @tf.function ...   def double(self, x):
+        ...     return x + x  >>> model = Adder() >>>
+        model.double.get_concrete_function( ...   tf.TensorSpec(shape=[],
+        dtype=tf.float32, name="float_input")) >>>
+        model.double.get_concrete_function( ...   tf.TensorSpec(shape=[],
+        dtype=tf.string, name="string_input"))  >>> options =
+        tf.saved_model.SaveOptions( ...   function_aliases={'double':
+        model.double}) >>> tf.saved_model.save(model, '/tmp/adder',
+        options=options)
+      experimental_debug_stripper: bool. If set to True, this strips the debug
+        nodes from the graph, from both the nodes and the function defs. Note
+        that this currently only strips the `Assert` nodes from the graph and
+        converts them into `NoOp`s instead.
+      experimental_io_device: string. Applies in a distributed setting.
+        Tensorflow device to use to access the filesystem. If `None` (default)
+        then for each variable the filesystem is accessed from the CPU:0 device
+        of the host where that variable is assigned. If specified, the
+        filesystem is instead accessed from that device for all variables.  This
+        is for example useful if you want to save to a local directory, such as
+        "/tmp" when running in a distributed setting. In that case pass a device
+        for the host where the "/tmp" directory is accessible.
+      experimental_variable_policy: The policy to apply to variables when
+        saving. This is either a `saved_model.experimental.VariablePolicy` enum
+        instance or one of its value strings (case is not important). See that
+        enum documentation for details. A value of `None` corresponds to the
+        default policy.
+      experimental_custom_gradients: Boolean. When True, will save traced
+        gradient functions for the functions decorated by `tf.custom_gradient`.
+        Defaults to `True`.
+      experimental_image_format: New (highly) experimental format that is
+        capable of saving models larger than the 2GB protobuf limit. Enabling
+        this option will likely break compatibility with downstream consumers.
+        This option is currently disabled in OSS.
+      experimental_skip_saver: If True, will prevent SavedModel from creating
+        its native checkpointing ops - this is for models that do not use
+        SavedModel's native checkpointing functionality to avoid the costs
+        associated with creating and serializing those ops.
+      experimental_sharding_callback: `tf.train.experimental.ShardingCallback`.
+        A pre-made or custom callback that determines how checkpoints are
+        sharded on disk. Pre-made callback options are
+        `tf.train.experimental.ShardByDevicePolicy` and
+        `tf.train.experimental.MaxShardSizePolicy`. You may also write a custom
+        callback, see `tf.train.experimental.ShardingCallback`.
+      extra_tags: Extra tags to be saved with the MetaGraph in the SavedModel.
+    """
+    self.namespace_whitelist = _validate_namespace_whitelist(
+        namespace_whitelist
+    )
+    self.save_debug_info = save_debug_info
+    self.function_aliases = function_aliases if function_aliases else dict()
+    self.experimental_custom_gradients = experimental_custom_gradients
+    self.experimental_debug_stripper = experimental_debug_stripper
+    self.experimental_io_device = experimental_io_device
+    self.experimental_variable_policy = VariablePolicy.from_obj(
+        experimental_variable_policy
+    )
+    self.experimental_skip_saver = experimental_skip_saver
+
+    # TODO(b/277279153): Enable image format in OSS after proto splitter is
+    #  public.
+    if experimental_image_format and is_oss:
+      raise ValueError(
+          "The option `experimental_image_format` is disabled in OSS."
+      )
+    self.experimental_image_format = experimental_image_format
+
+    if experimental_sharding_callback is not None:
+      if not isinstance(
+          experimental_sharding_callback, sharding_util.ShardingCallback
+      ):
+        raise ValueError(
+            "The experimental_sharding_callback checkpoint option"
+            "must be of type ShardingCallback. The option provided"
+            f"was of type {type(experimental_sharding_callback)}."
+        )
+    self.experimental_sharding_callback = experimental_sharding_callback
+    self.extra_tags = extra_tags
+
+
+def _validate_namespace_whitelist(namespace_whitelist):
+  """Validates namespace whitelist argument."""
+  if namespace_whitelist is None:
+    return None
+  if not isinstance(namespace_whitelist, list):
+    raise TypeError(
+        "`namespace_whitelist` must be a list of strings. Got: "
+        f"{namespace_whitelist} with type "
+        f"{type(namespace_whitelist)}."
+    )
+
+  processed = []
+  for namespace in namespace_whitelist:
+    if not isinstance(namespace, str):
+      raise ValueError(
+          "Whitelisted namespace must be a string. Got: "
+          f"{namespace} of type {type(namespace)}."
+      )
+    processed.append(compat.as_str(namespace))
+  return processed

tensorflow/python/saved_model/saved_model.py

@@ -0,0 +1,36 @@
+# 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.
+# ==============================================================================
+"""Convenience functions to save a model.
+"""
+
+
+# pylint: disable=unused-import
+from tensorflow.python.saved_model import builder
+from tensorflow.python.saved_model import constants
+from tensorflow.python.saved_model import loader
+from tensorflow.python.saved_model import main_op
+from tensorflow.python.saved_model import method_name_updater
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import tag_constants
+from tensorflow.python.saved_model import utils
+from tensorflow.python.saved_model.fingerprinting import Fingerprint
+from tensorflow.python.saved_model.fingerprinting import read_fingerprint
+from tensorflow.python.saved_model.load import load
+from tensorflow.python.saved_model.save import save
+# pylint: enable=unused-import
+# pylint: disable=wildcard-import
+from tensorflow.python.saved_model.simple_save import *
+# pylint: enable=wildcard-import

tensorflow/python/saved_model/signature_constants.py

@@ -0,0 +1,143 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Signature constants for SavedModel save and restore operations.
+
+"""
+from tensorflow.python.util.tf_export import tf_export
+
+
+# Key in the signature def map for `default` serving signatures. The default
+# signature is used in inference requests where a specific signature was not
+# specified.
+DEFAULT_SERVING_SIGNATURE_DEF_KEY = "serving_default"
+tf_export(
+    "saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY",
+    v1=[
+        "saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY",
+        "saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY"
+    ],
+).export_constant(__name__, "DEFAULT_SERVING_SIGNATURE_DEF_KEY")
+
+################################################################################
+# Classification API constants.
+
+# Classification inputs.
+CLASSIFY_INPUTS = "inputs"
+tf_export(
+    "saved_model.CLASSIFY_INPUTS",
+    v1=[
+        "saved_model.CLASSIFY_INPUTS",
+        "saved_model.signature_constants.CLASSIFY_INPUTS"
+    ]).export_constant(__name__, "CLASSIFY_INPUTS")
+
+# Classification method name used in a SignatureDef.
+CLASSIFY_METHOD_NAME = "tensorflow/serving/classify"
+tf_export(
+    "saved_model.CLASSIFY_METHOD_NAME",
+    v1=[
+        "saved_model.CLASSIFY_METHOD_NAME",
+        "saved_model.signature_constants.CLASSIFY_METHOD_NAME"
+    ]).export_constant(__name__, "CLASSIFY_METHOD_NAME")
+
+# Classification classes output.
+CLASSIFY_OUTPUT_CLASSES = "classes"
+tf_export(
+    "saved_model.CLASSIFY_OUTPUT_CLASSES",
+    v1=[
+        "saved_model.CLASSIFY_OUTPUT_CLASSES",
+        "saved_model.signature_constants.CLASSIFY_OUTPUT_CLASSES"
+    ]).export_constant(__name__, "CLASSIFY_OUTPUT_CLASSES")
+
+# Classification scores output.
+CLASSIFY_OUTPUT_SCORES = "scores"
+tf_export(
+    "saved_model.CLASSIFY_OUTPUT_SCORES",
+    v1=[
+        "saved_model.CLASSIFY_OUTPUT_SCORES",
+        "saved_model.signature_constants.CLASSIFY_OUTPUT_SCORES"
+    ]).export_constant(__name__, "CLASSIFY_OUTPUT_SCORES")
+
+################################################################################
+# Prediction API constants.
+
+# Predict inputs.
+PREDICT_INPUTS = "inputs"
+tf_export(
+    "saved_model.PREDICT_INPUTS",
+    v1=[
+        "saved_model.PREDICT_INPUTS",
+        "saved_model.signature_constants.PREDICT_INPUTS"
+    ]).export_constant(__name__, "PREDICT_INPUTS")
+
+# Prediction method name used in a SignatureDef.
+PREDICT_METHOD_NAME = "tensorflow/serving/predict"
+tf_export(
+    "saved_model.PREDICT_METHOD_NAME",
+    v1=[
+        "saved_model.PREDICT_METHOD_NAME",
+        "saved_model.signature_constants.PREDICT_METHOD_NAME"
+    ]).export_constant(__name__, "PREDICT_METHOD_NAME")
+
+# Predict outputs.
+PREDICT_OUTPUTS = "outputs"
+tf_export(
+    "saved_model.PREDICT_OUTPUTS",
+    v1=[
+        "saved_model.PREDICT_OUTPUTS",
+        "saved_model.signature_constants.PREDICT_OUTPUTS"
+    ]).export_constant(__name__, "PREDICT_OUTPUTS")
+
+################################################################################
+# Regression API constants.
+
+# Regression inputs.
+REGRESS_INPUTS = "inputs"
+tf_export(
+    "saved_model.REGRESS_INPUTS",
+    v1=[
+        "saved_model.REGRESS_INPUTS",
+        "saved_model.signature_constants.REGRESS_INPUTS"
+    ]).export_constant(__name__, "REGRESS_INPUTS")
+
+# Regression method name used in a SignatureDef.
+REGRESS_METHOD_NAME = "tensorflow/serving/regress"
+tf_export(
+    "saved_model.REGRESS_METHOD_NAME",
+    v1=[
+        "saved_model.REGRESS_METHOD_NAME",
+        "saved_model.signature_constants.REGRESS_METHOD_NAME"
+    ]).export_constant(__name__, "REGRESS_METHOD_NAME")
+
+# Regression outputs.
+REGRESS_OUTPUTS = "outputs"
+tf_export(
+    "saved_model.REGRESS_OUTPUTS",
+    v1=[
+        "saved_model.REGRESS_OUTPUTS",
+        "saved_model.signature_constants.REGRESS_OUTPUTS"
+    ]).export_constant(__name__, "REGRESS_OUTPUTS")
+
+################################################################################
+# LINT.IfChange
+# Train/Eval API constants.
+# Not exported while export_all_saved_models is experimental.
+DEFAULT_TRAIN_SIGNATURE_DEF_KEY = "train"
+
+DEFAULT_EVAL_SIGNATURE_DEF_KEY = "eval"
+
+SUPERVISED_TRAIN_METHOD_NAME = "tensorflow/supervised/training"
+
+SUPERVISED_EVAL_METHOD_NAME = "tensorflow/supervised/eval"
+# LINT.ThenChange(//tensorflow/python/keras/saving/utils_v1/unexported_constants.py)

tensorflow/python/saved_model/signature_def_utils.py

@@ -0,0 +1,29 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""SignatureDef utility functions.
+
+Utility functions for building and inspecting SignatureDef protos.
+"""
+# pylint: disable=unused-import
+from tensorflow.python.saved_model.signature_def_utils_impl import build_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import classification_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import is_valid_signature
+from tensorflow.python.saved_model.signature_def_utils_impl import load_op_from_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import op_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import predict_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import regression_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import supervised_eval_signature_def
+from tensorflow.python.saved_model.signature_def_utils_impl import supervised_train_signature_def
+# pylint: enable=unused-import

tensorflow/python/saved_model/signature_def_utils_impl.py

@@ -0,0 +1,420 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""SignatureDef utility functions implementation."""
+
+
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.python.framework import errors
+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.saved_model import signature_constants
+from tensorflow.python.saved_model import utils_impl as utils
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(
+    v1=[
+        'saved_model.build_signature_def',
+        'saved_model.signature_def_utils.build_signature_def'
+    ])
+@deprecation.deprecated_endpoints(
+    'saved_model.signature_def_utils.build_signature_def')
+def build_signature_def(
+    inputs=None, outputs=None, method_name=None, defaults=None
+):
+  """Utility function to build a SignatureDef protocol buffer.
+
+  Args:
+    inputs: Inputs of the SignatureDef defined as a proto map of string to
+      tensor info.
+    outputs: Outputs of the SignatureDef defined as a proto map of string to
+      tensor info.
+    method_name: Method name of the SignatureDef as a string.
+    defaults: Defaults of the SignatureDef defined as a proto map of string to
+      TensorProto.
+
+  Returns:
+    A SignatureDef protocol buffer constructed based on the supplied arguments.
+  """
+  signature_def = meta_graph_pb2.SignatureDef()
+  if inputs is not None:
+    for item in inputs:
+      signature_def.inputs[item].CopyFrom(inputs[item])
+  if outputs is not None:
+    for item in outputs:
+      signature_def.outputs[item].CopyFrom(outputs[item])
+  if method_name is not None:
+    signature_def.method_name = method_name
+  if defaults is not None:
+    for arg_name, default in defaults.items():
+      if isinstance(default, ops.EagerTensor):
+        signature_def.defaults[arg_name].CopyFrom(
+            tensor_util.make_tensor_proto(default.numpy())
+        )
+      else:
+        if default.op.type == 'Const':
+          signature_def.defaults[arg_name].CopyFrom(
+              default.op.get_attr('value')
+          )
+        else:
+          raise ValueError(
+              f'Unable to convert object {str(default)} of type {type(default)}'
+              ' to TensorProto.'
+          )
+  return signature_def
+
+
+@tf_export(
+    v1=[
+        'saved_model.regression_signature_def',
+        'saved_model.signature_def_utils.regression_signature_def'
+    ])
+@deprecation.deprecated_endpoints(
+    'saved_model.signature_def_utils.regression_signature_def')
+def regression_signature_def(examples, predictions):
+  """Creates regression signature from given examples and predictions.
+
+  This function produces signatures intended for use with the TensorFlow Serving
+  Regress API (tensorflow_serving/apis/prediction_service.proto), and so
+  constrains the input and output types to those allowed by TensorFlow Serving.
+
+  Args:
+    examples: A string `Tensor`, expected to accept serialized tf.Examples.
+    predictions: A float `Tensor`.
+
+  Returns:
+    A regression-flavored signature_def.
+
+  Raises:
+    ValueError: If examples is `None`.
+  """
+  if examples is None:
+    raise ValueError('Regression `examples` cannot be None.')
+  if not isinstance(examples, tensor_lib.Tensor):
+    raise ValueError('Expected regression `examples` to be of type Tensor. '
+                     f'Found `examples` of type {type(examples)}.')
+  if predictions is None:
+    raise ValueError('Regression `predictions` cannot be None.')
+
+  input_tensor_info = utils.build_tensor_info(examples)
+  if input_tensor_info.dtype != types_pb2.DT_STRING:
+    raise ValueError('Regression input tensors must be of type string. '
+                     f'Found tensors with type {input_tensor_info.dtype}.')
+  signature_inputs = {signature_constants.REGRESS_INPUTS: input_tensor_info}
+
+  output_tensor_info = utils.build_tensor_info(predictions)
+  if output_tensor_info.dtype != types_pb2.DT_FLOAT:
+    raise ValueError('Regression output tensors must be of type float. '
+                     f'Found tensors with type {output_tensor_info.dtype}.')
+  signature_outputs = {signature_constants.REGRESS_OUTPUTS: output_tensor_info}
+
+  signature_def = build_signature_def(
+      signature_inputs, signature_outputs,
+      signature_constants.REGRESS_METHOD_NAME)
+
+  return signature_def
+
+
+@tf_export(
+    v1=[
+        'saved_model.classification_signature_def',
+        'saved_model.signature_def_utils.classification_signature_def'
+    ])
+@deprecation.deprecated_endpoints(
+    'saved_model.signature_def_utils.classification_signature_def')
+def classification_signature_def(examples, classes, scores):
+  """Creates classification signature from given examples and predictions.
+
+  This function produces signatures intended for use with the TensorFlow Serving
+  Classify API (tensorflow_serving/apis/prediction_service.proto), and so
+  constrains the input and output types to those allowed by TensorFlow Serving.
+
+  Args:
+    examples: A string `Tensor`, expected to accept serialized tf.Examples.
+    classes: A string `Tensor`.  Note that the ClassificationResponse message
+      requires that class labels are strings, not integers or anything else.
+    scores: a float `Tensor`.
+
+  Returns:
+    A classification-flavored signature_def.
+
+  Raises:
+    ValueError: If examples is `None`.
+  """
+  if examples is None:
+    raise ValueError('Classification `examples` cannot be None.')
+  if not isinstance(examples, tensor_lib.Tensor):
+    raise ValueError('Classification `examples` must be a string Tensor. '
+                     f'Found `examples` of type {type(examples)}.')
+  if classes is None and scores is None:
+    raise ValueError('Classification `classes` and `scores` cannot both be '
+                     'None.')
+
+  input_tensor_info = utils.build_tensor_info(examples)
+  if input_tensor_info.dtype != types_pb2.DT_STRING:
+    raise ValueError('Classification input tensors must be of type string. '
+                     f'Found tensors of type {input_tensor_info.dtype}')
+  signature_inputs = {signature_constants.CLASSIFY_INPUTS: input_tensor_info}
+
+  signature_outputs = {}
+  if classes is not None:
+    classes_tensor_info = utils.build_tensor_info(classes)
+    if classes_tensor_info.dtype != types_pb2.DT_STRING:
+      raise ValueError('Classification classes must be of type string Tensor. '
+                       f'Found tensors of type {classes_tensor_info.dtype}.`')
+    signature_outputs[signature_constants.CLASSIFY_OUTPUT_CLASSES] = (
+        classes_tensor_info)
+  if scores is not None:
+    scores_tensor_info = utils.build_tensor_info(scores)
+    if scores_tensor_info.dtype != types_pb2.DT_FLOAT:
+      raise ValueError('Classification scores must be a float Tensor.')
+    signature_outputs[signature_constants.CLASSIFY_OUTPUT_SCORES] = (
+        scores_tensor_info)
+
+  signature_def = build_signature_def(
+      signature_inputs, signature_outputs,
+      signature_constants.CLASSIFY_METHOD_NAME)
+
+  return signature_def
+
+
+@tf_export(
+    v1=[
+        'saved_model.predict_signature_def',
+        'saved_model.signature_def_utils.predict_signature_def'
+    ])
+@deprecation.deprecated_endpoints(
+    'saved_model.signature_def_utils.predict_signature_def')
+def predict_signature_def(inputs, outputs):
+  """Creates prediction signature from given inputs and outputs.
+
+  This function produces signatures intended for use with the TensorFlow Serving
+  Predict API (tensorflow_serving/apis/prediction_service.proto). This API
+  imposes no constraints on the input and output types.
+
+  Args:
+    inputs: dict of string to `Tensor`.
+    outputs: dict of string to `Tensor`.
+
+  Returns:
+    A prediction-flavored signature_def.
+
+  Raises:
+    ValueError: If inputs or outputs is `None`.
+  """
+  if inputs is None or not inputs:
+    raise ValueError('Prediction `inputs` cannot be None or empty.')
+  if outputs is None or not outputs:
+    raise ValueError('Prediction `outputs` cannot be None or empty.')
+
+  signature_inputs = {key: utils.build_tensor_info(tensor)
+                      for key, tensor in inputs.items()}
+  signature_outputs = {key: utils.build_tensor_info(tensor)
+                       for key, tensor in outputs.items()}
+
+  signature_def = build_signature_def(
+      signature_inputs, signature_outputs,
+      signature_constants.PREDICT_METHOD_NAME)
+
+  return signature_def
+
+
+def supervised_train_signature_def(
+    inputs, loss, predictions=None, metrics=None):
+  return _supervised_signature_def(
+      signature_constants.SUPERVISED_TRAIN_METHOD_NAME, inputs, loss=loss,
+      predictions=predictions, metrics=metrics)
+
+
+def supervised_eval_signature_def(
+    inputs, loss, predictions=None, metrics=None):
+  return _supervised_signature_def(
+      signature_constants.SUPERVISED_EVAL_METHOD_NAME, inputs, loss=loss,
+      predictions=predictions, metrics=metrics)
+
+
+def _supervised_signature_def(
+    method_name, inputs, loss=None, predictions=None,
+    metrics=None):
+  """Creates a signature for training and eval data.
+
+  This function produces signatures that describe the inputs and outputs
+  of a supervised process, such as training or evaluation, that
+  results in loss, metrics, and the like. Note that this function only requires
+  inputs to be not None.
+
+  Args:
+    method_name: Method name of the SignatureDef as a string.
+    inputs: dict of string to `Tensor`.
+    loss: dict of string to `Tensor` representing computed loss.
+    predictions: dict of string to `Tensor` representing the output predictions.
+    metrics: dict of string to `Tensor` representing metric ops.
+
+  Returns:
+    A train- or eval-flavored signature_def.
+
+  Raises:
+    ValueError: If inputs or outputs is `None`.
+  """
+  if inputs is None or not inputs:
+    raise ValueError(f'{method_name} `inputs` cannot be None or empty.')
+
+  signature_inputs = {key: utils.build_tensor_info(tensor)
+                      for key, tensor in inputs.items()}
+
+  signature_outputs = {}
+  for output_set in (loss, predictions, metrics):
+    if output_set is not None:
+      sig_out = {key: utils.build_tensor_info(tensor)
+                 for key, tensor in output_set.items()}
+      signature_outputs.update(sig_out)
+
+  signature_def = build_signature_def(
+      signature_inputs, signature_outputs, method_name)
+
+  return signature_def
+
+
+@tf_export(
+    v1=[
+        'saved_model.is_valid_signature',
+        'saved_model.signature_def_utils.is_valid_signature'
+    ])
+@deprecation.deprecated_endpoints(
+    'saved_model.signature_def_utils.is_valid_signature')
+def is_valid_signature(signature_def):
+  """Determine whether a SignatureDef can be served by TensorFlow Serving."""
+  if signature_def is None:
+    return False
+  return (_is_valid_classification_signature(signature_def) or
+          _is_valid_regression_signature(signature_def) or
+          _is_valid_predict_signature(signature_def))
+
+
+def _is_valid_predict_signature(signature_def):
+  """Determine whether the argument is a servable 'predict' SignatureDef."""
+  if signature_def.method_name != signature_constants.PREDICT_METHOD_NAME:
+    return False
+  if not signature_def.inputs.keys():
+    return False
+  if not signature_def.outputs.keys():
+    return False
+  return True
+
+
+def _is_valid_regression_signature(signature_def):
+  """Determine whether the argument is a servable 'regress' SignatureDef."""
+  if signature_def.method_name != signature_constants.REGRESS_METHOD_NAME:
+    return False
+
+  if (set(signature_def.inputs.keys())
+      != set([signature_constants.REGRESS_INPUTS])):
+    return False
+  if (signature_def.inputs[signature_constants.REGRESS_INPUTS].dtype !=
+      types_pb2.DT_STRING):
+    return False
+
+  if (set(signature_def.outputs.keys())
+      != set([signature_constants.REGRESS_OUTPUTS])):
+    return False
+  if (signature_def.outputs[signature_constants.REGRESS_OUTPUTS].dtype !=
+      types_pb2.DT_FLOAT):
+    return False
+
+  return True
+
+
+def _is_valid_classification_signature(signature_def):
+  """Determine whether the argument is a servable 'classify' SignatureDef."""
+  if signature_def.method_name != signature_constants.CLASSIFY_METHOD_NAME:
+    return False
+
+  if (set(signature_def.inputs.keys())
+      != set([signature_constants.CLASSIFY_INPUTS])):
+    return False
+  if (signature_def.inputs[signature_constants.CLASSIFY_INPUTS].dtype !=
+      types_pb2.DT_STRING):
+    return False
+
+  allowed_outputs = set([signature_constants.CLASSIFY_OUTPUT_CLASSES,
+                         signature_constants.CLASSIFY_OUTPUT_SCORES])
+
+  if not signature_def.outputs.keys():
+    return False
+  if set(signature_def.outputs.keys()) - allowed_outputs:
+    return False
+  if (signature_constants.CLASSIFY_OUTPUT_CLASSES in signature_def.outputs
+      and
+      signature_def.outputs[signature_constants.CLASSIFY_OUTPUT_CLASSES].dtype
+      != types_pb2.DT_STRING):
+    return False
+  if (signature_constants.CLASSIFY_OUTPUT_SCORES in signature_def.outputs
+      and
+      signature_def.outputs[signature_constants.CLASSIFY_OUTPUT_SCORES].dtype !=
+      types_pb2.DT_FLOAT):
+    return False
+
+  return True
+
+
+def op_signature_def(op, key):
+  """Creates a signature def with the output pointing to an op.
+
+  Note that op isn't strictly enforced to be an Op object, and may be a Tensor.
+  It is recommended to use the build_signature_def() function for Tensors.
+
+  Args:
+    op: An Op (or possibly Tensor).
+    key: Key to graph element in the SignatureDef outputs.
+
+  Returns:
+    A SignatureDef with a single output pointing to the op.
+  """
+  # Use build_tensor_info_from_op, which creates a TensorInfo from the element's
+  # name.
+  return build_signature_def(outputs={key: utils.build_tensor_info_from_op(op)})
+
+
+def load_op_from_signature_def(signature_def, key, import_scope=None):
+  """Load an Op from a SignatureDef created by op_signature_def().
+
+  Args:
+    signature_def: a SignatureDef proto
+    key: string key to op in the SignatureDef outputs.
+    import_scope: Scope used to import the op
+
+  Returns:
+    Op (or possibly Tensor) in the graph with the same name as saved in the
+      SignatureDef.
+
+  Raises:
+    NotFoundError: If the op could not be found in the graph.
+  """
+  tensor_info = signature_def.outputs[key]
+  try:
+    # The init and train ops are not strictly enforced to be operations, so
+    # retrieve any graph element (can be either op or tensor).
+    return utils.get_element_from_tensor_info(
+        tensor_info, import_scope=import_scope)
+  except KeyError:
+    raise errors.NotFoundError(
+        None, None,
+        f'The key "{key}" could not be found in the graph. Please make sure the'
+        ' SavedModel was created by the internal _SavedModelBuilder. If you '
+        'are using the public API, please make sure the SignatureDef in the '
+        f'SavedModel does not contain the key "{key}".')

tensorflow/python/saved_model/signature_def_utils_test.py

@@ -0,0 +1,483 @@
+# 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.
+# ==============================================================================
+"""Tests for SignatureDef utils."""
+
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import signature_def_utils_impl
+from tensorflow.python.saved_model import utils
+
+
+# We'll reuse the same tensor_infos in multiple contexts just for the tests.
+# The validator doesn't check shapes so we just omit them.
+_STRING = meta_graph_pb2.TensorInfo(
+    name="foobar",
+    dtype=dtypes.string.as_datatype_enum
+)
+
+
+_FLOAT = meta_graph_pb2.TensorInfo(
+    name="foobar",
+    dtype=dtypes.float32.as_datatype_enum
+)
+
+
+def _make_signature(inputs, outputs, name=None):
+  input_info = {
+      input_name: utils.build_tensor_info(tensor)
+      for input_name, tensor in inputs.items()
+  }
+  output_info = {
+      output_name: utils.build_tensor_info(tensor)
+      for output_name, tensor in outputs.items()
+  }
+  return signature_def_utils_impl.build_signature_def(input_info, output_info,
+                                                      name)
+
+
+class SignatureDefUtilsTest(test.TestCase):
+
+  def testBuildSignatureDef(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that uses functionality that does not work
+    # with eager tensors (namely build_tensor_info).
+    with ops.Graph().as_default():
+      x = array_ops.placeholder(dtypes.float32, 1, name="x")
+      x_tensor_info = utils.build_tensor_info(x)
+      inputs = {}
+      inputs["foo-input"] = x_tensor_info
+
+      y = array_ops.placeholder(dtypes.float32, name="y")
+      y_tensor_info = utils.build_tensor_info(y)
+      outputs = {}
+      outputs["foo-output"] = y_tensor_info
+
+      default_tensor = constant_op.constant(1.0, name="w")
+      defaults = {}
+      defaults["w"] = default_tensor
+
+    signature_def = signature_def_utils_impl.build_signature_def(
+        inputs, outputs, "foo-method-name", defaults
+    )
+    self.assertEqual("foo-method-name", signature_def.method_name)
+
+    # Check inputs in signature def.
+    self.assertEqual(1, len(signature_def.inputs))
+    x_tensor_info_actual = signature_def.inputs["foo-input"]
+    self.assertEqual("x:0", x_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_FLOAT, x_tensor_info_actual.dtype)
+    self.assertEqual(1, len(x_tensor_info_actual.tensor_shape.dim))
+    self.assertEqual(1, x_tensor_info_actual.tensor_shape.dim[0].size)
+
+    # Check outputs in signature def.
+    self.assertEqual(1, len(signature_def.outputs))
+    y_tensor_info_actual = signature_def.outputs["foo-output"]
+    self.assertEqual("y:0", y_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_FLOAT, y_tensor_info_actual.dtype)
+    self.assertEqual(0, len(y_tensor_info_actual.tensor_shape.dim))
+
+    self.assertEqual(1, len(signature_def.defaults))
+    self.assertEqual(types_pb2.DT_FLOAT, signature_def.defaults["w"].dtype)
+    self.assertEqual(1.0, signature_def.defaults["w"].float_val[0])
+
+  def testRegressionSignatureDef(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that uses functionality that does not work
+    # with eager tensors (namely build_tensor_info).
+    with ops.Graph().as_default():
+      input1 = constant_op.constant("a", name="input-1")
+      output1 = constant_op.constant(2.2, name="output-1")
+      signature_def = signature_def_utils_impl.regression_signature_def(
+          input1, output1)
+
+    self.assertEqual(signature_constants.REGRESS_METHOD_NAME,
+                     signature_def.method_name)
+
+    # Check inputs in signature def.
+    self.assertEqual(1, len(signature_def.inputs))
+    x_tensor_info_actual = (
+        signature_def.inputs[signature_constants.REGRESS_INPUTS])
+    self.assertEqual("input-1:0", x_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, x_tensor_info_actual.dtype)
+    self.assertEqual(0, len(x_tensor_info_actual.tensor_shape.dim))
+
+    # Check outputs in signature def.
+    self.assertEqual(1, len(signature_def.outputs))
+    y_tensor_info_actual = (
+        signature_def.outputs[signature_constants.REGRESS_OUTPUTS])
+    self.assertEqual("output-1:0", y_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_FLOAT, y_tensor_info_actual.dtype)
+    self.assertEqual(0, len(y_tensor_info_actual.tensor_shape.dim))
+
+  def testClassificationSignatureDef(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that uses functionality that does not work
+    # with eager tensors (namely build_tensor_info).
+    with ops.Graph().as_default():
+      input1 = constant_op.constant("a", name="input-1")
+      output1 = constant_op.constant("b", name="output-1")
+      output2 = constant_op.constant(3.3, name="output-2")
+      signature_def = signature_def_utils_impl.classification_signature_def(
+          input1, output1, output2)
+
+    self.assertEqual(signature_constants.CLASSIFY_METHOD_NAME,
+                     signature_def.method_name)
+
+    # Check inputs in signature def.
+    self.assertEqual(1, len(signature_def.inputs))
+    x_tensor_info_actual = (
+        signature_def.inputs[signature_constants.CLASSIFY_INPUTS])
+    self.assertEqual("input-1:0", x_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, x_tensor_info_actual.dtype)
+    self.assertEqual(0, len(x_tensor_info_actual.tensor_shape.dim))
+
+    # Check outputs in signature def.
+    self.assertEqual(2, len(signature_def.outputs))
+    classes_tensor_info_actual = (
+        signature_def.outputs[signature_constants.CLASSIFY_OUTPUT_CLASSES])
+    self.assertEqual("output-1:0", classes_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, classes_tensor_info_actual.dtype)
+    self.assertEqual(0, len(classes_tensor_info_actual.tensor_shape.dim))
+    scores_tensor_info_actual = (
+        signature_def.outputs[signature_constants.CLASSIFY_OUTPUT_SCORES])
+    self.assertEqual("output-2:0", scores_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_FLOAT, scores_tensor_info_actual.dtype)
+    self.assertEqual(0, len(scores_tensor_info_actual.tensor_shape.dim))
+
+  def testPredictionSignatureDef(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that uses functionality that does not work
+    # with eager tensors (namely build_tensor_info).
+    with ops.Graph().as_default():
+      input1 = constant_op.constant("a", name="input-1")
+      input2 = constant_op.constant("b", name="input-2")
+      output1 = constant_op.constant("c", name="output-1")
+      output2 = constant_op.constant("d", name="output-2")
+      signature_def = signature_def_utils_impl.predict_signature_def(
+          {
+              "input-1": input1,
+              "input-2": input2
+          }, {
+              "output-1": output1,
+              "output-2": output2
+          })
+
+    self.assertEqual(signature_constants.PREDICT_METHOD_NAME,
+                     signature_def.method_name)
+
+    # Check inputs in signature def.
+    self.assertEqual(2, len(signature_def.inputs))
+    input1_tensor_info_actual = (signature_def.inputs["input-1"])
+    self.assertEqual("input-1:0", input1_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, input1_tensor_info_actual.dtype)
+    self.assertEqual(0, len(input1_tensor_info_actual.tensor_shape.dim))
+    input2_tensor_info_actual = (signature_def.inputs["input-2"])
+    self.assertEqual("input-2:0", input2_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, input2_tensor_info_actual.dtype)
+    self.assertEqual(0, len(input2_tensor_info_actual.tensor_shape.dim))
+
+    # Check outputs in signature def.
+    self.assertEqual(2, len(signature_def.outputs))
+    output1_tensor_info_actual = (signature_def.outputs["output-1"])
+    self.assertEqual("output-1:0", output1_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, output1_tensor_info_actual.dtype)
+    self.assertEqual(0, len(output1_tensor_info_actual.tensor_shape.dim))
+    output2_tensor_info_actual = (signature_def.outputs["output-2"])
+    self.assertEqual("output-2:0", output2_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, output2_tensor_info_actual.dtype)
+    self.assertEqual(0, len(output2_tensor_info_actual.tensor_shape.dim))
+
+  def testTrainSignatureDef(self):
+    self._testSupervisedSignatureDef(
+        signature_def_utils_impl.supervised_train_signature_def,
+        signature_constants.SUPERVISED_TRAIN_METHOD_NAME)
+
+  def testEvalSignatureDef(self):
+    self._testSupervisedSignatureDef(
+        signature_def_utils_impl.supervised_eval_signature_def,
+        signature_constants.SUPERVISED_EVAL_METHOD_NAME)
+
+  def _testSupervisedSignatureDef(self, fn_to_test, method_name):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that uses functionality that does not work
+    # with eager tensors (namely build_tensor_info).
+    with ops.Graph().as_default():
+      inputs = {
+          "input-1": constant_op.constant("a", name="input-1"),
+          "input-2": constant_op.constant("b", name="input-2"),
+      }
+      loss = {"loss-1": constant_op.constant(0.45, name="loss-1")}
+      predictions = {
+          "classes": constant_op.constant([100], name="classes"),
+      }
+      metrics_val = constant_op.constant(100.0, name="metrics_val")
+      metrics = {
+          "metrics/value":
+              metrics_val,
+          "metrics/update_op":
+              array_ops.identity(metrics_val, name="metrics_op"),
+      }
+
+      signature_def = fn_to_test(inputs, loss, predictions, metrics)
+
+    self.assertEqual(method_name, signature_def.method_name)
+
+    # Check inputs in signature def.
+    self.assertEqual(2, len(signature_def.inputs))
+    input1_tensor_info_actual = (signature_def.inputs["input-1"])
+    self.assertEqual("input-1:0", input1_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, input1_tensor_info_actual.dtype)
+    self.assertEqual(0, len(input1_tensor_info_actual.tensor_shape.dim))
+    input2_tensor_info_actual = (signature_def.inputs["input-2"])
+    self.assertEqual("input-2:0", input2_tensor_info_actual.name)
+    self.assertEqual(types_pb2.DT_STRING, input2_tensor_info_actual.dtype)
+    self.assertEqual(0, len(input2_tensor_info_actual.tensor_shape.dim))
+
+    # Check outputs in signature def.
+    self.assertEqual(4, len(signature_def.outputs))
+    self.assertEqual("loss-1:0", signature_def.outputs["loss-1"].name)
+    self.assertEqual(types_pb2.DT_FLOAT, signature_def.outputs["loss-1"].dtype)
+
+    self.assertEqual("classes:0", signature_def.outputs["classes"].name)
+    self.assertEqual(1, len(signature_def.outputs["classes"].tensor_shape.dim))
+
+    self.assertEqual(
+        "metrics_val:0", signature_def.outputs["metrics/value"].name)
+    self.assertEqual(
+        types_pb2.DT_FLOAT, signature_def.outputs["metrics/value"].dtype)
+
+    self.assertEqual(
+        "metrics_op:0", signature_def.outputs["metrics/update_op"].name)
+    self.assertEqual(
+        types_pb2.DT_FLOAT, signature_def.outputs["metrics/value"].dtype)
+
+  def testTrainSignatureDefMissingInputs(self):
+    self._testSupervisedSignatureDefMissingInputs(
+        signature_def_utils_impl.supervised_train_signature_def,
+        signature_constants.SUPERVISED_TRAIN_METHOD_NAME)
+
+  def testEvalSignatureDefMissingInputs(self):
+    self._testSupervisedSignatureDefMissingInputs(
+        signature_def_utils_impl.supervised_eval_signature_def,
+        signature_constants.SUPERVISED_EVAL_METHOD_NAME)
+
+  def _testSupervisedSignatureDefMissingInputs(self, fn_to_test, method_name):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that uses functionality that does not work
+    # with eager tensors (namely build_tensor_info).
+    with ops.Graph().as_default():
+      inputs = {
+          "input-1": constant_op.constant("a", name="input-1"),
+          "input-2": constant_op.constant("b", name="input-2"),
+      }
+      loss = {"loss-1": constant_op.constant(0.45, name="loss-1")}
+      predictions = {
+          "classes": constant_op.constant([100], name="classes"),
+      }
+      metrics_val = constant_op.constant(100, name="metrics_val")
+      metrics = {
+          "metrics/value":
+              metrics_val,
+          "metrics/update_op":
+              array_ops.identity(metrics_val, name="metrics_op"),
+      }
+
+      with self.assertRaises(ValueError):
+        signature_def = fn_to_test({},
+                                   loss=loss,
+                                   predictions=predictions,
+                                   metrics=metrics)
+
+      signature_def = fn_to_test(inputs, loss=loss)
+      self.assertEqual(method_name, signature_def.method_name)
+      self.assertEqual(1, len(signature_def.outputs))
+
+      signature_def = fn_to_test(inputs, metrics=metrics, loss=loss)
+      self.assertEqual(method_name, signature_def.method_name)
+      self.assertEqual(3, len(signature_def.outputs))
+
+  def _assertValidSignature(self, inputs, outputs, method_name):
+    signature_def = signature_def_utils_impl.build_signature_def(
+        inputs, outputs, method_name)
+    self.assertTrue(
+        signature_def_utils_impl.is_valid_signature(signature_def))
+
+  def _assertInvalidSignature(self, inputs, outputs, method_name):
+    signature_def = signature_def_utils_impl.build_signature_def(
+        inputs, outputs, method_name)
+    self.assertFalse(
+        signature_def_utils_impl.is_valid_signature(signature_def))
+
+  def testValidSignaturesAreAccepted(self):
+    self._assertValidSignature(
+        {"inputs": _STRING},
+        {"classes": _STRING, "scores": _FLOAT},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertValidSignature(
+        {"inputs": _STRING},
+        {"classes": _STRING},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertValidSignature(
+        {"inputs": _STRING},
+        {"scores": _FLOAT},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertValidSignature(
+        {"inputs": _STRING},
+        {"outputs": _FLOAT},
+        signature_constants.REGRESS_METHOD_NAME)
+
+    self._assertValidSignature(
+        {"foo": _STRING, "bar": _FLOAT},
+        {"baz": _STRING, "qux": _FLOAT},
+        signature_constants.PREDICT_METHOD_NAME)
+
+  def testInvalidMethodNameSignatureIsRejected(self):
+    # WRONG METHOD
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {"classes": _STRING, "scores": _FLOAT},
+        "WRONG method name")
+
+  def testInvalidClassificationSignaturesAreRejected(self):
+    # CLASSIFY: wrong types
+    self._assertInvalidSignature(
+        {"inputs": _FLOAT},
+        {"classes": _STRING, "scores": _FLOAT},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {"classes": _FLOAT, "scores": _FLOAT},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {"classes": _STRING, "scores": _STRING},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    # CLASSIFY: wrong keys
+    self._assertInvalidSignature(
+        {},
+        {"classes": _STRING, "scores": _FLOAT},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs_WRONG": _STRING},
+        {"classes": _STRING, "scores": _FLOAT},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {"classes_WRONG": _STRING, "scores": _FLOAT},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {"classes": _STRING, "scores": _FLOAT, "extra_WRONG": _STRING},
+        signature_constants.CLASSIFY_METHOD_NAME)
+
+  def testInvalidRegressionSignaturesAreRejected(self):
+    # REGRESS: wrong types
+    self._assertInvalidSignature(
+        {"inputs": _FLOAT},
+        {"outputs": _FLOAT},
+        signature_constants.REGRESS_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {"outputs": _STRING},
+        signature_constants.REGRESS_METHOD_NAME)
+
+    # REGRESS: wrong keys
+    self._assertInvalidSignature(
+        {},
+        {"outputs": _FLOAT},
+        signature_constants.REGRESS_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs_WRONG": _STRING},
+        {"outputs": _FLOAT},
+        signature_constants.REGRESS_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {"outputs_WRONG": _FLOAT},
+        signature_constants.REGRESS_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {},
+        signature_constants.REGRESS_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"inputs": _STRING},
+        {"outputs": _FLOAT, "extra_WRONG": _STRING},
+        signature_constants.REGRESS_METHOD_NAME)
+
+  def testInvalidPredictSignaturesAreRejected(self):
+    # PREDICT: wrong keys
+    self._assertInvalidSignature(
+        {},
+        {"baz": _STRING, "qux": _FLOAT},
+        signature_constants.PREDICT_METHOD_NAME)
+
+    self._assertInvalidSignature(
+        {"foo": _STRING, "bar": _FLOAT},
+        {},
+        signature_constants.PREDICT_METHOD_NAME)
+
+  def testOpSignatureDef(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that uses functionality that does not work
+    # with eager tensors (namely build_tensor_info_from_op).
+    with ops.Graph().as_default():
+      key = "adding_1_and_2_key"
+      add_op = math_ops.add(1, 2, name="adding_1_and_2")
+      signature_def = signature_def_utils_impl.op_signature_def(add_op, key)
+
+    self.assertIn(key, signature_def.outputs)
+    self.assertEqual(add_op.name, signature_def.outputs[key].name)
+
+  def testLoadOpFromSignatureDef(self):
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that uses functionality that does not work
+    # with eager tensors (namely build_tensor_info_from_op).
+    with ops.Graph().as_default():
+      key = "adding_1_and_2_key"
+      add_op = math_ops.add(1, 2, name="adding_1_and_2")
+      signature_def = signature_def_utils_impl.op_signature_def(add_op, key)
+      self.assertEqual(
+          add_op,
+          signature_def_utils_impl.load_op_from_signature_def(
+              signature_def, key))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/signature_serialization.py

@@ -0,0 +1,377 @@
+# 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.
+# ==============================================================================
+"""Helpers for working with signatures in tf.saved_model.save."""
+
+from absl import logging
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as defun
+from tensorflow.python.eager.polymorphic_function import attributes
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.saved_model import function_serialization
+from tensorflow.python.saved_model import revived_types
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.trackable import base
+from tensorflow.python.types import core
+from tensorflow.python.util import compat
+from tensorflow.python.util import nest
+from tensorflow.python.util.compat import collections_abc
+
+DEFAULT_SIGNATURE_ATTR = "_default_save_signature"
+SIGNATURE_ATTRIBUTE_NAME = "signatures"
+# Max number of warnings to show if signature contains normalized input names.
+_NUM_DISPLAY_NORMALIZED_SIGNATURES = 5
+
+
+def _get_signature(function):
+  if (
+      isinstance(function, def_function.Function)
+      and function.input_signature is not None
+  ):
+    function = function._get_concrete_function_garbage_collected()  # pylint: disable=protected-access
+  if not isinstance(function, defun.ConcreteFunction):
+    return None
+  return function
+
+
+def _valid_signature(concrete_function):
+  """Returns whether concrete function can be converted to a signature."""
+  if not concrete_function.outputs:
+    # Functions without outputs don't make sense as signatures. We just don't
+    # have any way to run an Operation with no outputs as a SignatureDef in the
+    # 1.x style.
+    return False
+  try:
+    _validate_inputs(concrete_function)
+    _normalize_outputs(concrete_function.structured_outputs, "unused", "unused")
+  except ValueError:
+    return False
+  return True
+
+
+def _validate_inputs(concrete_function):
+  """Raises error if input type is tf.Variable."""
+  if any(
+      isinstance(inp, resource_variable_ops.VariableSpec)
+      for inp in nest.flatten(concrete_function.structured_input_signature)
+  ):
+    raise ValueError(
+        f"Unable to serialize concrete_function '{concrete_function.name}'"
+        "with tf.Variable input. Functions that expect tf.Variable "
+        "inputs cannot be exported as signatures."
+    )
+
+
+def _get_signature_name_changes(concrete_function):
+  """Checks for user-specified signature input names that are normalized."""
+  # Map of {user-given name: normalized name} if the names are un-identical.
+  name_changes = {}
+  for signature_input_name, graph_input in zip(
+      concrete_function.function_def.signature.input_arg,
+      concrete_function.graph.inputs,
+  ):
+    try:
+      user_specified_name = compat.as_str(
+          graph_input.op.get_attr("_user_specified_name")
+      )
+      if signature_input_name.name != user_specified_name:
+        name_changes[user_specified_name] = signature_input_name.name
+    except ValueError:
+      # Signature input does not have a user-specified name.
+      pass
+  return name_changes
+
+
+def find_function_to_export(saveable_view):
+  """Function to export, None if no suitable function was found."""
+  # If the user did not specify signatures, check the root object for a function
+  # that can be made into a signature.
+  children = saveable_view.list_children(saveable_view.root)
+
+  # TODO(b/205014194): Discuss removing this behaviour. It can lead to WTFs when
+  # a user decides to annotate more functions with tf.function and suddenly
+  # serving that model way later in the process stops working.
+  possible_signatures = []
+  for name, child in children:
+    if not isinstance(child, (def_function.Function, defun.ConcreteFunction)):
+      continue
+    if name == DEFAULT_SIGNATURE_ATTR:
+      return child
+    concrete = _get_signature(child)
+    if concrete is not None and _valid_signature(concrete):
+      possible_signatures.append(concrete)
+
+  if len(possible_signatures) == 1:
+    single_function = possible_signatures[0]
+    signature = _get_signature(single_function)
+    if signature and _valid_signature(signature):
+      return signature
+  return None
+
+
+def canonicalize_signatures(signatures):
+  """Converts `signatures` into a dictionary of concrete functions."""
+  if signatures is None:
+    return {}, {}, {}
+  if not isinstance(signatures, collections_abc.Mapping):
+    signatures = {
+        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signatures
+    }
+  num_normalized_signatures_counter = 0
+  concrete_signatures = {}
+  wrapped_functions = {}
+  defaults = {}
+  for signature_key, function in signatures.items():
+    original_function = signature_function = _get_signature(function)
+    if signature_function is None:
+      raise ValueError(
+          "Expected a TensorFlow function for which to generate a signature, "
+          f"but got {function}. Only `tf.functions` with an input signature or "
+          "concrete functions can be used as a signature."
+      )
+
+    wrapped_functions[original_function] = signature_function = (
+        wrapped_functions.get(original_function)
+        or function_serialization.wrap_cached_variables(original_function)
+    )
+    _validate_inputs(signature_function)
+    if num_normalized_signatures_counter < _NUM_DISPLAY_NORMALIZED_SIGNATURES:
+      signature_name_changes = _get_signature_name_changes(signature_function)
+      if signature_name_changes:
+        num_normalized_signatures_counter += 1
+        logging.info(
+            "Function `%s` contains input name(s) %s with unsupported "
+            "characters which will be renamed to %s in the SavedModel.",
+            compat.as_str(signature_function.graph.name),
+            ", ".join(signature_name_changes.keys()),
+            ", ".join(signature_name_changes.values()),
+        )
+
+    # Re-wrap the function so that it returns a dictionary of Tensors. This
+    # matches the format of 1.x-style signatures.
+    # pylint: disable=cell-var-from-loop
+    def signature_wrapper(**kwargs):
+      structured_outputs = signature_function(**kwargs)
+      return _normalize_outputs(
+          structured_outputs, signature_function.name, signature_key
+      )
+
+    if hasattr(function, "__name__"):
+      signature_wrapper.__name__ = "signature_wrapper_" + function.__name__
+
+    # Extract experimental attributes and propagate it to the wrapped_function.
+    # At the moment only the `disable_summaries_at_runtime` attr needs to be
+    # propagated.
+    experimental_attributes = {}
+    for attr in attributes.POLYMORPHIC_FUNCTION_ALLOWLIST:
+      attr_value = signature_function.function_def.attr.get(attr, None)
+      if attr != attributes.NO_INLINE and attr_value is not None:
+        experimental_attributes[attr] = attr_value
+    if not experimental_attributes:
+      experimental_attributes = None
+
+    wrapped_function = def_function.function(
+        signature_wrapper, experimental_attributes=experimental_attributes
+    )
+    tensor_spec_signature = {}
+    if signature_function.structured_input_signature is not None:
+      # The structured input signature may contain other non-tensor arguments.
+      inputs = filter(
+          lambda x: isinstance(x, tensor.TensorSpec),
+          nest.flatten(
+              signature_function.structured_input_signature,
+              expand_composites=True,
+          ),
+      )
+    else:
+      # Structured input signature isn't always defined for some functions.
+      inputs = signature_function.inputs
+
+    for keyword, inp in zip(
+        signature_function._arg_keywords,  # pylint: disable=protected-access
+        inputs,
+    ):
+      keyword = compat.as_str(keyword)
+      if isinstance(inp, tensor.TensorSpec):
+        spec = tensor.TensorSpec(inp.shape, inp.dtype, name=keyword)
+      else:
+        spec = tensor.TensorSpec.from_tensor(inp, name=keyword)
+      tensor_spec_signature[keyword] = spec
+    final_concrete = wrapped_function._get_concrete_function_garbage_collected(  # pylint: disable=protected-access
+        **tensor_spec_signature
+    )
+    # pylint: disable=protected-access
+    if len(final_concrete._arg_keywords) == 1:
+      # If there is only one input to the signature, a very common case, then
+      # ordering is unambiguous and we can let people pass a positional
+      # argument. Since SignatureDefs are unordered (protobuf "map") multiple
+      # arguments means we need to be keyword-only.
+      final_concrete._num_positional_args = 1
+    else:
+      final_concrete._num_positional_args = 0
+    # pylint: enable=protected-access
+    concrete_signatures[signature_key] = final_concrete
+    # pylint: enable=cell-var-from-loop
+    if isinstance(function, core.PolymorphicFunction):
+      flattened_defaults = nest.flatten(
+          function.function_spec.fullargspec.defaults  # pylint: disable=protected-access
+      )
+      len_default = len(flattened_defaults or [])
+      arg_names = list(tensor_spec_signature.keys())
+      if len_default > 0:
+        # tensor_spec_signature uses the same nest.flatten() as
+        # flattened_defaults.
+        for arg, default in zip(
+            arg_names[-len_default:],  # pylint: disable=protected-access
+            flattened_defaults or [],
+        ):
+          if not isinstance(default, tensor.Tensor):
+            continue
+          defaults.setdefault(signature_key, {})[arg] = default
+  return concrete_signatures, wrapped_functions, defaults
+
+
+def _normalize_outputs(outputs, function_name, signature_key):
+  """Normalize outputs if necessary and check that they are tensors."""
+  # Convert `outputs` to a dictionary (if it's not one already).
+  if not isinstance(outputs, collections_abc.Mapping):
+    # Check if `outputs` is a namedtuple.
+    if hasattr(outputs, "_asdict"):
+      outputs = outputs._asdict()
+    else:
+      if not isinstance(outputs, collections_abc.Sequence):
+        outputs = [outputs]
+      outputs = {
+          "output_{}".format(output_index): output
+          for output_index, output in enumerate(outputs)
+      }
+
+  # Check that the keys of `outputs` are strings and the values are Tensors.
+  for key, value in outputs.items():
+    if not isinstance(key, compat.bytes_or_text_types):
+      raise ValueError(
+          f"Got a dictionary with a non-string key {key!r} in the output of "
+          f"the function {compat.as_str_any(function_name)} used to generate "
+          f"the SavedModel signature {signature_key!r}."
+      )
+    if not isinstance(value, (tensor.Tensor, composite_tensor.CompositeTensor)):
+      raise ValueError(
+          f"Got a non-Tensor value {value!r} for key {key!r} in the output of "
+          f"the function {compat.as_str_any(function_name)} used to generate "
+          f"the SavedModel signature {signature_key!r}. "
+          "Outputs for functions used as signatures must be a single Tensor, "
+          "a sequence of Tensors, or a dictionary from string to Tensor."
+      )
+  return outputs
+
+
+# _SignatureMap is immutable to ensure that users do not expect changes to be
+# reflected in the SavedModel. Using public APIs, tf.saved_model.load() is the
+# only way to create a _SignatureMap and there is no way to modify it. So we can
+# safely ignore/overwrite ".signatures" attributes attached to objects being
+# saved if they contain a _SignatureMap. A ".signatures" attribute containing
+# any other type (e.g. a regular dict) will raise an exception asking the user
+# to first "del obj.signatures" if they want it overwritten.
+class _SignatureMap(collections_abc.Mapping, base.Trackable):
+  """A collection of SavedModel signatures."""
+
+  def __init__(self):
+    self._signatures = {}
+
+  def _add_signature(self, name, concrete_function):
+    """Adds a signature to the _SignatureMap."""
+    # Ideally this object would be immutable, but restore is streaming so we do
+    # need a private API for adding new signatures to an existing object.
+    self._signatures[name] = concrete_function
+
+  def __getitem__(self, key):
+    return self._signatures[key]
+
+  def __iter__(self):
+    return iter(self._signatures)
+
+  def __len__(self):
+    return len(self._signatures)
+
+  def __repr__(self):
+    return "_SignatureMap({})".format(self._signatures)
+
+  def _trackable_children(self, save_type=base.SaveType.CHECKPOINT, **kwargs):
+    if save_type != base.SaveType.SAVEDMODEL:
+      return {}
+
+    return {
+        key: value
+        for key, value in self.items()
+        if isinstance(value, (def_function.Function, defun.ConcreteFunction))
+    }
+
+
+revived_types.register_revived_type(
+    "signature_map",
+    lambda obj: isinstance(obj, _SignatureMap),
+    versions=[
+        revived_types.VersionedTypeRegistration(
+            # Standard dependencies are enough to reconstruct the trackable
+            # items in dictionaries, so we don't need to save any extra
+            # information.
+            object_factory=lambda proto: _SignatureMap(),
+            version=1,
+            min_producer_version=1,
+            min_consumer_version=1,
+            setter=_SignatureMap._add_signature,  # pylint: disable=protected-access
+        )
+    ],
+)
+
+
+def create_signature_map(signatures):
+  """Creates an object containing `signatures`."""
+  signature_map = _SignatureMap()
+  for name, func in signatures.items():
+    # This true of any signature that came from canonicalize_signatures. Here as
+    # a sanity check on saving; crashing on load (e.g. in _add_signature) would
+    # be more problematic in case future export changes violated these
+    # assertions.
+    assert isinstance(func, defun.ConcreteFunction)
+    assert isinstance(func.structured_outputs, collections_abc.Mapping)
+    # pylint: disable=protected-access
+    if len(func._arg_keywords) == 1:
+      assert 1 == func._num_positional_args
+    else:
+      assert 0 == func._num_positional_args
+    signature_map._add_signature(name, func)
+    # pylint: enable=protected-access
+  return signature_map
+
+
+def validate_augmented_graph_view(augmented_graph_view):
+  """Performs signature-related sanity checks on `augmented_graph_view`."""
+  for name, dep in augmented_graph_view.list_children(
+      augmented_graph_view.root
+  ):
+    if name == SIGNATURE_ATTRIBUTE_NAME:
+      if not isinstance(dep, _SignatureMap):
+        raise ValueError(
+            f"Exporting an object {augmented_graph_view.root} which has an"
+            f" attribute named '{SIGNATURE_ATTRIBUTE_NAME}'. This is a reserved"
+            " attribute used to store SavedModel signatures in objects which"
+            " come from `tf.saved_model.load`. Delete this attribute (e.g."
+            f" `del obj.{SIGNATURE_ATTRIBUTE_NAME}`) before saving if this"
+            " shadowing is acceptable."
+        )
+      break

tensorflow/python/saved_model/simple_save.py

@@ -0,0 +1,84 @@
+# 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.
+# ==============================================================================
+"""SavedModel simple save functionality."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.saved_model import builder
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import signature_def_utils
+from tensorflow.python.saved_model import tag_constants
+from tensorflow.python.util import deprecation
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=['saved_model.simple_save'])
+@deprecation.deprecated(
+    None,
+    'This API was designed for TensorFlow v1. See https://www.tensorflow.org/guide/migrate '
+    'for instructions on how to migrate your code to TensorFlow v2.')
+def simple_save(session, export_dir, inputs, outputs, legacy_init_op=None):
+  """Convenience function to build a SavedModel suitable for serving.
+
+  In many common cases, saving models for serving will be as simple as:
+
+      simple_save(session,
+                  export_dir,
+                  inputs={"x": x, "y": y},
+                  outputs={"z": z})
+
+  Although in many cases it's not necessary to understand all of the many ways
+      to configure a SavedModel, this method has a few practical implications:
+    - It will be treated as a graph for inference / serving (i.e. uses the tag
+      `saved_model.SERVING`)
+    - The SavedModel will load in TensorFlow Serving and supports the
+      [Predict
+      API](https://github.com/tensorflow/serving/blob/master/tensorflow_serving/apis/predict.proto).
+      To use the Classify, Regress, or MultiInference APIs, please see the
+      [SavedModel
+      APIs](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md).
+    - Some TensorFlow ops depend on information on disk or other information
+      called "assets". These are generally handled automatically by adding the
+      assets to the `GraphKeys.ASSET_FILEPATHS` collection. Only assets in that
+      collection are exported; if you need more custom behavior, you'll need to
+      use the
+      [SavedModelBuilder](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/builder.py).
+
+  More information about SavedModel and signatures can be found here:
+  https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/saved_model/README.md.
+
+  Args:
+    session: The TensorFlow session from which to save the meta graph and
+        variables.
+    export_dir: The path to which the SavedModel will be stored.
+    inputs: dict mapping string input names to tensors. These are added
+        to the SignatureDef as the inputs.
+    outputs:  dict mapping string output names to tensors. These are added
+        to the SignatureDef as the outputs.
+    legacy_init_op: Legacy support for op or group of ops to execute after the
+        restore op upon a load.
+  """
+  signature_def_map = {
+      signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
+          signature_def_utils.predict_signature_def(inputs, outputs)
+  }
+  b = builder.SavedModelBuilder(export_dir)
+  b.add_meta_graph_and_variables(
+      session,
+      tags=[tag_constants.SERVING],
+      signature_def_map=signature_def_map,
+      assets_collection=ops.get_collection(ops.GraphKeys.ASSET_FILEPATHS),
+      main_op=legacy_init_op,
+      clear_devices=True)
+  b.save()

tensorflow/python/saved_model/simple_save_test.py

@@ -0,0 +1,103 @@
+# 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.
+# ==============================================================================
+"""Tests for SavedModel simple save functionality."""
+
+import os
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import loader
+from tensorflow.python.saved_model import signature_constants
+from tensorflow.python.saved_model import simple_save
+from tensorflow.python.saved_model import tag_constants
+
+
+class SimpleSaveTest(test.TestCase):
+
+  def _init_and_validate_variable(self, variable_name, variable_value):
+    v = variables.Variable(variable_value, name=variable_name)
+    self.evaluate(variables.global_variables_initializer())
+    self.assertEqual(variable_value, self.evaluate(v))
+    return v
+
+  def _check_variable_info(self, actual_variable, expected_variable):
+    self.assertEqual(actual_variable.name, expected_variable.name)
+    self.assertEqual(actual_variable.dtype, expected_variable.dtype)
+    self.assertEqual(len(actual_variable.shape), len(expected_variable.shape))
+    for i in range(len(actual_variable.shape)):
+      self.assertEqual(actual_variable.shape[i], expected_variable.shape[i])
+
+  def _check_tensor_info(self, actual_tensor_info, expected_tensor):
+    self.assertEqual(actual_tensor_info.name, expected_tensor.name)
+    self.assertEqual(actual_tensor_info.dtype, expected_tensor.dtype)
+    self.assertEqual(
+        len(actual_tensor_info.tensor_shape.dim), len(expected_tensor.shape))
+    for i in range(len(actual_tensor_info.tensor_shape.dim)):
+      self.assertEqual(actual_tensor_info.tensor_shape.dim[i].size,
+                       expected_tensor.shape[i])
+
+  def testSimpleSave(self):
+    """Test simple_save that uses the default parameters."""
+    export_dir = os.path.join(test.get_temp_dir(),
+                              "test_simple_save")
+
+    # Force the test to run in graph mode.
+    # This tests a deprecated v1 API that both requires a session and uses
+    # functionality that does not work with eager tensors (such as
+    # build_tensor_info as called by predict_signature_def).
+    with ops.Graph().as_default():
+      # Initialize input and output variables and save a prediction graph using
+      # the default parameters.
+      with self.session(graph=ops.Graph()) as sess:
+        var_x = self._init_and_validate_variable("var_x", 1)
+        var_y = self._init_and_validate_variable("var_y", 2)
+        inputs = {"x": var_x}
+        outputs = {"y": var_y}
+        simple_save.simple_save(sess, export_dir, inputs, outputs)
+
+      # Restore the graph with a valid tag and check the global variables and
+      # signature def map.
+      with self.session(graph=ops.Graph()) as sess:
+        graph = loader.load(sess, [tag_constants.SERVING], export_dir)
+        collection_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+
+        # Check value and metadata of the saved variables.
+        self.assertEqual(len(collection_vars), 2)
+        self.assertEqual(1, collection_vars[0].eval())
+        self.assertEqual(2, collection_vars[1].eval())
+        self._check_variable_info(collection_vars[0], var_x)
+        self._check_variable_info(collection_vars[1], var_y)
+
+        # Check that the appropriate signature_def_map is created with the
+        # default key and method name, and the specified inputs and outputs.
+        signature_def_map = graph.signature_def
+        self.assertEqual(1, len(signature_def_map))
+        self.assertEqual(signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
+                         list(signature_def_map.keys())[0])
+
+        signature_def = signature_def_map[
+            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY]
+        self.assertEqual(signature_constants.PREDICT_METHOD_NAME,
+                         signature_def.method_name)
+
+        self.assertEqual(1, len(signature_def.inputs))
+        self._check_tensor_info(signature_def.inputs["x"], var_x)
+        self.assertEqual(1, len(signature_def.outputs))
+        self._check_tensor_info(signature_def.outputs["y"], var_y)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/tag_constants.py

@@ -0,0 +1,54 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Common tags used for graphs in SavedModel.
+
+"""
+from tensorflow.python.util.tf_export import tf_export
+
+
+# Tag for the `serving` graph.
+SERVING = "serve"
+tf_export(
+    "saved_model.SERVING",
+    v1=["saved_model.SERVING",
+        "saved_model.tag_constants.SERVING"]).export_constant(
+            __name__, "SERVING")
+
+# Tag for the `training` graph.
+TRAINING = "train"
+tf_export(
+    "saved_model.TRAINING",
+    v1=["saved_model.TRAINING",
+        "saved_model.tag_constants.TRAINING"]).export_constant(
+            __name__, "TRAINING")
+
+# LINT.IfChange
+# Tag for the `eval` graph. Not exported while the export logic is in contrib.
+EVAL = "eval"
+# LINT.ThenChange(//tensorflow/python/keras/saving/utils_v1/unexported_constants.py)
+
+# Tag for the `gpu` graph.
+GPU = "gpu"
+tf_export(
+    "saved_model.GPU", v1=["saved_model.GPU",
+                           "saved_model.tag_constants.GPU"]).export_constant(
+                               __name__, "GPU")
+
+# Tag for the `tpu` graph.
+TPU = "tpu"
+tf_export(
+    "saved_model.TPU", v1=["saved_model.TPU",
+                           "saved_model.tag_constants.TPU"]).export_constant(
+                               __name__, "TPU")

tensorflow/python/saved_model/tests/BUILD

@@ -0,0 +1,24 @@
+# Description:
+# Comprehensive tests for the TensorFlow SavedModels.
+
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_py_strict_test(
+    name = "variable_wrapper_test",
+    srcs = ["variable_wrapper_test.py"],
+    deps = [
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/trackable:autotrackable",
+        "//tensorflow/python/trackable:base",
+    ],
+)

tensorflow/python/saved_model/tests/variable_wrapper_test.py

@@ -0,0 +1,90 @@
+# Copyright 2023 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 serialization of Variable wrapper class using the Trackable API.
+
+A similar implementation is used in Keras Variable, and this test exists to
+ensure the use case continues to work.
+"""
+
+import os
+
+from tensorflow.python.checkpoint import checkpoint
+from tensorflow.python.eager import test
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import save
+from tensorflow.python.trackable import autotrackable
+from tensorflow.python.trackable import base
+
+
+class VariableWrapper(base.Trackable):
+  _should_act_as_resource_variable = True
+
+  def __init__(self, value):
+    self.value = variables.Variable(value)
+
+  @property
+  def _shared_name(self):
+    return self.value._shared_name
+
+  def _serialize_to_tensors(self):
+    return self.value._serialize_to_tensors()
+
+  def _restore_from_tensors(self, restored_tensors):
+    return self.value._restore_from_tensors(restored_tensors)
+
+  def _export_to_saved_model_graph(
+      self, object_map, tensor_map, options, **kwargs
+  ):
+    resource_list = self.value._export_to_saved_model_graph(  # pylint:disable=protected-access
+        object_map, tensor_map, options, **kwargs
+    )
+    object_map[self] = VariableWrapper(object_map[self.value])
+    return resource_list
+
+  def _write_object_proto(self, proto, options):
+    return self.value._write_object_proto(proto, options)
+
+
+class VariableWrapperTest(test.TestCase):
+
+  def test_checkpoint(self):
+    v = VariableWrapper(5.0)
+    root = autotrackable.AutoTrackable()
+    root.v = v
+
+    save_prefix = os.path.join(self.get_temp_dir(), "checkpoint")
+    ckpt = checkpoint.Checkpoint(v=v)
+    save_path = ckpt.save(save_prefix)
+
+    v.value.assign(100)
+    ckpt.restore(save_path)
+    self.assertEqual(5, v.value.numpy())
+
+  def test_export(self):
+    v = VariableWrapper(15.0)
+    root = autotrackable.AutoTrackable()
+    root.v = v
+
+    save_dir = os.path.join(self.get_temp_dir(), "saved_model")
+    save.save(root, save_dir)
+    loaded = load.load(save_dir)
+    self.assertTrue(resource_variable_ops.is_resource_variable(loaded.v))
+    self.assertEqual(15, loaded.v.numpy())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/tracing_utils.py

@@ -0,0 +1,72 @@
+# 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.
+# ==============================================================================
+"""Tracing utilities used by SavedModel."""
+
+from tensorflow.python.checkpoint import saveable_compat
+from tensorflow.python.checkpoint import tensor_callable
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import function as defun
+
+
+def trace_save_and_restore(obj):
+  """Traces `Trackable` serialize- and restore-from-tensors functions.
+
+  Args:
+    obj: A `Trackable` object.
+
+  Returns:
+    A concrete Function.
+  """
+  legacy_name = saveable_compat.get_saveable_name(obj)
+
+  obj_save_fn = obj._serialize_to_tensors  # pylint: disable=protected-access
+  obj_restore_fn = obj._restore_from_tensors  # pylint: disable=protected-access
+
+  if isinstance(obj_save_fn, defun.ConcreteFunction):
+    concrete_save = obj_save_fn
+  else:
+    @def_function.function
+    def save_fn():
+      tensor_dict = obj_save_fn()
+      if any(isinstance(v, tensor_callable.Callable)
+             for v in tensor_dict.values()):
+        raise NotImplementedError(
+            f"Unable to export SavedModel with object of type {type(obj)} "
+            "because it returns a Callable in `_serialize_to_tensors`. "
+            "If you need this functionality please file a feature request.")
+
+      if legacy_name:
+        # If there is a legacy decorator, append the name to the keys.
+        return {f"{legacy_name}{key}": value
+                for key, value in tensor_dict.items()}
+      return tensor_dict
+
+    concrete_save = save_fn.get_concrete_function()
+
+  if isinstance(obj_restore_fn, defun.ConcreteFunction):
+    concrete_restore = obj_restore_fn
+  else:
+    @def_function.function
+    def restore_fn(restored_tensors):
+      if legacy_name:
+        # Do the opposite operation of save_fn()
+        restored_tensors = {key[len(legacy_name):]: value
+                            for key, value in restored_tensors.items()}
+      obj_restore_fn(restored_tensors)
+
+    concrete_restore = restore_fn.get_concrete_function(
+        concrete_save.structured_outputs)
+
+  return concrete_save, concrete_restore

tensorflow/python/saved_model/tracing_utils_test.py

@@ -0,0 +1,67 @@
+# 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 tracing_utils."""
+
+from tensorflow.python.eager import def_function
+
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import variables
+
+from tensorflow.python.saved_model import tracing_utils
+from tensorflow.python.trackable import base
+
+
+class MyTrackable(base.Trackable):
+
+  def __init__(self):
+    self.a = variables.Variable(0)
+    self.b = variables.Variable(1)
+
+  def _serialize_to_tensors(self):
+    return {"a": self.a, "b": self.b}
+
+  def _restore_from_tensors(self, restored_tensors):
+    return control_flow_ops.group(
+        self.a.assign(restored_tensors["a"]),
+        self.b.assign(restored_tensors["b"]))
+
+
+class TracingUtilsTest(test.TestCase):
+
+  def test_trace_save_and_restore(self):
+    t = MyTrackable()
+    save_fn, restore_fn = tracing_utils.trace_save_and_restore(t)
+    self.assertDictEqual({"a": 0, "b": 1}, self.evaluate(save_fn()))
+    restore_fn({"a": constant_op.constant(2), "b": constant_op.constant(3)})
+    self.assertDictEqual({"a": 2, "b": 3}, self.evaluate(save_fn()))
+
+  def test_trace_save_and_restore_concrete(self):
+    t = MyTrackable()
+    t._serialize_to_tensors = (def_function.function(t._serialize_to_tensors)
+                               .get_concrete_function())
+    restored_tensor_spec = t._serialize_to_tensors.structured_outputs
+    # The wrapped tf.function doesn't matter.
+    t._restore_from_tensors = (def_function.function(lambda x: x)
+                               .get_concrete_function(restored_tensor_spec))
+
+    save_fn, restore_fn = tracing_utils.trace_save_and_restore(t)
+    self.assertIs(t._serialize_to_tensors, save_fn)
+    self.assertIs(t._restore_from_tensors, restore_fn)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/utils.py

@@ -0,0 +1,23 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""SavedModel utility functions.
+
+Utility functions to assist with setup and construction of the SavedModel proto.
+"""
+# pylint: disable=unused-import
+from tensorflow.python.saved_model.utils_impl import build_tensor_info
+from tensorflow.python.saved_model.utils_impl import build_tensor_info_from_op
+from tensorflow.python.saved_model.utils_impl import get_tensor_from_tensor_info
+# pylint: enable=unused-import

tensorflow/python/saved_model/utils_impl.py

@@ -0,0 +1,214 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""SavedModel utility functions implementation."""
+
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.protobuf import meta_graph_pb2
+from tensorflow.core.protobuf import struct_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import byte_swap_tensor as bst
+from tensorflow.python.framework import composite_tensor
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+# TensorInfo helpers.
+_DEPRECATION_MSG = (
+    "This API was designed for TensorFlow v1. See "
+    "https://www.tensorflow.org/guide/migrate for instructions on how to "
+    "migrate your code to TensorFlow v2.")
+
+
+@tf_export(
+    v1=["saved_model.build_tensor_info", "saved_model.utils.build_tensor_info"])
+@deprecation.deprecated(None, _DEPRECATION_MSG)
+def build_tensor_info(tensor):
+  """Utility function to build TensorInfo proto from a Tensor.
+
+  Args:
+    tensor: Tensor or SparseTensor whose name, dtype and shape are used to
+        build the TensorInfo. For SparseTensors, the names of the three
+        constituent Tensors are used.
+
+  Returns:
+    A TensorInfo protocol buffer constructed based on the supplied argument.
+
+  Raises:
+    RuntimeError: If eager execution is enabled.
+
+  @compatibility(TF2)
+  This API is not compatible with eager execution as `tensor` needs to be a
+  graph tensor, and there is no replacement for it in TensorFlow 2.x. To start
+  writing programs using TensorFlow 2.x, please refer to the [Effective
+  TensorFlow 2](https://www.tensorflow.org/guide/effective_tf2) guide.
+  @end_compatibility
+  """
+  if context.executing_eagerly():
+    raise RuntimeError("`build_tensor_info` is not supported in eager "
+                       "execution.")
+  return build_tensor_info_internal(tensor)
+
+
+def build_tensor_info_internal(tensor):
+  """Utility function to build TensorInfo proto from a Tensor."""
+  if (isinstance(tensor, composite_tensor.CompositeTensor) and
+      not isinstance(tensor, sparse_tensor.SparseTensor) and
+      not isinstance(tensor, resource_variable_ops.ResourceVariable)):
+    return _build_composite_tensor_info_internal(tensor)
+
+  tensor_info = meta_graph_pb2.TensorInfo(
+      dtype=dtypes.as_dtype(tensor.dtype).as_datatype_enum,
+      tensor_shape=tensor.get_shape().as_proto())
+  if isinstance(tensor, sparse_tensor.SparseTensor):
+    tensor_info.coo_sparse.values_tensor_name = tensor.values.name
+    tensor_info.coo_sparse.indices_tensor_name = tensor.indices.name
+    tensor_info.coo_sparse.dense_shape_tensor_name = tensor.dense_shape.name
+  else:
+    tensor_info.name = tensor.name
+  return tensor_info
+
+
+def _build_composite_tensor_info_internal(tensor):
+  """Utility function to build TensorInfo proto from a CompositeTensor."""
+  spec = tensor._type_spec  # pylint: disable=protected-access
+  tensor_info = meta_graph_pb2.TensorInfo()
+  spec_proto = nested_structure_coder.encode_structure(spec)
+  tensor_info.composite_tensor.type_spec.CopyFrom(spec_proto.type_spec_value)
+  for component in nest.flatten(tensor, expand_composites=True):
+    tensor_info.composite_tensor.components.add().CopyFrom(
+        build_tensor_info_internal(component))
+  return tensor_info
+
+
+def build_tensor_info_from_op(op):
+  """Utility function to build TensorInfo proto from an Op.
+
+  Note that this function should be used with caution. It is strictly restricted
+  to TensorFlow internal use-cases only. Please make sure you do need it before
+  using it.
+
+  This utility function overloads the TensorInfo proto by setting the name to
+  the Op's name, dtype to DT_INVALID and tensor_shape as None. One typical usage
+  is for the Op of the call site for the defunned function:
+  ```python
+    @function.defun
+    def some_variable_initialization_fn(value_a, value_b):
+      a = value_a
+      b = value_b
+
+    value_a = constant_op.constant(1, name="a")
+    value_b = constant_op.constant(2, name="b")
+    op_info = utils.build_op_info(
+        some_variable_initialization_fn(value_a, value_b))
+  ```
+
+  Args:
+    op: An Op whose name is used to build the TensorInfo. The name that points
+        to the Op could be fetched at run time in the Loader session.
+
+  Returns:
+    A TensorInfo protocol buffer constructed based on the supplied argument.
+
+  Raises:
+    RuntimeError: If eager execution is enabled.
+  """
+  if context.executing_eagerly():
+    raise RuntimeError(
+        "`build_tensor_info_from_op` is not supported in eager execution.")
+  return meta_graph_pb2.TensorInfo(
+      dtype=types_pb2.DT_INVALID,
+      tensor_shape=tensor_shape.unknown_shape().as_proto(),
+      name=op.name)
+
+
+@tf_export(v1=["saved_model.get_tensor_from_tensor_info",
+               "saved_model.utils.get_tensor_from_tensor_info"])
+@deprecation.deprecated(None, _DEPRECATION_MSG)
+def get_tensor_from_tensor_info(tensor_info, graph=None, import_scope=None):
+  """Returns the Tensor or CompositeTensor described by a TensorInfo proto.
+
+  Args:
+    tensor_info: A TensorInfo proto describing a Tensor or SparseTensor or
+      CompositeTensor.
+    graph: The tf.Graph in which tensors are looked up. If None, the
+        current default graph is used.
+    import_scope: If not None, names in `tensor_info` are prefixed with this
+        string before lookup.
+
+  Returns:
+    The Tensor or SparseTensor or CompositeTensor in `graph` described by
+    `tensor_info`.
+
+  Raises:
+    KeyError: If `tensor_info` does not correspond to a tensor in `graph`.
+    ValueError: If `tensor_info` is malformed.
+  """
+  graph = graph or ops.get_default_graph()
+  def _get_tensor(name):
+    return graph.get_tensor_by_name(
+        ops.prepend_name_scope(name, import_scope=import_scope))
+  encoding = tensor_info.WhichOneof("encoding")
+  if encoding == "name":
+    return _get_tensor(tensor_info.name)
+  elif encoding == "coo_sparse":
+    return sparse_tensor.SparseTensor(
+        _get_tensor(tensor_info.coo_sparse.indices_tensor_name),
+        _get_tensor(tensor_info.coo_sparse.values_tensor_name),
+        _get_tensor(tensor_info.coo_sparse.dense_shape_tensor_name))
+  elif encoding == "composite_tensor":
+    spec_proto = struct_pb2.StructuredValue(
+        type_spec_value=tensor_info.composite_tensor.type_spec)
+    spec = nested_structure_coder.decode_proto(spec_proto)
+    components = [_get_tensor(component.name) for component in
+                  tensor_info.composite_tensor.components]
+    return nest.pack_sequence_as(spec, components, expand_composites=True)
+  else:
+    raise ValueError(f"Invalid TensorInfo.encoding: {encoding}. Expected `"
+                     "coo_sparse`, `composite_tensor`, or `name` for a dense "
+                     "tensor.")
+
+
+def get_element_from_tensor_info(tensor_info, graph=None, import_scope=None):
+  """Returns the element in the graph described by a TensorInfo proto.
+
+  Args:
+    tensor_info: A TensorInfo proto describing an Op or Tensor by name.
+    graph: The tf.Graph in which tensors are looked up. If None, the current
+      default graph is used.
+    import_scope: If not None, names in `tensor_info` are prefixed with this
+      string before lookup.
+
+  Returns:
+    Op or tensor in `graph` described by `tensor_info`.
+
+  Raises:
+    KeyError: If `tensor_info` does not correspond to an op or tensor in `graph`
+  """
+  graph = graph or ops.get_default_graph()
+  return graph.as_graph_element(
+      ops.prepend_name_scope(tensor_info.name, import_scope=import_scope))
+
+
+def swap_function_tensor_content(meta_graph_def, from_endiness, to_endiness):
+  bst.swap_tensor_content_in_graph_function(
+      meta_graph_def, from_endiness, to_endiness
+  )

tensorflow/python/saved_model/utils_test.py

@@ -0,0 +1,173 @@
+# 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.
+# ==============================================================================
+"""Tests for SavedModel utils."""
+
+from tensorflow.core.framework import types_pb2
+from tensorflow.core.protobuf import struct_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import nested_structure_coder
+from tensorflow.python.saved_model import utils
+
+
+class UtilsTest(test.TestCase):
+
+  @test_util.run_v1_only(
+      "b/120545219: `build_tensor_info` is only available in graph mode.")
+  def testBuildTensorInfoOp(self):
+    x = constant_op.constant(1, name="x")
+    y = constant_op.constant(2, name="y")
+    z = control_flow_ops.group([x, y], name="op_z")
+    z_op_info = utils.build_tensor_info_from_op(z)
+    self.assertEqual("op_z", z_op_info.name)
+    self.assertEqual(types_pb2.DT_INVALID, z_op_info.dtype)
+    self.assertEqual(0, len(z_op_info.tensor_shape.dim))
+
+  @test_util.run_v1_only(
+      "b/120545219: `build_tensor_info` is only available in graph mode.")
+  def testBuildTensorInfoDense(self):
+    x = array_ops.placeholder(dtypes.float32, 1, name="x")
+    x_tensor_info = utils.build_tensor_info(x)
+    self.assertEqual("x:0", x_tensor_info.name)
+    self.assertEqual(types_pb2.DT_FLOAT, x_tensor_info.dtype)
+    self.assertEqual(1, len(x_tensor_info.tensor_shape.dim))
+    self.assertEqual(1, x_tensor_info.tensor_shape.dim[0].size)
+
+  @test_util.run_v1_only(
+      "b/120545219: `build_tensor_info` is only available in graph mode.")
+  def testBuildTensorInfoSparse(self):
+    x = array_ops.sparse_placeholder(dtypes.float32, [42, 69], name="x")
+    x_tensor_info = utils.build_tensor_info(x)
+    self.assertEqual(x.values.name,
+                     x_tensor_info.coo_sparse.values_tensor_name)
+    self.assertEqual(x.indices.name,
+                     x_tensor_info.coo_sparse.indices_tensor_name)
+    self.assertEqual(x.dense_shape.name,
+                     x_tensor_info.coo_sparse.dense_shape_tensor_name)
+    self.assertEqual(types_pb2.DT_FLOAT, x_tensor_info.dtype)
+    self.assertEqual(2, len(x_tensor_info.tensor_shape.dim))
+    self.assertEqual(42, x_tensor_info.tensor_shape.dim[0].size)
+    self.assertEqual(69, x_tensor_info.tensor_shape.dim[1].size)
+
+  @test_util.run_v1_only(
+      "b/120545219: `build_tensor_info` is only available in graph mode.")
+  def testBuildTensorInfoRagged(self):
+    x = ragged_factory_ops.constant([[1, 2], [3]])
+    x_tensor_info = utils.build_tensor_info(x)
+    # Check components
+    self.assertEqual(x.values.name,
+                     x_tensor_info.composite_tensor.components[0].name)
+    self.assertEqual(types_pb2.DT_INT32,
+                     x_tensor_info.composite_tensor.components[0].dtype)
+    self.assertEqual(x.row_splits.name,
+                     x_tensor_info.composite_tensor.components[1].name)
+    self.assertEqual(types_pb2.DT_INT64,
+                     x_tensor_info.composite_tensor.components[1].dtype)
+    # Check type_spec.
+    spec_proto = struct_pb2.StructuredValue(
+        type_spec_value=x_tensor_info.composite_tensor.type_spec)
+    spec = nested_structure_coder.decode_proto(spec_proto)
+    self.assertEqual(spec, x._type_spec)
+
+  def testBuildTensorInfoEager(self):
+    x = constant_op.constant(1, name="x")
+    with context.eager_mode(), self.assertRaisesRegex(
+        RuntimeError, "`build_tensor_info` is not supported"):
+      utils.build_tensor_info(x)
+
+  @test_util.run_v1_only(
+      "b/120545219: `build_tensor_info` is only available in graph mode.")
+  def testGetTensorFromInfoDense(self):
+    expected = array_ops.placeholder(dtypes.float32, 1, name="x")
+    tensor_info = utils.build_tensor_info(expected)
+    actual = utils.get_tensor_from_tensor_info(tensor_info)
+    self.assertIsInstance(actual, tensor.Tensor)
+    self.assertEqual(expected.name, actual.name)
+
+  @test_util.run_v1_only(
+      "b/120545219: `build_tensor_info` is only available in graph mode.")
+  def testGetTensorFromInfoSparse(self):
+    expected = array_ops.sparse_placeholder(dtypes.float32, name="x")
+    tensor_info = utils.build_tensor_info(expected)
+    actual = utils.get_tensor_from_tensor_info(tensor_info)
+    self.assertIsInstance(actual, sparse_tensor.SparseTensor)
+    self.assertEqual(expected.values.name, actual.values.name)
+    self.assertEqual(expected.indices.name, actual.indices.name)
+    self.assertEqual(expected.dense_shape.name, actual.dense_shape.name)
+
+  @test_util.run_v1_only(
+      "b/120545219: `build_tensor_info` is only available in graph mode.")
+  def testGetTensorFromInfoRagged(self):
+    expected = ragged_factory_ops.constant([[1, 2], [3]], name="x")
+    tensor_info = utils.build_tensor_info(expected)
+    actual = utils.get_tensor_from_tensor_info(tensor_info)
+    self.assertIsInstance(actual, ragged_tensor.RaggedTensor)
+    self.assertEqual(expected.values.name, actual.values.name)
+    self.assertEqual(expected.row_splits.name, actual.row_splits.name)
+
+  def testGetTensorFromInfoInOtherGraph(self):
+    with ops.Graph().as_default() as expected_graph:
+      expected = array_ops.placeholder(dtypes.float32, 1, name="right")
+      tensor_info = utils.build_tensor_info(expected)
+    with ops.Graph().as_default():  # Some other graph.
+      array_ops.placeholder(dtypes.float32, 1, name="other")
+    actual = utils.get_tensor_from_tensor_info(tensor_info,
+                                               graph=expected_graph)
+    self.assertIsInstance(actual, tensor.Tensor)
+    self.assertIs(actual.graph, expected_graph)
+    self.assertEqual(expected.name, actual.name)
+
+  def testGetTensorFromInfoInScope(self):
+    # Build a TensorInfo with name "bar/x:0".
+    with ops.Graph().as_default():
+      with ops.name_scope("bar"):
+        unscoped = array_ops.placeholder(dtypes.float32, 1, name="x")
+        tensor_info = utils.build_tensor_info(unscoped)
+        self.assertEqual("bar/x:0", tensor_info.name)
+    # Build a graph with node "foo/bar/x:0", akin to importing into scope foo.
+    with ops.Graph().as_default():
+      with ops.name_scope("foo"):
+        with ops.name_scope("bar"):
+          expected = array_ops.placeholder(dtypes.float32, 1, name="x")
+      self.assertEqual("foo/bar/x:0", expected.name)
+      # Test that tensor is found by prepending the import scope.
+      actual = utils.get_tensor_from_tensor_info(tensor_info,
+                                                 import_scope="foo")
+      self.assertEqual(expected.name, actual.name)
+
+  @test_util.run_v1_only(
+      "b/120545219: `build_tensor_info` is only available in graph mode.")
+  def testGetTensorFromInfoRaisesErrors(self):
+    expected = array_ops.placeholder(dtypes.float32, 1, name="x")
+    tensor_info = utils.build_tensor_info(expected)
+    tensor_info.name = "blah:0"  # Nonexistent name.
+    with self.assertRaises(KeyError):
+      utils.get_tensor_from_tensor_info(tensor_info)
+    tensor_info.ClearField("name")  # Malformed (missing encoding).
+    with self.assertRaises(ValueError):
+      utils.get_tensor_from_tensor_info(tensor_info)
+
+if __name__ == "__main__":
+  test.main()