chore: import upstream snapshot with attribution

tensorflow/python/keras/saving/saved_model/BUILD

@@ -0,0 +1,70 @@
+# 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.
+# ==============================================================================
+
+# Description:
+#   Keras saving and loading files for SavedModel.
+
+load("@xla//third_party/rules_python/python:defs.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow/python/keras/saving:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "load_context",
+    srcs = ["load_context.py"],
+    strict_deps = False,
+    visibility = ["//tensorflow/python/distribute:__pkg__"],
+)
+
+py_library(
+    name = "saved_model",
+    srcs = [
+        "base_serialization.py",
+        "constants.py",
+        "json_utils.py",
+        "layer_serialization.py",
+        "load.py",
+        "load_context.py",
+        "metric_serialization.py",
+        "model_serialization.py",
+        "network_serialization.py",
+        "save.py",
+        "save_impl.py",
+        "serialized_attributes.py",
+        "utils.py",
+    ],
+    srcs_version = "PY3",
+    strict_deps = False,
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_util",
+        "//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:tf_decorator_py",
+    ],
+)

tensorflow/python/keras/saving/saved_model/README.md

@@ -0,0 +1,131 @@
+# Keras SavedModel
+
+For questions, feedback, and feature requests please file a bug/contact kathywu@
+
+## TensorFlow Core SavedModel implementation
+
+In TensorFlow 2.0, all saving and loading implementations revolve around the
+object graph generated from a root trackable object, and all trackable objects
+connected to it through attributes. Program building blocks such as variables,
+assets, and tables, and high level objects like Optimizers and Layers all
+subclass the trackable class. Other objects like TensorFlow functions and
+concrete functions are also saved as nodes in the object graph. When loading a
+SavedModel, the object graph is used to recreate the structure of the original
+object.
+
+Please see the links below for more details:
+
+- [Saved Model Guide](https://www.tensorflow.org/guide/saved_model)
+- [Checkpoint Guide](https://www.tensorflow.org/guide/checkpoint)
+
+## Keras SavedModel implementation
+
+### Overview
+
+Keras object serialization is built on top of the core serialization.
+
+All attributes that impact model execution or inspection are saved to the
+SavedModel to allow the model to be recreated. These attributes are divided into
+three categories:
+
+1. python properties (e.g., layer name, layer config)
+2. objects (e.g. data structures like list of variables or layers)
+3. functions (e.g. call function, loss functions)
+
+Trackable objects and TensorFlow functions are represented as nodes in the
+trackable object graph, and each node in the graph stores information about
+their python properties.
+
+Since many attributes in Keras Layers/Models are not Trackable objects or
+tf.functions, these attributes are wrapped as trackable objects/tf.functions at
+serialization time. For example, `layer.variables` is implemented as a python
+property that appends the lists of trainable/nontrainable variables. During
+serialization, a new Trackable List object is created and saved to the
+`variables` attribute. Another example is the call function. Most models do not
+decorate their call function with `tf.function`, since Keras will take care of
+the graph/function management. When the model is saved, the call function is
+wrapped in a `tf.function` and added to the `__call__` attribute.
+
+
+### `keras_api` attribute
+
+Many attributes are only relevant for revivability. Instead of attaching these
+directly to the exported object, they are saved to a new `keras_api` trackable
+object that is then attached to the exported object. This avoids cluttering the
+exported object with objects/functions that are only used by the Keras library.
+
+For example, `__call__` and `call_and_return_conditional_losses` are functions
+saved for all models. The `__call__` function is attached directly to the
+exported object, while `call_and_return_conditional_losses` is attached to a
+separate object. Say a user saves the model, then loads the SavedModel using the
+core loader (tf.saved_model.load which does not rely on the Keras library to
+revive the model).
+
+The loaded object will have a structure that looks like:
+
+```
+  loaded object -- __call__
+                -- keras_api -- __call__
+                             -- call_and_return_conditional_losses
+```
+
+The two call functions may be accessed through:
+
+  - `loaded.__call__` or `loaded.keras_api.__call__`
+  - `loaded.keras_api.call_and_return_conditional_losses`.
+
+
+### Saving details
+
+Keras Layers use a helper abstract class and an attribute validator class to
+define and standardize the serialization implementation:
+
+- [`SerializationImpl`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/base_serialization.py):
+Ensures that layer python properties are saved as a serialized JSON string in
+the metadata field, and gathers all attributes to save with the Keras object.
+Please see the docstrings in each of the abstract methods/properties to see what
+is required.
+- [`SerializedAttributes`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/serialized_attributes.py?):
+Tracks all of the attributes that must be saved with a Keras object. Objects and
+functions may be specified to be "keras_only", meaning that they will only
+appear in the `keras_api` attribute.
+
+The base `Layer` serialization is defined in
+[`layer_serialization.py`](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/layer_serialization.py).
+See `LayerAttributes` and `LayerSerializationImpl`.
+
+**Adding a new attribute to base Layer SavedModel**
+
+1. Add a new attributes to `LayerAttributes`.
+2. Modify `LayerSerializationImpl` internal methods:
+
+   a. If adding a python property, add the key-value item to the dictionary
+   returned by `_python_properties_internal`
+
+   b.If adding a new object/function, modify the dictionary returned by
+   `_get_serialized_attributes_internal`.
+
+
+**Adding custom serialization for a Layer subclass.**
+
+1. Create a new attribute validator by copying `LayerAttributes`, and add any
+new attributes to serialize.
+2. Subclass `LayerSerializationImpl`
+3. Implement `_python_properties_internal` and/or
+`_get_serialized_attributes_internal` to return the new attributes.
+
+Unless you are modifying the loader (see section below on loading), please keep
+the `object_identifier` the same.
+
+These instructions also carry over for modifying
+[Model](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/model_serialization.py)
+and
+[Network](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/saving/saved_model/network_serialization.py)
+serialization.
+
+
+### Loading details
+
+TODO(kathywu): Will write this section when the loading code is moved into
+\*_serialization.py files.
+

tensorflow/python/keras/saving/saved_model/base_serialization.py

@@ -0,0 +1,134 @@
+# 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.
+# ==============================================================================
+"""Helper classes that list&validate all attributes to serialize to SavedModel."""
+
+import abc
+
+from tensorflow.python.keras.saving.saved_model import json_utils
+from tensorflow.python.keras.saving.saved_model import utils
+
+
+class SavedModelSaver(object, metaclass=abc.ABCMeta):
+  """Saver defining the methods and properties used to serialize Keras objects.
+  """
+
+  def __init__(self, obj):
+    self.obj = obj
+
+  @abc.abstractproperty
+  def object_identifier(self):
+    """String stored in object identifier field in the SavedModel proto.
+
+    Returns:
+      A string with the object identifier, which is used at load time.
+    """
+    raise NotImplementedError
+
+  @property
+  def tracking_metadata(self):
+    """String stored in metadata field in the SavedModel proto.
+
+    Returns:
+      A serialized JSON storing information necessary for recreating this layer.
+    """
+    # TODO(kathywu): check that serialized JSON can be loaded (e.g., if an
+    # object is in the python property)
+    return json_utils.Encoder().encode(self.python_properties)
+
+  def trackable_children(self, serialization_cache):
+    """Lists all Trackable children connected to this object."""
+    if not utils.should_save_traces():
+      return {}
+
+    children = self.objects_to_serialize(serialization_cache)
+    children.update(self.functions_to_serialize(serialization_cache))
+    return children
+
+  @abc.abstractproperty
+  def python_properties(self):
+    """Returns dictionary of python properties to save in the metadata.
+
+    This dictionary must be serializable and deserializable to/from JSON.
+
+    When loading, the items in this dict are used to initialize the object and
+    define attributes in the revived object.
+    """
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def objects_to_serialize(self, serialization_cache):
+    """Returns dictionary of extra checkpointable objects to serialize.
+
+    See `functions_to_serialize` for an explanation of this function's
+    effects.
+
+    Args:
+      serialization_cache: Dictionary passed to all objects in the same object
+        graph during serialization.
+
+    Returns:
+        A dictionary mapping attribute names to checkpointable objects.
+    """
+    raise NotImplementedError
+
+  @abc.abstractmethod
+  def functions_to_serialize(self, serialization_cache):
+    """Returns extra functions to include when serializing a Keras object.
+
+    Normally, when calling exporting an object to SavedModel, only the
+    functions and objects defined by the user are saved. For example:
+
+    ```
+    obj = tf.Module()
+    obj.v = tf.Variable(1.)
+
+    @tf.function
+    def foo(...): ...
+
+    obj.foo = foo
+
+    w = tf.Variable(1.)
+
+    tf.saved_model.save(obj, 'path/to/saved/model')
+    loaded = tf.saved_model.load('path/to/saved/model')
+
+    loaded.v  # Variable with the same value as obj.v
+    loaded.foo  # Equivalent to obj.foo
+    loaded.w  # AttributeError
+    ```
+
+    Assigning trackable objects to attributes creates a graph, which is used for
+    both checkpointing and SavedModel serialization.
+
+    When the graph generated from attribute tracking is insufficient, extra
+    objects and functions may be added at serialization time. For example,
+    most models do not have their call function wrapped with a @tf.function
+    decorator. This results in `model.call` not being saved. Since Keras objects
+    should be revivable from the SavedModel format, the call function is added
+    as an extra function to serialize.
+
+    This function and `objects_to_serialize` is called multiple times when
+    exporting to SavedModel. Please use the cache to avoid generating new
+    functions and objects. A fresh cache is created for each SavedModel export.
+
+    Args:
+      serialization_cache: Dictionary passed to all objects in the same object
+        graph during serialization.
+
+    Returns:
+        A dictionary mapping attribute names to `Function` or
+        `ConcreteFunction`.
+    """
+    raise NotImplementedError

tensorflow/python/keras/saving/saved_model/constants.py

@@ -0,0 +1,47 @@
+# 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.
+# ==============================================================================
+"""Constants for Keras SavedModel serialization."""
+
+# Namespace used to store all attributes added during serialization.
+# e.g. the list of layers can be accessed using `loaded.keras_api.layers`, in an
+# object loaded from `tf.saved_model.load()`.
+KERAS_ATTR = 'keras_api'
+
+# Keys for the serialization cache.
+# Maps to the keras serialization dict {Layer --> SerializedAttributes object}
+KERAS_CACHE_KEY = 'keras_serialized_attributes'
+
+
+# Name of Keras metadata file stored in the SavedModel.
+SAVED_METADATA_PATH = 'keras_metadata.pb'
+
+# Names of SavedObject Keras identifiers.
+INPUT_LAYER_IDENTIFIER = '_tf_keras_input_layer'
+LAYER_IDENTIFIER = '_tf_keras_layer'
+METRIC_IDENTIFIER = '_tf_keras_metric'
+MODEL_IDENTIFIER = '_tf_keras_model'
+NETWORK_IDENTIFIER = '_tf_keras_network'
+RNN_LAYER_IDENTIFIER = '_tf_keras_rnn_layer'
+SEQUENTIAL_IDENTIFIER = '_tf_keras_sequential'
+
+KERAS_OBJECT_IDENTIFIERS = (
+    INPUT_LAYER_IDENTIFIER,
+    LAYER_IDENTIFIER,
+    METRIC_IDENTIFIER,
+    MODEL_IDENTIFIER,
+    NETWORK_IDENTIFIER,
+    RNN_LAYER_IDENTIFIER,
+    SEQUENTIAL_IDENTIFIER,
+)

tensorflow/python/keras/saving/saved_model/json_utils.py

@@ -0,0 +1,144 @@
+# 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.
+# ==============================================================================
+"""Utils for creating and loading the Layer metadata for SavedModel.
+
+These are required to retain the original format of the build input shape, since
+layers and models may have different build behaviors depending on if the shape
+is a list, tuple, or TensorShape. For example, Network.build() will create
+separate inputs if the given input_shape is a list, and will create a single
+input if the given shape is a tuple.
+"""
+
+import collections
+import enum
+import json
+import numpy as np
+import wrapt
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import type_spec_registry
+from tensorflow.python.types import internal
+
+
+class Encoder(json.JSONEncoder):
+  """JSON encoder and decoder that handles TensorShapes and tuples."""
+
+  def default(self, obj):  # pylint: disable=method-hidden
+    """Encodes objects for types that aren't handled by the default encoder."""
+    if isinstance(obj, tensor_shape.TensorShape):
+      items = obj.as_list() if obj.rank is not None else None
+      return {'class_name': 'TensorShape', 'items': items}
+    return get_json_type(obj)
+
+  def encode(self, obj):
+    return super(Encoder, self).encode(_encode_tuple(obj))
+
+
+def _encode_tuple(x):
+  if isinstance(x, tuple):
+    return {'class_name': '__tuple__',
+            'items': tuple(_encode_tuple(i) for i in x)}
+  elif isinstance(x, list):
+    return [_encode_tuple(i) for i in x]
+  elif isinstance(x, dict):
+    return {key: _encode_tuple(value) for key, value in x.items()}
+  else:
+    return x
+
+
+def decode(json_string):
+  return json.loads(json_string, object_hook=_decode_helper)
+
+
+def _decode_helper(obj):
+  """A decoding helper that is TF-object aware."""
+  if isinstance(obj, dict) and 'class_name' in obj:
+    if obj['class_name'] == 'TensorShape':
+      return tensor_shape.TensorShape(obj['items'])
+    elif obj['class_name'] == 'TypeSpec':
+      return type_spec_registry.lookup(obj['type_spec'])._deserialize(  # pylint: disable=protected-access
+          _decode_helper(obj['serialized']))
+    elif obj['class_name'] == '__tuple__':
+      return tuple(_decode_helper(i) for i in obj['items'])
+    elif obj['class_name'] == '__ellipsis__':
+      return Ellipsis
+  return obj
+
+
+def get_json_type(obj):
+  """Serializes any object to a JSON-serializable structure.
+
+  Args:
+      obj: the object to serialize
+
+  Returns:
+      JSON-serializable structure representing `obj`.
+
+  Raises:
+      TypeError: if `obj` cannot be serialized.
+  """
+  # if obj is a serializable Keras class instance
+  # e.g. optimizer, layer
+  if hasattr(obj, 'get_config'):
+    return {'class_name': obj.__class__.__name__, 'config': obj.get_config()}
+
+  # if obj is any numpy type
+  if type(obj).__module__ == np.__name__:
+    if isinstance(obj, np.ndarray):
+      return obj.tolist()
+    else:
+      return obj.item()
+
+  # misc functions (e.g. loss function)
+  if callable(obj):
+    return obj.__name__
+
+  # if obj is a python 'type'
+  if type(obj).__name__ == type.__name__:
+    return obj.__name__
+
+  if isinstance(obj, tensor_shape.Dimension):
+    return obj.value
+
+  if isinstance(obj, tensor_shape.TensorShape):
+    return obj.as_list()
+
+  if isinstance(obj, dtypes.DType):
+    return obj.name
+
+  if isinstance(obj, collections.abc.Mapping):
+    return dict(obj)
+
+  if obj is Ellipsis:
+    return {'class_name': '__ellipsis__'}
+
+  if isinstance(obj, wrapt.ObjectProxy):
+    return obj.__wrapped__
+
+  if isinstance(obj, internal.TypeSpec):
+    try:
+      type_spec_name = type_spec_registry.get_name(type(obj))
+      return {'class_name': 'TypeSpec', 'type_spec': type_spec_name,
+              'serialized': obj._serialize()}  # pylint: disable=protected-access
+    except ValueError:
+      raise ValueError('Unable to serialize {} to JSON, because the TypeSpec '
+                       'class {} has not been registered.'
+                       .format(obj, type(obj)))
+
+  if isinstance(obj, enum.Enum):
+    return obj.value
+
+  raise TypeError('Not JSON Serializable:', obj)

tensorflow/python/keras/saving/saved_model/layer_serialization.py

@@ -0,0 +1,175 @@
+# 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.
+# ==============================================================================
+"""Classes and functions implementing Layer SavedModel serialization."""
+
+from tensorflow.python.keras.mixed_precision import policy
+from tensorflow.python.keras.saving.saved_model import base_serialization
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import save_impl
+from tensorflow.python.keras.saving.saved_model import serialized_attributes
+from tensorflow.python.keras.utils import generic_utils
+from tensorflow.python.trackable import data_structures
+from tensorflow.python.util import nest
+
+
+class LayerSavedModelSaver(base_serialization.SavedModelSaver):
+  """Implements Layer SavedModel serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.LAYER_IDENTIFIER
+
+  @property
+  def python_properties(self):
+    # TODO(kathywu): Add python property validator
+    return self._python_properties_internal()
+
+  def _python_properties_internal(self):
+    """Returns dictionary of all python properties."""
+    # TODO(kathywu): Add support for metrics serialization.
+    # TODO(kathywu): Synchronize with the keras spec (go/keras-json-spec) once
+    # the python config serialization has caught up.
+    metadata = dict(
+        name=self.obj.name,
+        trainable=self.obj.trainable,
+        expects_training_arg=self.obj._expects_training_arg,  # pylint: disable=protected-access
+        dtype=policy.serialize(self.obj._dtype_policy),  # pylint: disable=protected-access
+        batch_input_shape=getattr(self.obj, '_batch_input_shape', None),
+        stateful=self.obj.stateful,
+        must_restore_from_config=self.obj._must_restore_from_config,  # pylint: disable=protected-access
+    )
+
+    metadata.update(get_serialized(self.obj))
+    if self.obj.input_spec is not None:
+      # Layer's input_spec has already been type-checked in the property setter.
+      metadata['input_spec'] = nest.map_structure(
+          lambda x: generic_utils.serialize_keras_object(x) if x else None,
+          self.obj.input_spec)
+    if (self.obj.activity_regularizer is not None and
+        hasattr(self.obj.activity_regularizer, 'get_config')):
+      metadata['activity_regularizer'] = generic_utils.serialize_keras_object(
+          self.obj.activity_regularizer)
+    if self.obj._build_input_shape is not None:  # pylint: disable=protected-access
+      metadata['build_input_shape'] = self.obj._build_input_shape  # pylint: disable=protected-access
+    return metadata
+
+  def objects_to_serialize(self, serialization_cache):
+    return (self._get_serialized_attributes(
+        serialization_cache).objects_to_serialize)
+
+  def functions_to_serialize(self, serialization_cache):
+    return (self._get_serialized_attributes(
+        serialization_cache).functions_to_serialize)
+
+  def _get_serialized_attributes(self, serialization_cache):
+    """Generates or retrieves serialized attributes from cache."""
+    keras_cache = serialization_cache.setdefault(constants.KERAS_CACHE_KEY, {})
+    if self.obj in keras_cache:
+      return keras_cache[self.obj]
+
+    serialized_attr = keras_cache[self.obj] = (
+        serialized_attributes.SerializedAttributes.new(self.obj))
+
+    if (save_impl.should_skip_serialization(self.obj) or
+        self.obj._must_restore_from_config):  # pylint: disable=protected-access
+      return serialized_attr
+
+    object_dict, function_dict = self._get_serialized_attributes_internal(
+        serialization_cache)
+
+    serialized_attr.set_and_validate_objects(object_dict)
+    serialized_attr.set_and_validate_functions(function_dict)
+    return serialized_attr
+
+  def _get_serialized_attributes_internal(self, serialization_cache):
+    """Returns dictionary of serialized attributes."""
+    objects = save_impl.wrap_layer_objects(self.obj, serialization_cache)
+    functions = save_impl.wrap_layer_functions(self.obj, serialization_cache)
+    # Attribute validator requires that the default save signature is added to
+    # function dict, even if the value is None.
+    functions['_default_save_signature'] = None
+    return objects, functions
+
+
+# TODO(kathywu): Move serialization utils (and related utils from
+# generic_utils.py) to a separate file.
+def get_serialized(obj):
+  with generic_utils.skip_failed_serialization():
+    # Store the config dictionary, which may be used when reviving the object.
+    # When loading, the program will attempt to revive the object from config,
+    # and if that fails, the object will be revived from the SavedModel.
+    return generic_utils.serialize_keras_object(obj)
+
+
+class InputLayerSavedModelSaver(base_serialization.SavedModelSaver):
+  """InputLayer serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.INPUT_LAYER_IDENTIFIER
+
+  @property
+  def python_properties(self):
+
+    return dict(
+        class_name=type(self.obj).__name__,
+        name=self.obj.name,
+        dtype=self.obj.dtype,
+        sparse=self.obj.sparse,
+        ragged=self.obj.ragged,
+        batch_input_shape=self.obj._batch_input_shape,  # pylint: disable=protected-access
+        config=self.obj.get_config())
+
+  def objects_to_serialize(self, serialization_cache):
+    return {}
+
+  def functions_to_serialize(self, serialization_cache):
+    return {}
+
+
+class RNNSavedModelSaver(LayerSavedModelSaver):
+  """RNN layer serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.RNN_LAYER_IDENTIFIER
+
+  def _get_serialized_attributes_internal(self, serialization_cache):
+    objects, functions = (
+        super(RNNSavedModelSaver, self)._get_serialized_attributes_internal(
+            serialization_cache))
+    states = data_structures.wrap_or_unwrap(self.obj.states)
+    # SaveModel require all the objects to be Trackable when saving.
+    # If the states is still a tuple after wrap_or_unwrap, it means it doesn't
+    # contain any trackable item within it, eg empty tuple or (None, None) for
+    # stateless ConvLSTM2D. We convert them to list so that wrap_or_unwrap can
+    # make it a Trackable again for saving. When loaded, ConvLSTM2D is
+    # able to handle the tuple/list conversion.
+    if isinstance(states, tuple):
+      states = data_structures.wrap_or_unwrap(list(states))
+    objects['states'] = states
+    return objects, functions
+
+
+class IndexLookupLayerSavedModelSaver(LayerSavedModelSaver):
+  """Index lookup layer serialization."""
+
+  @property
+  def python_properties(self):
+    # TODO(kathywu): Add python property validator
+    metadata = self._python_properties_internal()
+    if metadata['config'].get('has_static_table', False):
+      metadata['config']['vocabulary'] = None
+    return metadata

tensorflow/python/keras/saving/saved_model/load_context.py

@@ -0,0 +1,63 @@
+# 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 storing options for loading a SavedModel."""
+
+import contextlib
+import threading
+
+
+class LoadContext(threading.local):
+  """A context for loading a model."""
+
+  def __init__(self):
+    super(LoadContext, self).__init__()
+    self._entered_load_context = []
+    self._load_options = None
+
+  def set_load_options(self, load_options):
+    self._load_options = load_options
+    self._entered_load_context.append(True)
+
+  def clear_load_options(self):
+    self._load_options = None
+    self._entered_load_context.pop()
+
+  def load_options(self):
+    return self._load_options
+
+  def in_load_context(self):
+    return self._entered_load_context
+
+
+_load_context = LoadContext()
+
+
+@contextlib.contextmanager
+def load_context(load_options):
+  _load_context.set_load_options(load_options)
+  try:
+    yield
+  finally:
+    _load_context.clear_load_options()
+
+
+def get_load_options():
+  """Returns the load options under a load context."""
+  return _load_context.load_options()
+
+
+def in_load_context():
+  """Returns whether under a load context."""
+  return _load_context.in_load_context()

tensorflow/python/keras/saving/saved_model/metric_serialization.py

@@ -0,0 +1,43 @@
+# 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.
+# ==============================================================================
+"""Classes and functions implementing Metrics SavedModel serialization."""
+
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import layer_serialization
+from tensorflow.python.keras.utils import generic_utils
+from tensorflow.python.trackable import data_structures
+
+
+class MetricSavedModelSaver(layer_serialization.LayerSavedModelSaver):
+  """Metric serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.METRIC_IDENTIFIER
+
+  def _python_properties_internal(self):
+    metadata = dict(
+        class_name=generic_utils.get_registered_name(type(self.obj)),
+        name=self.obj.name,
+        dtype=self.obj.dtype)
+    metadata.update(layer_serialization.get_serialized(self.obj))
+    if self.obj._build_input_shape is not None:  # pylint: disable=protected-access
+      metadata['build_input_shape'] = self.obj._build_input_shape  # pylint: disable=protected-access
+    return metadata
+
+  def _get_serialized_attributes_internal(self, unused_serialization_cache):
+    return (dict(variables=data_structures.wrap_or_unwrap(self.obj.variables)),
+            dict())  # TODO(b/135550038): save functions to enable saving
+                     # custom metrics.

tensorflow/python/keras/saving/saved_model/model_serialization.py

@@ -0,0 +1,63 @@
+# 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.
+# ==============================================================================
+"""Classes and functions implementing to Model SavedModel serialization."""
+
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import layer_serialization
+from tensorflow.python.keras.saving.saved_model import save_impl
+
+
+class ModelSavedModelSaver(layer_serialization.LayerSavedModelSaver):
+  """Model SavedModel serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.MODEL_IDENTIFIER
+
+  def _python_properties_internal(self):
+    metadata = super(ModelSavedModelSaver, self)._python_properties_internal()
+    # Network stateful property is dependent on the child layers.
+    metadata.pop('stateful')
+    metadata['is_graph_network'] = self.obj._is_graph_network  # pylint: disable=protected-access
+    metadata['save_spec'] = self.obj._get_save_spec(dynamic_batch=False)  # pylint: disable=protected-access
+
+    metadata.update(
+        saving_utils.model_metadata(
+            self.obj, include_optimizer=True, require_config=False))
+    return metadata
+
+  def _get_serialized_attributes_internal(self, serialization_cache):
+    default_signature = None
+
+    # Create a default signature function if this is the only object in the
+    # cache (i.e. this is the root level object).
+    if len(serialization_cache[constants.KERAS_CACHE_KEY]) == 1:
+      default_signature = save_impl.default_save_signature(self.obj)
+
+    # Other than the default signature function, all other attributes match with
+    # the ones serialized by Layer.
+    objects, functions = (
+        super(ModelSavedModelSaver, self)._get_serialized_attributes_internal(
+            serialization_cache))
+    functions['_default_save_signature'] = default_signature
+    return objects, functions
+
+
+class SequentialSavedModelSaver(ModelSavedModelSaver):
+
+  @property
+  def object_identifier(self):
+    return constants.SEQUENTIAL_IDENTIFIER

tensorflow/python/keras/saving/saved_model/network_serialization.py

@@ -0,0 +1,27 @@
+# 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.
+# ==============================================================================
+"""Classes and functions implementing to Network SavedModel serialization."""
+
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import model_serialization
+
+
+# FunctionalModel serialization is pretty much the same as Model serialization.
+class NetworkSavedModelSaver(model_serialization.ModelSavedModelSaver):
+  """Network serialization."""
+
+  @property
+  def object_identifier(self):
+    return constants.NETWORK_IDENTIFIER

tensorflow/python/keras/saving/saved_model/save.py

@@ -0,0 +1,124 @@
+# 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.
+# ==============================================================================
+"""Keras SavedModel serialization."""
+
+import os
+
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras.protobuf import saved_metadata_pb2
+from tensorflow.python.keras.protobuf import versions_pb2
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import save_impl
+from tensorflow.python.keras.saving.saved_model import utils
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.keras.utils.io_utils import ask_to_proceed_with_overwrite
+from tensorflow.python.platform import gfile
+from tensorflow.python.saved_model import save as save_lib
+
+# To avoid circular dependencies between keras/engine and keras/saving,
+# code in keras/saving must delay imports.
+
+base_layer = LazyLoader(
+    "base_layer", globals(),
+    "tensorflow.python.keras.engine.base_layer")
+training_lib = LazyLoader(
+    "training_lib", globals(),
+    "tensorflow.python.keras.engine.training")
+
+
+def save(model, filepath, overwrite, include_optimizer, signatures=None,
+         options=None, save_traces=True):
+  """Saves a model as a SavedModel to the filepath.
+
+  Args:
+    model: Keras model instance to be saved.
+    filepath: String path to save the model.
+    overwrite: whether to overwrite the existing filepath.
+    include_optimizer: If True, save the model's optimizer state.
+    signatures: Signatures to save with the SavedModel. Applicable to the 'tf'
+      format only. Please see the `signatures` argument in `tf.saved_model.save`
+      for details.
+    options: (only applies to SavedModel format) `tf.saved_model.SaveOptions`
+      object that specifies options for saving to SavedModel.
+    save_traces: (only applies to SavedModel format) When enabled, the
+      SavedModel will store the function traces for each layer. This
+      can be disabled, so that only the configs of each layer are stored.
+      Defaults to `True`. Disabling this will decrease serialization time
+      and reduce file size, but it requires that all custom layers/models
+      implement a `get_config()` method.
+
+  Raises:
+    ValueError: if the model's inputs have not been defined.
+  """
+  # If file exists and should not be overwritten.
+  if not overwrite and os.path.exists(filepath):
+    proceed = ask_to_proceed_with_overwrite(filepath)
+    if not proceed:
+      return
+
+  if save_traces:
+    if save_impl.should_skip_serialization(model):
+      saving_utils.raise_model_input_error(model)
+
+  if not include_optimizer:
+    orig_optimizer = model.optimizer
+    model.optimizer = None
+    # TODO(b/180760306) Change to del model.optimizer if Layer's __delattr__
+    # calls AutoTrackable's __delattr__.
+    model._delete_tracking("optimizer")  # pylint: disable=protected-access
+
+  # Trace all functions and signatures with `training=0` instead of using an
+  # already-set learning phase placeholder.
+  # This is needed for compatibility reasons until learning phase setting
+  # is removed from the public apis.
+  with K.deprecated_internal_learning_phase_scope(0):
+    with utils.keras_option_scope(save_traces):
+      saved_nodes, node_paths = save_lib.save_and_return_nodes(
+          model, filepath, signatures, options)
+
+    # Save all metadata to a separate file in the SavedModel directory.
+    metadata = generate_keras_metadata(saved_nodes, node_paths)
+
+  with gfile.GFile(
+      os.path.join(filepath, constants.SAVED_METADATA_PATH), "wb") as w:
+    w.write(metadata.SerializeToString(deterministic=True))
+
+  if not include_optimizer:
+    model.optimizer = orig_optimizer
+
+
+def generate_keras_metadata(saved_nodes, node_paths):
+  """Constructs a KerasMetadata proto with the metadata of each keras object."""
+  metadata = saved_metadata_pb2.SavedMetadata()
+
+  for node_id, node in enumerate(saved_nodes):
+    if isinstance(node, base_layer.Layer):
+      path = node_paths[node]
+      if not path:
+        node_path = "root"
+      else:
+        node_path = "root.{}".format(
+            ".".join([ref.name for ref in path]))
+
+      metadata.nodes.add(
+          node_id=node_id,
+          node_path=node_path,
+          version=versions_pb2.VersionDef(
+              producer=1, min_consumer=1, bad_consumers=[]),
+          identifier=node._object_identifier,  # pylint: disable=protected-access
+          metadata=node._tracking_metadata)  # pylint: disable=protected-access
+
+  return metadata

tensorflow/python/keras/saving/saved_model/save_impl.py

@@ -0,0 +1,734 @@
+# 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.
+# ==============================================================================
+"""Keras SavedModel serialization.
+
+TODO (kathywu): Move to layer_serialization.py. Some model-specific logic should
+go to model_serialization.py.
+"""
+
+import functools
+import threading
+import weakref
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras.engine import base_layer_utils
+from tensorflow.python.keras.engine import input_spec
+from tensorflow.python.keras.mixed_precision import autocast_variable
+from tensorflow.python.keras.saving import saving_utils
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import load as keras_load
+from tensorflow.python.keras.saving.saved_model import serialized_attributes
+from tensorflow.python.keras.saving.saved_model import utils
+from tensorflow.python.keras.utils import tf_contextlib
+from tensorflow.python.keras.utils import tf_inspect
+from tensorflow.python.keras.utils import tf_utils
+from tensorflow.python.keras.utils import version_utils
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.trackable import data_structures
+from tensorflow.python.util import nest
+from tensorflow.python.util import tf_decorator
+
+
+# To avoid circular dependencies between keras/engine and keras/saving,
+# code in keras/saving must delay imports.
+
+# TODO(b/134426265): Switch back to single-quotes to match the rest of the file
+# once the issue with copybara is fixed.
+# pylint:disable=g-inconsistent-quotes
+base_layer = LazyLoader(
+    "base_layer", globals(),
+    "tensorflow.python.keras.engine.base_layer")
+metrics = LazyLoader("metrics", globals(),
+                     "tensorflow.python.keras.metrics")
+input_layer = LazyLoader(
+    "input_layer", globals(),
+    "tensorflow.python.keras.engine.input_layer")
+training_lib = LazyLoader(
+    "training_lib", globals(),
+    "tensorflow.python.keras.engine.training")
+sequential_lib = LazyLoader(
+    "sequential_lib", globals(),
+    "tensorflow.python.keras.engine.sequential")
+# pylint:enable=g-inconsistent-quotes
+
+
+def should_skip_serialization(layer):
+  """Skip serializing extra objects and functions if layer inputs aren't set."""
+  saved_model_input_spec_set = (isinstance(layer, training_lib.Model) and
+                                layer._saved_model_inputs_spec is not None)  # pylint: disable=protected-access
+  if not layer.built and not saved_model_input_spec_set:
+    logging.warning('Skipping full serialization of Keras layer {}, because '
+                    'it is not built.'.format(layer))
+    return True
+  return False
+
+
+def wrap_layer_objects(layer, serialization_cache):
+  """Returns extra trackable objects to attach to the serialized layer.
+
+  Args:
+    layer: Keras Layer object.
+    serialization_cache: Dictionary shared between all objects during
+      serialization.
+
+  Returns:
+    A dictionary containing all checkpointable objects from a
+    SerializedAttributes object. See LayerAttributes and ModelAttributes for
+    entire list of objects
+  """
+  # Wrap all regularization losses as tf.functions.
+  # First, generate list of all regularization losses in this layer and
+  # sublayers.
+  all_losses = layer._callable_losses[:]  # pylint: disable=protected-access
+  for child_layer in utils.list_all_layers(layer):
+    all_losses.extend(child_layer._callable_losses)  # pylint: disable=protected-access
+  # Next, wrap all loss functions as tf.functions. Use the serialization cache
+  # to store already-wrapped functions.
+  keras_loss_cache = serialization_cache.setdefault('keras_losses', {})
+  wrapped_loss_functions = []
+  for loss_fn in all_losses:
+    if loss_fn in keras_loss_cache:
+      wrapped_loss_functions.append(keras_loss_cache[loss_fn])
+    else:
+      wrapped_loss = _wrap_unconditional_loss(loss_fn, len(keras_loss_cache))
+      keras_loss_cache[loss_fn] = wrapped_loss
+      wrapped_loss_functions.append(wrapped_loss)
+  wrapped_layer_losses = [keras_loss_cache[fn]
+                          for fn in layer._callable_losses[:]]  # pylint: disable=protected-access
+
+  layer_metrics = data_structures.wrap_or_unwrap(
+      {m.name: m for m in layer._metrics})  # pylint: disable=protected-access
+  return dict(
+      variables=data_structures.wrap_or_unwrap(layer.variables),
+      trainable_variables=data_structures.wrap_or_unwrap(
+          layer.trainable_variables),
+      non_trainable_variables=data_structures.wrap_or_unwrap(
+          layer.non_trainable_variables),
+      layers=data_structures.wrap_or_unwrap(utils.list_all_layers(layer)),
+      metrics=data_structures.wrap_or_unwrap(layer.metrics),
+      regularization_losses=data_structures.wrap_or_unwrap(
+          wrapped_loss_functions),
+      layer_regularization_losses=data_structures.wrap_or_unwrap(
+          wrapped_layer_losses),
+      layer_metrics=layer_metrics)
+  # pylint: disable=protected-access
+
+
+def wrap_layer_functions(layer, serialization_cache):
+  """Returns dict of wrapped layer call function and losses in tf.functions.
+
+  Args:
+    layer: Keras Layer object.
+    serialization_cache: Dictionary shared between all objects during
+      serialization.
+
+  Returns:
+    A dictionary containing all keras tf.functions to serialize. See
+    LayerAttributes and ModelAttributes for the list of all attributes.
+  """
+  # Since Sequential models may be modified in place using model.add() or
+  # model.pop(), don't use saved functions.
+  if (isinstance(layer, keras_load.RevivedLayer) and
+      not isinstance(layer, sequential_lib.Sequential)):
+    return {fn_name: getattr(layer.keras_api, fn_name, None)
+            for fn_name in serialized_attributes.LayerAttributes.all_functions}
+
+  # Reset the losses of the layer and its children. The call function in each
+  # child layer is replaced with tf.functions.
+  original_fns = _replace_child_layer_functions(layer, serialization_cache)
+  original_losses = _reset_layer_losses(layer)
+
+  # Wrap all the layer call and activity regularizer functions.
+
+  # Use LayerCallCollection to ensure that all layer call functions (__call__,
+  # call with losses) are traced with the same inputs.
+  call_collection = LayerCallCollection(layer)
+  call_fn_with_losses = call_collection.add_function(
+      _wrap_call_and_conditional_losses(layer),
+      '{}_layer_call_and_return_conditional_losses'.format(layer.name),
+      # If any of this layer's child layers use the training arg, the traced
+      # call functions of this layer will have a training keyword argument. If
+      # the original layer does not expect the training arg, then it will have
+      # to be removed (by setting `match_layer_training_arg`).
+      match_layer_training_arg=True)
+  call_fn = call_collection.add_function(
+      _extract_outputs_from_fn(layer, call_fn_with_losses),
+      '{}_layer_call_fn'.format(layer.name),
+      # Since `call_fn` wraps call_fn_with_losses and not the original call
+      # function, `match_layer_training_arg` should be set to False.
+      match_layer_training_arg=False)
+
+  fns = {'call_and_return_conditional_losses': call_fn_with_losses,
+         '__call__': call_fn}
+
+  if layer._activity_regularizer is not None:  # pylint: disable=protected-access
+    fns['activity_regularizer_fn'] = _wrap_activity_regularizer(layer)
+    fns['call_and_return_all_conditional_losses'] = (
+        call_collection.add_function(
+            _append_activity_regularizer_loss(
+                layer, call_fn_with_losses, fns['activity_regularizer_fn']),
+            '{}_layer_call_and_return_all_conditional_losses'.format(
+                layer.name),
+            match_layer_training_arg=False))
+  else:
+    fns['activity_regularizer_fn'] = None
+    fns['call_and_return_all_conditional_losses'] = call_fn_with_losses
+
+  # Manually trigger traces before restoring the overwritten functions. The
+  # functions are traced within the layer call context to ensure that layer
+  # functions (e.g. add_loss) behave as though running in graph mode.
+  with tracing_scope():
+    call_collection.trace_with_input_signature()
+    with base_layer_utils.call_context().enter(
+        layer, inputs=None, build_graph=True, training=None, saving=True):
+      for fn in fns.values():
+        if fn is not None and fn.input_signature is not None:
+          if isinstance(fn, LayerCall):
+            fn = fn.wrapped_call
+          fn.get_concrete_function()
+
+  # Restore overwritten functions and losses
+  _restore_child_layer_functions(original_fns)
+  _restore_layer_losses(original_losses)
+
+  return fns
+
+
+def default_save_signature(layer):
+  original_losses = _reset_layer_losses(layer)
+  fn = saving_utils.trace_model_call(layer)
+  fn.get_concrete_function()
+  _restore_layer_losses(original_losses)
+  return fn
+
+
+def _replace_child_layer_functions(layer, serialization_cache):
+  """Replaces functions in the children layers with wrapped tf.functions.
+
+  This step allows functions from parent layers to reference the wrapped
+  functions from their children layers instead of retracing the ops.
+
+  This function also resets all losses stored in the layer. These are stored in
+  the returned dictionary. Use `_restore_child_layer_functions` to restore
+  the original attributes.
+
+  Args:
+    layer: Keras Layer object.
+    serialization_cache: Dictionary shared between all objects during
+      serialization.
+
+  Returns:
+    Dictionary mapping layer objects -> original functions and losses:
+      { Child layer 1: {
+          'losses': Original losses,
+          'call': Original call function
+          '_activity_regularizer': Original activity regularizer},
+        Child layer 2: ...
+      }
+  """
+  # pylint: disable=protected-access
+  original_fns = {}
+
+  def replace_layer_functions(child_layer, serialized_fns):
+    """Replaces layer call and activity regularizer with wrapped functions."""
+    original_fns[child_layer] = {
+        'call': child_layer.call,
+        '_activity_regularizer': child_layer._activity_regularizer
+    }
+    with utils.no_automatic_dependency_tracking_scope(child_layer):
+      try:
+        child_layer._activity_regularizer = serialized_fns.get(
+            'activity_regularizer_fn')
+      except AttributeError:
+        # Some layers have an unsettable activity regularizer.
+        pass
+      child_layer.call = utils.use_wrapped_call(
+          child_layer,
+          serialized_fns['call_and_return_conditional_losses'],
+          default_training_value=False)
+
+  def replace_metric_functions(child_layer, serialized_fns):
+    """Replaces metric functions with wrapped functions."""
+    original_fns[child_layer] = {
+        '__call__': child_layer.__call__,
+        'result': child_layer.result,
+        'update_state': child_layer.update_state
+    }
+    with utils.no_automatic_dependency_tracking_scope(child_layer):
+      child_layer.__call__ = serialized_fns['__call__']
+      child_layer.result = serialized_fns['result']
+      child_layer.update_state = serialized_fns['update_state']
+
+  for child_layer in utils.list_all_layers(layer):
+    if isinstance(child_layer, input_layer.InputLayer):
+      continue
+
+    if child_layer not in serialization_cache[constants.KERAS_CACHE_KEY]:
+      serialized_functions = (
+          child_layer._trackable_saved_model_saver._get_serialized_attributes(
+              serialization_cache).functions)
+    else:
+      serialized_functions = (
+          serialization_cache[constants.KERAS_CACHE_KEY][child_layer].functions)
+    if not serialized_functions:
+      # This indicates either:
+      #   - circular dependency, which means the current layer's functions
+      #     should be wrapped first.
+      #   - Child layer's inputs are not defined, so its functions have not been
+      #     wrapped. In this case, no replacement is necessary so move on to the
+      #     next child.
+      continue
+
+    if isinstance(child_layer, metrics.Metric):
+      replace_metric_functions(child_layer, serialized_functions)
+    else:
+      replace_layer_functions(child_layer, serialized_functions)
+
+  return original_fns
+  # pylint: enable=protected-access
+
+
+def _restore_child_layer_functions(original_fns):
+  """Restores attributes replaced with `_replace_child_layer_functions`."""
+  for child_layer, fns in original_fns.items():
+    with utils.no_automatic_dependency_tracking_scope(child_layer):
+      for fn_name, fn in fns.items():
+        try:
+          setattr(child_layer, fn_name, fn)  # pylint: disable=protected-access
+        except AttributeError:
+          pass  # In the case of _activity_regularizer, setting the attribute
+          # may be disallowed.
+
+
+# pylint: disable=protected-access
+def _reset_layer_losses(parent_layer):
+  """Resets losses of layer and its sublayers, and returns original losses."""
+  losses_dict = {}
+  for layer in utils.list_all_layers_and_sublayers(parent_layer):
+    losses_dict[layer] = {'losses': layer._losses[:],
+                          'eager_losses': layer._eager_losses[:]}
+    with utils.no_automatic_dependency_tracking_scope(layer):
+      layer._losses = []
+      layer._eager_losses = []
+  return losses_dict
+
+
+def _restore_layer_losses(losses_dict):
+  for layer in losses_dict:
+    with utils.no_automatic_dependency_tracking_scope(layer):
+      layer._losses = losses_dict[layer]['losses']
+      layer._eager_losses = losses_dict[layer]['eager_losses']
+# pylint: enable=protected-access
+
+
+class LayerTracingContext(threading.local):
+
+  def __init__(self):
+    super(LayerTracingContext, self).__init__()
+    self.enable_call_tracing = False
+    self.trace_queue = []
+
+_thread_local_data = LayerTracingContext()
+
+
+@tf_contextlib.contextmanager
+def tracing_scope():
+  """Enables tracing scope."""
+  # This enables the LayerCallCollection's tracing mechanism to trace all call
+  # functions in the collection.
+  previous_value = _thread_local_data.enable_call_tracing
+  previous_queue = _thread_local_data.trace_queue
+  try:
+    _thread_local_data.enable_call_tracing = True
+    _thread_local_data.trace_queue = []
+    yield
+  finally:
+    # Run traces from the queue.
+    while _thread_local_data.trace_queue:
+      fn, args, kwargs, training = _thread_local_data.trace_queue.pop()
+      if training is not None:
+        with K.deprecated_internal_learning_phase_scope(training):
+          fn.get_concrete_function(*args, **kwargs)
+      else:
+        fn.get_concrete_function(*args, **kwargs)
+    _thread_local_data.trace_queue = previous_queue
+    _thread_local_data.enable_call_tracing = previous_value
+
+
+def add_trace_to_queue(fn, args, kwargs, training=None):
+  if tracing_enabled():
+    _thread_local_data.trace_queue.append(
+        (fn, args[:], kwargs.copy(), training))
+
+
+def tracing_enabled():
+  """Whether to add extra traces to the queue."""
+  return _thread_local_data.enable_call_tracing
+
+
+class LayerCallCollection(object):
+  """Groups wrapped layer call functions.
+
+  This is used to ensure that all layer call functions are traced with the same
+  inputs-
+    - call
+    - call_and_return_conditional_losses
+    - call_and_return_all_conditional_losses
+  """
+
+  def __init__(self, layer):
+    self.layer = layer
+
+    self.layer_call_method = _get_layer_call_method(layer)
+    self._expects_training_arg = utils.layer_uses_training_bool(layer)
+    self._training_arg_index = utils.get_training_arg_index(
+        self.layer_call_method)
+
+    # If the layer call function has kwargs, then the traced function cannot
+    # have an input signature.
+    arg_spec = tf_inspect.getfullargspec(self.layer_call_method)
+    self._has_kwargs = bool(self._expects_training_arg or
+                            arg_spec.defaults or
+                            arg_spec.kwonlyargs or
+                            arg_spec.varkw)
+
+    self._input_signature = self._generate_input_signature(layer)
+    self._functions = weakref.WeakValueDictionary()
+
+    # Get the input argument name from the args.
+    args = arg_spec.args
+    if tf_inspect.ismethod(self.layer_call_method):
+      args = args[1:]
+    self._input_arg_name = args[0] if args else 'inputs'
+
+  def _generate_input_signature(self, layer):
+    """Inspects layer object and returns the inferred input signature.
+
+    Args:
+      layer: Layer object.
+
+    Returns:
+      List of possibly nested TensorSpecs of the layer call function inputs.
+      The list does not contain the `training` argument.
+    """
+    if (isinstance(layer.call, def_function.Function) and
+        layer.call.input_signature is not None):
+      return layer.call.input_signature
+    elif isinstance(layer, training_lib.Model):
+      return saving_utils.model_input_signature(layer)
+    elif (layer.input_spec is not None and
+          layer._use_input_spec_as_call_signature):  # pylint: disable=protected-access
+
+      def to_tensor_spec_or_none(x):
+        spec = input_spec.to_tensor_spec(x, layer._compute_dtype)  # pylint: disable=protected-access
+        # If the shape is too general (e.g. multiple dimensions are allowed),
+        # return None so that separate functions can be generated for each
+        # inferred input signature.
+        # TODO(b/134962016): currently partial signatures are not supported.
+        if spec.shape == tensor_shape.TensorShape(None):
+          return None
+        return spec
+      input_signature = [nest.map_structure(
+          to_tensor_spec_or_none, layer.input_spec)]
+
+      return input_signature
+    else:
+      return None
+
+  def add_trace(self, *args, **kwargs):
+    """Traces all functions with the same args and kwargs.
+
+    Args:
+      *args: Positional args passed to the original function.
+      **kwargs: Keyword args passed to the original function.
+    """
+    args = list(args)
+    kwargs = kwargs.copy()
+
+    for fn in self._functions.values():
+      # TODO(kathywu): Replace arguments with broader shapes defined in the
+      # input signature.
+      if self._expects_training_arg:
+        def trace_with_training(value, fn=fn):
+          utils.set_training_arg(value, self._training_arg_index, args, kwargs)
+          add_trace_to_queue(fn, args, kwargs, value)
+
+        trace_with_training(True)
+        trace_with_training(False)
+      else:
+        add_trace_to_queue(fn, args, kwargs)
+
+  @property
+  def fn_input_signature(self):
+    """Returns input signature for the wrapped layer call function."""
+    if self._has_kwargs:
+      # Input signatures may only describe tensor arguments and kwargs are not
+      # supported.
+      return None
+    if None in nest.flatten(self._input_signature):
+      # TODO(b/134962016): If input signature cannot be partially defined.
+      return None
+    return self._input_signature
+
+  def training_arg_was_passed(self, args, kwargs):
+    if not self.layer._expects_training_arg and self._expects_training_arg:  # pylint: disable=protected-access
+      return (utils.get_training_arg(self._training_arg_index, args, kwargs)
+              is not None)
+    else:
+      return self.layer._call_arg_was_passed(  # pylint: disable=protected-access
+          'training', args, kwargs, inputs_in_args=True)
+
+  def get_training_arg_value(self, args, kwargs):
+    if not self.layer._expects_training_arg and self._expects_training_arg:  # pylint: disable=protected-access
+      return utils.get_training_arg(self._training_arg_index, args, kwargs)
+    else:
+      return self.layer._get_call_arg_value(  # pylint: disable=protected-access
+          'training', args, kwargs, inputs_in_args=True)
+
+  def get_input_arg_value(self, args, kwargs):
+    return self.layer._get_call_arg_value(  # pylint: disable=protected-access
+        self._input_arg_name, args, kwargs, inputs_in_args=True)
+
+  def _maybe_wrap_with_training_arg(self, call_fn, match_layer_training_arg):
+    """Wraps call function with added training argument if necessary."""
+    if not self.layer._expects_training_arg and self._expects_training_arg:  # pylint: disable=protected-access
+      # Add training arg to wrapper function.
+      arg_spec = tf_inspect.getfullargspec(call_fn)
+      args = arg_spec.args + ['training']
+      defaults = list(arg_spec.defaults or [])
+      defaults.append(False)
+      new_arg_spec = tf_inspect.FullArgSpec(
+          args=args,
+          varargs=arg_spec.varargs,
+          varkw=arg_spec.varkw,
+          defaults=defaults,
+          kwonlyargs=arg_spec.kwonlyargs,
+          kwonlydefaults=arg_spec.kwonlydefaults,
+          annotations=arg_spec.annotations)
+
+      # Set new training arg index
+      self._training_arg_index = len(args) - 1
+      if tf_inspect.ismethod(call_fn):
+        self._training_arg_index -= 1
+
+      def wrap_with_training_arg(*args, **kwargs):
+        if match_layer_training_arg:
+          # Remove the training value, since the original call_fn does not
+          # expect a training arg. Instead, the training value will be
+          # propagated using the call context created in LayerCall.
+          args = list(args)
+          kwargs = kwargs.copy()
+          utils.remove_training_arg(self._training_arg_index, args, kwargs)
+        return call_fn(*args, **kwargs)
+
+      return tf_decorator.make_decorator(
+          target=call_fn,
+          decorator_func=wrap_with_training_arg,
+          decorator_argspec=new_arg_spec)
+
+    return call_fn
+
+  def add_function(self, call_fn, name, match_layer_training_arg):
+    """Adds a layer call function to the collection.
+
+    Args:
+      call_fn: a python function
+      name: Name of call function
+      match_layer_training_arg: If True, removes the `training` from the
+        function arguments when calling `call_fn`.
+
+    Returns:
+      LayerCall (tf.function)
+    """
+    fn = LayerCall(
+        self,
+        self._maybe_wrap_with_training_arg(call_fn, match_layer_training_arg),
+        name,
+        input_signature=self.fn_input_signature)
+    self._functions[name] = fn.wrapped_call
+    return fn
+
+  def trace_with_input_signature(self):
+    """Trace with the layer/models inferred input signature if possible."""
+    if (None not in nest.flatten(self._input_signature) and self._has_kwargs):
+      # Manually add traces for layers that have keyword arguments and have
+      # a fully defined input signature.
+      self.add_trace(*self._input_signature)
+
+
+def _filtered_inputs(inputs):
+  return list(filter(tf_utils.is_tensor_or_variable, nest.flatten(inputs)))
+
+
+def layer_call_wrapper(call_collection, method, name):
+  """Ensures layer losses are kept the same, and runs method in call context."""
+
+  # Create wrapper that deals with losses and call context.
+  def wrapper(*args, **kwargs):
+    """Calls method within call context."""
+    layer = call_collection.layer
+    training = None
+    inputs = _filtered_inputs([args, kwargs])
+    # pylint: disable=protected-access
+    if (args or kwargs) and call_collection.training_arg_was_passed(
+        args, kwargs):
+      training = call_collection.get_training_arg_value(args, kwargs)
+    # pylint: enable=protected-access
+    original_losses = _reset_layer_losses(layer)
+    with base_layer_utils.call_context().enter(
+        layer, inputs=inputs, build_graph=False, training=training,
+        saving=True):
+      with autocast_variable.enable_auto_cast_variables(
+          layer._compute_dtype_object):  # pylint: disable=protected-access
+        ret = method(*args, **kwargs)
+    _restore_layer_losses(original_losses)
+    return ret
+
+  # Rename to `name`, since tf.function doesn't have a name argument. Without
+  # this, all functions returned by this method will be named "call", which
+  # would be a nightmare to debug.
+  fn = tf_decorator.make_decorator(target=method, decorator_func=wrapper)
+  fn.__name__ = name
+  return fn
+
+
+class LayerCall(object):
+  """Function that triggers traces of other functions in the same collection."""
+
+  def __init__(self, call_collection, call_fn, name, input_signature):
+    """Initializes a LayerCall object.
+
+    Args:
+      call_collection: a LayerCallCollection, which contains the other layer
+        call functions (e.g. call_with_conditional_losses, call). These
+        functions should be traced with the same arguments.
+      call_fn: A call function.
+      name: Name of the call function.
+      input_signature: Input signature of call_fn (can be None).
+    """
+    self.call_collection = call_collection
+    self.input_signature = input_signature
+    self.wrapped_call = def_function.function(
+        layer_call_wrapper(call_collection, call_fn, name),
+        input_signature=input_signature)
+    self.original_layer_call = call_collection.layer_call_method
+
+  def _maybe_trace(self, args, kwargs):
+    # Trigger traces of other call functions + extra training-arg traces.
+    if tracing_enabled():
+      self.call_collection.add_trace(*args, **kwargs)
+
+  def __call__(self, *args, **kwargs):
+    self._maybe_trace(args, kwargs)
+    return self.wrapped_call(*args, **kwargs)
+
+  def get_concrete_function(self, *args, **kwargs):
+    self._maybe_trace(args, kwargs)
+    return self.wrapped_call.get_concrete_function(*args, **kwargs)
+
+
+def _wrap_call_and_conditional_losses(layer):
+  """Wraps call function that returns a tuple of (outputs, losses).
+
+  The losses returned are conditional on the inputs passed to the call function.
+  Unconditional losses (e.g. weight regularizeration) are wrapped separately.
+
+  Args:
+    layer: a Keras layer object
+
+  Returns:
+    python call function that returns outputs and conditional losses -- excludes
+    activity regularizer
+  """
+  # Create function that generates both outputs and losses
+  layer_call = _get_layer_call_method(layer)
+  def call_and_return_conditional_losses(*args, **kwargs):
+    """Returns layer (call_output, conditional losses) tuple."""
+    call_output = layer_call(*args, **kwargs)
+    if version_utils.is_v1_layer_or_model(layer):
+      conditional_losses = layer.get_losses_for(
+          _filtered_inputs([args, kwargs]))
+    else:
+      conditional_losses = [
+          l for l in layer.losses if not hasattr(l, '_unconditional_loss')
+      ]
+    return call_output, conditional_losses
+
+  return _create_call_fn_decorator(layer, call_and_return_conditional_losses)
+
+
+def _extract_outputs_from_fn(layer, call_and_return_conditional_losses):
+  """Returns a function that returns only call function outputs."""
+  if isinstance(layer, keras_load.RevivedLayer):
+    return layer.keras_api.__call__  # pylint: disable=protected-access
+  def call(inputs, *args, **kwargs):
+    return call_and_return_conditional_losses(inputs, *args, **kwargs)[0]
+  return _create_call_fn_decorator(layer, call)
+
+
+def _append_activity_regularizer_loss(
+    layer, call_fn_with_losses, activity_regularizer_fn):
+  """Appends activity regularizer loss to losses returned by the wrapped fn."""
+  def fn(inputs, *args, **kwargs):
+    outputs, losses = call_fn_with_losses(inputs, *args, **kwargs)
+    losses.append(activity_regularizer_fn(outputs))
+    return outputs, losses
+  return _create_call_fn_decorator(layer, fn)
+
+
+def _create_call_fn_decorator(layer, wrapped_call):
+  call_fn = _get_layer_call_method(layer)
+  fn, arg_spec = utils.maybe_add_training_arg(
+      call_fn, wrapped_call, layer._expects_training_arg,  # pylint: disable=protected-access
+      default_training_value=False)
+  return tf_decorator.make_decorator(
+      target=call_fn,
+      decorator_func=fn,
+      decorator_argspec=arg_spec)
+
+
+def _wrap_unconditional_loss(loss_fn, index):
+  """Wraps callable/unconditional loss, returning a serializable function."""
+  # Extract original loss function from partial function
+  fn = loss_fn.args[0] if isinstance(loss_fn, functools.partial) else loss_fn
+  if isinstance(fn, def_function.Function):
+    return fn
+  else:
+    return def_function.Function(
+        fn, 'loss_fn_{}'.format(index), input_signature=[])
+
+
+def _wrap_activity_regularizer(layer):
+  """Wraps the activity regularizer."""
+  # pylint: disable=protected-access
+  if isinstance(layer._activity_regularizer, def_function.Function):
+    return layer._activity_regularizer
+  return def_function.Function(
+      layer._activity_regularizer,
+      '{}_activity_regularizer'.format(layer.name),
+      input_signature=[
+          tensor_spec.TensorSpec(None, layer._compute_dtype or K.floatx())
+      ])
+  # pylint: enable=protected-access
+
+
+def _get_layer_call_method(layer):
+  if isinstance(layer.call, (def_function.Function)):
+    return layer.call.python_function
+  return layer.call

tensorflow/python/keras/saving/saved_model/serialized_attributes.py

@@ -0,0 +1,311 @@
+# 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.
+# ==============================================================================
+"""Helper classes that list&validate all attributes to serialize to SavedModel.
+"""
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.keras.saving.saved_model import constants
+from tensorflow.python.keras.saving.saved_model import save_impl
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.trackable import base as trackable
+from tensorflow.python.trackable.autotrackable import AutoTrackable
+
+# TODO(b/134426265): Switch back to single-quotes to match the rest of the file
+# once the issue with copybara is fixed.
+# pylint:disable=g-inconsistent-quotes
+base_layer = LazyLoader(
+    "base_layer", globals(),
+    "tensorflow.python.keras.engine.base_layer")
+training_lib = LazyLoader(
+    "training_lib", globals(),
+    "tensorflow.python.keras.engine.training")
+metrics = LazyLoader("metrics", globals(),
+                     "tensorflow.python.keras.metrics")
+recurrent = LazyLoader(
+    "recurrent", globals(),
+    "tensorflow.python.keras.layers.recurrent")
+# pylint:enable=g-inconsistent-quotes
+
+
+class SerializedAttributes(object):
+  """Class that tracks and validates all serialization attributes.
+
+  Keras models contain many Python-defined components. For example, the
+  trainable_variable property lists the model's trainable variables by
+  recursively retrieving the trainable variables from each of the child layers.
+  Another example is model.call, a python function that calls child layers and
+  adds ops to the backend graph.
+
+  Only Tensorflow checkpointable objects and functions can be serialized to
+  SavedModel. Serializing a Keras model as-is results in a checkpointable object
+  that does not resemble a Keras model at all. Thus, extra checkpointable
+  objects and functions must be created during serialization.
+
+  **Defining new serialized attributes**
+  Child classes should be defined using:
+    SerializedAttributes.with_attributes(
+        'name', checkpointable_objects=[...], functions=[...], copy_from=[...])
+  This class is used to cache generated checkpointable objects and functions,
+  ensuring that new objects and functions are generated a single time.
+
+  **Usage during serialization**
+  Each Layer/Model object should have a corresponding instance of
+  SerializedAttributes. Create a new instance by calling
+  `SerializedAttributes.new(obj)`. Objects and functions may be saved using
+  `.set_and_validate_checkpointable_objects`/`.set_and_and_validate_functions`.
+  The properties `.checkpointable_objects` and `.functions` returns the cached
+  values.
+
+  **Adding/changing attributes to save to SavedModel**
+  1. Change the call to `SerializedAttributes.with_attributes` in the correct
+     class:
+     - CommonEndpoints: Base attributes to be added during serialization. If
+       these attributes are present in a Trackable object, it can be
+       deserialized to a Keras Model.
+     - LayerAttributes: Attributes to serialize for Layer objects.
+     - ModelAttributes: Attributes to serialize for Model objects.
+  2. Update class docstring
+  3. Update arguments to any calls to `set_and_validate_*`. For example, if
+     `call_raw_tensors` is added to the ModelAttributes function list, then
+     a `call_raw_tensors` function should be passed to
+     `set_and_validate_functions`.
+
+  **Common endpoints vs other attributes**
+  Only common endpoints are attached directly to the root object. Keras-specific
+  attributes are saved to a separate trackable object with the name "keras_api".
+  The number of objects attached to the root is limited because any naming
+  conflicts will cause user code to break.
+
+  Another reason is that this will only affect users who call
+  `tf.saved_model.load` instead of `tf.keras.models.load_model`. These are
+  advanced users who are likely to have defined their own tf.functions and
+  trackable objects. The added Keras-specific attributes are kept out of the way
+  in the "keras_api" namespace.
+
+  Properties defined in this class may be used to filter out keras-specific
+  attributes:
+  - `functions_to_serialize`: Returns dict of functions to attach to the root
+      object.
+  - `checkpointable_objects_to_serialize`: Returns dict of objects to attach to
+      the root object (including separate trackable object containing
+      keras-specific attributes)
+
+  All changes to the serialized attributes must be backwards-compatible, so
+  attributes should not be removed or modified without sufficient justification.
+  """
+
+  @staticmethod
+  def with_attributes(
+      name, checkpointable_objects=None, functions=None, copy_from=None):
+    """Creates a subclass with all attributes as specified in the arguments.
+
+    Args:
+      name: Name of subclass
+      checkpointable_objects: List of checkpointable objects to be serialized
+        in the SavedModel.
+      functions: List of functions to be serialized in the SavedModel.
+      copy_from: List of other SerializedAttributes subclasses. The returned
+        class will copy checkpoint objects/functions from each subclass.
+
+    Returns:
+      Child class with attributes as defined in the `checkpointable_objects`
+      and `functions` lists.
+    """
+    checkpointable_objects = checkpointable_objects or []
+    functions = functions or []
+
+    if copy_from is not None:
+      for cls in copy_from:
+        checkpointable_objects.extend(cls.all_checkpointable_objects)
+        functions.extend(cls.all_functions)
+
+    classdict = {
+        'all_checkpointable_objects': set(checkpointable_objects),
+        'all_functions': set(functions)}
+    return type(name, (SerializedAttributes,), classdict)
+
+  @staticmethod
+  def new(obj):
+    """Returns a new SerializedAttribute object."""
+    if isinstance(obj, training_lib.Model):
+      return ModelAttributes()
+    elif isinstance(obj, metrics.Metric):
+      return MetricAttributes()
+    elif isinstance(obj, recurrent.RNN):
+      return RNNAttributes()
+    elif isinstance(obj, base_layer.Layer):
+      return LayerAttributes()
+    else:
+      raise TypeError('Internal error during serialization: Expected Keras '
+                      'Layer object, got {} of type {}'.format(obj, type(obj)))
+
+  def __init__(self):
+    self._object_dict = {}
+    self._function_dict = {}
+    self._keras_trackable = AutoTrackable()
+
+  @property
+  def functions(self):
+    """Returns dictionary of all functions."""
+    return {key: value for key, value in self._function_dict.items()
+            if value is not None}
+
+  @property
+  def checkpointable_objects(self):
+    """Returns dictionary of all checkpointable objects."""
+    return {key: value for key, value in self._object_dict.items()
+            if value is not None}
+
+  @property
+  def functions_to_serialize(self):
+    """Returns functions to attach to the root object during serialization."""
+    functions = {}
+    for key, v in self.functions.items():
+      if key in CommonEndpoints.all_functions:
+        functions[key] = (v.wrapped_call if isinstance(v, save_impl.LayerCall)
+                          else v)
+    return functions
+
+  @property
+  def objects_to_serialize(self):
+    """Returns objects to attach to the root object during serialization."""
+    objects = {key: value for key, value in self.checkpointable_objects.items()
+               if key in CommonEndpoints.all_checkpointable_objects}
+    objects[constants.KERAS_ATTR] = self._keras_trackable
+    return objects
+
+  def set_and_validate_functions(self, function_dict):
+    """Saves function dictionary, and validates dictionary values."""
+    for key in self.all_functions:
+      if key in function_dict:
+        if (function_dict[key] is not None and  # Not all functions are required
+            not isinstance(function_dict[key],
+                           (def_function.Function, save_impl.LayerCall))):
+          raise ValueError(
+              'Function dictionary contained a non-function object: {} (for key'
+              ' {})'.format(function_dict[key], key))
+        fn = function_dict[key]
+        self._function_dict[key] = fn
+
+        # Extract TensorFlow `Function` from LayerCall.
+        tf_fn = fn.wrapped_call if isinstance(fn, save_impl.LayerCall) else fn
+        setattr(self._keras_trackable, key, tf_fn)
+      else:
+        raise ValueError('Function {} missing from serialized function dict.'
+                         .format(key))
+    return self.functions
+
+  def set_and_validate_objects(self, object_dict):
+    """Saves objects to a dictionary, and validates the values."""
+    for key in self.all_checkpointable_objects:
+      if key in object_dict:
+        if not isinstance(object_dict[key], trackable.Trackable):
+          raise ValueError(
+              'Object dictionary contained a non-trackable object: {} (for key'
+              ' {})'.format(object_dict[key], key))
+        self._object_dict[key] = object_dict[key]
+        setattr(self._keras_trackable, key, object_dict[key])
+      else:
+        raise ValueError(
+            'Object {} missing from serialized object dict.'.format(key))
+    return self.checkpointable_objects
+
+
+class CommonEndpoints(SerializedAttributes.with_attributes(
+    'CommonEndpoints',
+    checkpointable_objects=['variables', 'trainable_variables',
+                            'regularization_losses'],
+    functions=['__call__', 'call_and_return_all_conditional_losses',
+               '_default_save_signature'])):
+  """Common endpoints shared by all models loadable by Keras.
+
+  List of all attributes:
+    variables: List of all variables in the model and its sublayers.
+    trainable_variables: List of all trainable variables in the model and its
+      sublayers.
+    regularization_losses: List of all unconditional losses (losses not
+      dependent on the inputs) in the model and its sublayers.
+    __call__: Function that takes inputs and returns the outputs of the model
+      call function.
+    call_and_return_all_conditional_losses: Function that returns a tuple of
+      (call function outputs, list of all losses that depend on the inputs).
+    _default_save_signature: Traced model call function. This is only included
+      if the top level exported object is a Keras model.
+  """
+
+
+class LayerAttributes(SerializedAttributes.with_attributes(
+    'LayerAttributes',
+    checkpointable_objects=['non_trainable_variables', 'layers', 'metrics',
+                            'layer_regularization_losses', 'layer_metrics'],
+    functions=['call_and_return_conditional_losses', 'activity_regularizer_fn'],
+    copy_from=[CommonEndpoints]
+    )):
+  """Layer checkpointable objects + functions that are saved to the SavedModel.
+
+  List of all attributes:
+    All attributes from CommonEndpoints
+    non_trainable_variables: List of non-trainable variables in the layer and
+      its sublayers.
+    layers: List of all sublayers.
+    metrics: List of all metrics in the layer and its sublayers.
+    call_and_return_conditional_losses: Function that takes inputs and returns a
+      tuple of (outputs of the call function, list of input-dependent losses).
+      The list of losses excludes the activity regularizer function, which is
+      separate to allow the deserialized Layer object to define a different
+      activity regularizer.
+    activity_regularizer_fn: Callable that returns the activity regularizer loss
+    layer_regularization_losses: List of losses owned only by this layer.
+    layer_metrics: List of metrics owned by this layer.
+  """
+
+
+class ModelAttributes(SerializedAttributes.with_attributes(
+    'ModelAttributes',
+    copy_from=[LayerAttributes])):
+  """Model checkpointable objects + functions that are saved to the SavedModel.
+
+  List of all attributes:
+    All attributes from LayerAttributes (including CommonEndpoints)
+  """
+  # TODO(kathywu): Add attributes `compile_losses` and `compile_metrics`, which
+  #  list all losses and metrics defined by `model.compile`.
+
+
+class MetricAttributes(
+    SerializedAttributes.with_attributes(
+        'MetricAttributes',
+        checkpointable_objects=['variables'],
+        functions=[],
+    )):
+  """Attributes that are added to Metric objects when saved to SavedModel.
+
+  List of all attributes:
+    variables: list of all variables
+  """
+  pass
+
+
+class RNNAttributes(SerializedAttributes.with_attributes(
+    'RNNAttributes',
+    checkpointable_objects=['states'],
+    copy_from=[LayerAttributes])):
+  """RNN checkpointable objects + functions that are saved to the SavedModel.
+
+  List of all attributes:
+    All attributes from LayerAttributes (including CommonEndpoints)
+    states: List of state variables
+  """
+

tensorflow/python/keras/saving/saved_model/utils.py

@@ -0,0 +1,306 @@
+# 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.
+# ==============================================================================
+"""Utility functions shared between SavedModel saving/loading implementations."""
+
+import itertools
+import threading
+import types
+
+from tensorflow.python.eager import context
+from tensorflow.python.keras import backend as K
+from tensorflow.python.keras.engine import base_layer_utils
+from tensorflow.python.keras.utils import control_flow_util
+from tensorflow.python.keras.utils import tf_contextlib
+from tensorflow.python.keras.utils import tf_inspect
+from tensorflow.python.keras.utils.generic_utils import LazyLoader
+from tensorflow.python.util import tf_decorator
+
+
+# pylint:disable=g-inconsistent-quotes
+training_lib = LazyLoader(
+    "training_lib", globals(),
+    "tensorflow.python.keras.engine.training")
+# pylint:enable=g-inconsistent-quotes
+
+
+def use_wrapped_call(layer, call_fn, default_training_value=None,
+                     return_method=False):
+  """Creates fn that adds the losses returned by call_fn & returns the outputs.
+
+  Args:
+    layer: A Keras layer object
+    call_fn: tf.function that takes layer inputs (and possibly a training arg),
+      and returns a tuple of (outputs, list of losses).
+    default_training_value: Default value of the training kwarg. If `None`, the
+      default is `K.learning_phase()`.
+    return_method: Whether to return a method bound to the layer.
+
+  Returns:
+    function that calls call_fn and returns the outputs. Losses returned by
+    call_fn are added to the layer losses.
+  """
+  expects_training_arg = layer_uses_training_bool(layer)
+  if hasattr(call_fn, 'original_layer_call'):  # call_fn is a LayerCall object
+    original_call = call_fn.original_layer_call
+    # In Python 3, callable objects are not compatible with inspect.getargspec
+    call_fn = call_fn.__call__
+  else:
+    original_call = call_fn
+  fn, arg_spec = maybe_add_training_arg(
+      original_call, call_fn, expects_training_arg, default_training_value)
+
+  def return_outputs_and_add_losses(*args, **kwargs):
+    """Returns the outputs from the layer call function, and adds the losses."""
+    if return_method:
+      args = args[1:]
+
+    outputs, losses = fn(*args, **kwargs)
+    layer.add_loss(losses, inputs=True)
+
+    # TODO(kathywu): This is a temporary hack. When a network of layers is
+    # revived from SavedModel, only the top-level layer will have losses. This
+    # causes issues in eager mode because the child layers may have graph losses
+    # (thus model.losses returns a mix of Eager and graph tensors). To fix this,
+    # whenever eager losses are added to one layer, add eager losses to all
+    # child layers. This causes `.losses` to only return eager losses.
+    # pylint: disable=protected-access
+    if context.executing_eagerly():
+      for i in layer._flatten_layers():
+        if i is not layer:
+          i._eager_losses = [base_layer_utils.REVIVED_LOSS_PLACEHOLDER]
+    # pylint: enable=protected-access
+    return outputs
+
+  decorated = tf_decorator.make_decorator(
+      target=call_fn,
+      decorator_func=return_outputs_and_add_losses,
+      decorator_argspec=arg_spec)
+
+  if return_method:
+    return types.MethodType(decorated, layer)
+  else:
+    return decorated
+
+
+def layer_uses_training_bool(layer):
+  """Returns whether this layer or any of its children uses the training arg."""
+  if layer._expects_training_arg:  # pylint: disable=protected-access
+    return True
+  visited = {layer}
+  to_visit = list_all_layers(layer)
+  while to_visit:
+    layer = to_visit.pop()
+    if layer in visited:
+      continue
+    if getattr(layer, '_expects_training_arg', True):
+      return True
+    visited.add(layer)
+    to_visit.extend(list_all_layers(layer))
+  return False
+
+
+def list_all_layers(obj):
+  if isinstance(obj, training_lib.Model):
+    # Handle special case of Sequential, which doesn't return
+    # the `Input` layer.
+    return obj.layers
+  else:
+    return list(obj._flatten_layers(include_self=False, recursive=False))  # pylint: disable=protected-access
+
+
+def list_all_layers_and_sublayers(obj):
+  s = set([obj])
+  s.update(itertools.chain.from_iterable(
+      list_all_layers_and_sublayers(layer) for layer in list_all_layers(obj)))
+  return s
+
+
+def maybe_add_training_arg(
+    original_call, wrapped_call, expects_training_arg, default_training_value):
+  """Decorate call and optionally adds training argument.
+
+  If a layer expects a training argument, this function ensures that 'training'
+  is present in the layer args or kwonly args, with the default training value.
+
+  Args:
+    original_call: Original call function.
+    wrapped_call: Wrapped call function.
+    expects_training_arg: Whether to include 'training' argument.
+    default_training_value: Default value of the training kwarg to include in
+      the arg spec. If `None`, the default is `K.learning_phase()`.
+
+  Returns:
+    Tuple of (
+      function that calls `wrapped_call` and sets the training arg,
+      Argspec of returned function or `None` if the argspec is unchanged)
+  """
+  if not expects_training_arg:
+    return wrapped_call, None
+  def wrap_with_training_arg(*args, **kwargs):
+    """Wrap the `wrapped_call` function, and set training argument."""
+    training_arg_index = get_training_arg_index(original_call)
+    training = get_training_arg(training_arg_index, args, kwargs)
+    if training is None:
+      training = default_training_value or K.learning_phase()
+
+    args = list(args)
+    kwargs = kwargs.copy()
+
+    def replace_training_and_call(training):
+      set_training_arg(training, training_arg_index, args, kwargs)
+      return wrapped_call(*args, **kwargs)
+
+    return control_flow_util.smart_cond(
+        training, lambda: replace_training_and_call(True),
+        lambda: replace_training_and_call(False))
+
+  # Create arg spec for decorated function. If 'training' is not defined in the
+  # args of the original arg spec, then add it to kwonlyargs.
+  arg_spec = tf_inspect.getfullargspec(original_call)
+  defaults = list(arg_spec.defaults) if arg_spec.defaults is not None else []
+
+  kwonlyargs = arg_spec.kwonlyargs
+  kwonlydefaults = arg_spec.kwonlydefaults or {}
+  # Add training arg if it does not exist, or set the default training value.
+  if 'training' not in arg_spec.args:
+    kwonlyargs.append('training')
+    kwonlydefaults['training'] = default_training_value
+  else:
+    index = arg_spec.args.index('training')
+    training_default_index = len(arg_spec.args) - index
+    if (arg_spec.defaults and
+        len(arg_spec.defaults) >= training_default_index and
+        defaults[-training_default_index] is None):
+      defaults[-training_default_index] = default_training_value
+
+  decorator_argspec = tf_inspect.FullArgSpec(
+      args=arg_spec.args,
+      varargs=arg_spec.varargs,
+      varkw=arg_spec.varkw,
+      defaults=defaults,
+      kwonlyargs=kwonlyargs,
+      kwonlydefaults=kwonlydefaults,
+      annotations=arg_spec.annotations)
+  return wrap_with_training_arg, decorator_argspec
+
+
+def get_training_arg_index(call_fn):
+  """Returns the index of 'training' in the layer call function arguments.
+
+  Args:
+    call_fn: Call function.
+
+  Returns:
+    - n: index of 'training' in the call function arguments.
+    - -1: if 'training' is not found in the arguments, but layer.call accepts
+          variable keyword arguments
+    - None: if layer doesn't expect a training argument.
+  """
+  argspec = tf_inspect.getfullargspec(call_fn)
+  if argspec.varargs:
+    # When there are variable args, training must be a keyword arg.
+    if 'training' in argspec.kwonlyargs or argspec.varkw:
+      return -1
+    return None
+  else:
+    # Try to find 'training' in the list of args or kwargs.
+    arg_list = argspec.args
+    if tf_inspect.ismethod(call_fn):
+      arg_list = arg_list[1:]
+
+    if 'training' in arg_list:
+      return arg_list.index('training')
+    elif 'training' in argspec.kwonlyargs or argspec.varkw:
+      return -1
+    return None
+
+
+def set_training_arg(training, index, args, kwargs):
+  if index is None or index < 0 or len(args) <= index:  # index is invalid
+    kwargs['training'] = training
+  else:
+    args[index] = training
+  return args, kwargs
+
+
+def get_training_arg(index, args, kwargs):
+  if index is None or index < 0 or len(args) <= index:  # index is invalid
+    return kwargs.get('training', None)
+  else:
+    return args[index]
+
+
+def remove_training_arg(index, args, kwargs):
+  if index is None or index < 0 or len(args) <= index:  # index is invalid
+    kwargs.pop('training', None)
+  else:
+    args.pop(index)
+
+
+class SaveOptionsContext(threading.local):
+
+  def __init__(self):
+    super(SaveOptionsContext, self).__init__()
+    self.save_traces = True
+
+
+_save_options_context = SaveOptionsContext()
+
+
+@tf_contextlib.contextmanager
+def keras_option_scope(save_traces):
+  previous_value = _save_options_context.save_traces
+  try:
+    _save_options_context.save_traces = save_traces
+    yield
+  finally:
+    _save_options_context.save_traces = previous_value
+
+
+def should_save_traces():
+  """Whether to trace layer functions-can be disabled in the save_traces arg."""
+  return _save_options_context.save_traces
+
+
+@tf_contextlib.contextmanager
+def no_automatic_dependency_tracking_scope(obj):
+  """A context that disables automatic dependency tracking when assigning attrs.
+
+  Objects that inherit from Autotrackable automatically creates dependencies
+  to trackable objects through attribute assignments, and wraps data structures
+  (lists or dicts) with trackable classes. This scope may be used to temporarily
+  disable this behavior. This works similar to the decorator
+  `no_automatic_dependency_tracking`.
+
+  Example usage:
+  ```
+  model = tf.keras.Model()
+  model.arr1 = []  # Creates a ListWrapper object
+  with no_automatic_dependency_tracking_scope(model):
+    model.arr2 = []  # Creates a regular, untracked python list
+  ```
+
+  Args:
+    obj: A trackable object.
+
+  Yields:
+    a scope in which the object doesn't track dependencies.
+  """
+  previous_value = getattr(obj, '_setattr_tracking', True)
+  obj._setattr_tracking = False  # pylint: disable=protected-access
+  try:
+    yield
+  finally:
+    obj._setattr_tracking = previous_value  # pylint: disable=protected-access