chore: import upstream snapshot with attribution

tensorflow/python/saved_model/registration/BUILD

@@ -0,0 +1,100 @@
+# Description:
+# TensorFlow SavedModel Registration.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:internal"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "registration",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+    deps = [
+        ":registration_lib",
+    ],
+)
+
+py_library(
+    name = "registration_lib",
+    srcs = [
+        "registration.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/util:tf_inspect",
+        "@absl_py//absl/logging",
+    ],
+)
+
+tf_py_strict_test(
+    name = "registration_test",
+    srcs = ["registration_test.py"],
+    deps = [
+        ":registration",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/trackable:base",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tf_py_strict_test(
+    name = "registration_saving_test",
+    srcs = ["registration_saving_test.py"],
+    deps = [
+        ":registration",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/client:session",
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:io_ops",
+        "//tensorflow/python/ops:resource_variable_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:loader",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/trackable:autotrackable",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+py_library(
+    name = "test_util",
+    srcs = [
+        "test_util.py",
+    ],
+    strict_deps = True,
+    deps = [
+        ":registration_lib",
+    ],
+)
+
+tf_py_strict_test(
+    name = "tf_registration_test",
+    srcs = ["tf_registration_test.py"],
+    data = [
+        "tf_checkpoint_saver_allowlist.txt",
+        "tf_serializable_allowlist.txt",
+    ],
+    tags = ["no_pip"],
+    deps = [
+        ":registration",
+        ":test_util",
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python/lib/io:file_io",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:resource_loader",
+    ],
+)

tensorflow/python/saved_model/registration/README.md

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

tensorflow/python/saved_model/registration/__init__.py

@@ -0,0 +1,49 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utilities for registering the saving/loading steps for advanced objects."""
+
+from tensorflow.python.saved_model.registration.registration import get_registered_class
+from tensorflow.python.saved_model.registration.registration import get_registered_class_name
+from tensorflow.python.saved_model.registration.registration import get_registered_saver_name
+from tensorflow.python.saved_model.registration.registration import get_restore_function
+from tensorflow.python.saved_model.registration.registration import get_save_function
+from tensorflow.python.saved_model.registration.registration import get_strict_predicate_restore
+
+# These are currently an evolving feature. Use with care.
+from tensorflow.python.saved_model.registration.registration import register_checkpoint_saver
+from tensorflow.python.saved_model.registration.registration import register_serializable
+
+from tensorflow.python.saved_model.registration.registration import RegisteredSaver
+from tensorflow.python.saved_model.registration.registration import validate_restore_function
+
+
+def register_tf_serializable(name=None, predicate=None):
+  """See the docstring for `register_serializable`."""
+  return register_serializable(package="tf", name=name, predicate=predicate)
+
+
+def register_tf_checkpoint_saver(name=None,
+                                 predicate=None,
+                                 save_fn=None,
+                                 restore_fn=None,
+                                 strict_predicate_restore=True):
+  """See the docstring for `register_checkpoint_saver`."""
+  return register_checkpoint_saver(
+      package="tf",
+      name=name,
+      predicate=predicate,
+      save_fn=save_fn,
+      restore_fn=restore_fn,
+      strict_predicate_restore=strict_predicate_restore)

tensorflow/python/saved_model/registration/registration.py

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

tensorflow/python/saved_model/registration/registration_saving_test.py

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

tensorflow/python/saved_model/registration/registration_test.py

@@ -0,0 +1,220 @@
+# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests the registration functions.
+
+For integration tests that use save and load functions, see
+registration_saving_test.py.
+"""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import test
+from tensorflow.python.saved_model import registration
+from tensorflow.python.trackable import base
+
+
+@registration.register_serializable()
+class RegisteredClass(base.Trackable):
+  pass
+
+
+@registration.register_serializable(name="Subclass")
+class RegisteredSubclass(RegisteredClass):
+  pass
+
+
+@registration.register_serializable(package="testing")
+class CustomPackage(base.Trackable):
+  pass
+
+
+@registration.register_serializable(package="testing", name="name")
+class CustomPackageAndName(base.Trackable):
+  pass
+
+
+class SerializableRegistrationTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      (RegisteredClass, "Custom.RegisteredClass"),
+      (RegisteredSubclass, "Custom.Subclass"),
+      (CustomPackage, "testing.CustomPackage"),
+      (CustomPackageAndName, "testing.name"),
+  ])
+  def test_registration(self, expected_cls, expected_name):
+    obj = expected_cls()
+    self.assertEqual(registration.get_registered_class_name(obj), expected_name)
+    self.assertIs(
+        registration.get_registered_class(expected_name), expected_cls)
+
+  def test_get_invalid_name(self):
+    self.assertIsNone(registration.get_registered_class("invalid name"))
+
+  def test_get_unregistered_class(self):
+
+    class NotRegistered(base.Trackable):
+      pass
+
+    no_register = NotRegistered
+    self.assertIsNone(registration.get_registered_class_name(no_register))
+
+  def test_duplicate_registration(self):
+
+    @registration.register_serializable()
+    class Duplicate(base.Trackable):
+      pass
+
+    dup = Duplicate()
+    self.assertEqual(
+        registration.get_registered_class_name(dup), "Custom.Duplicate")
+    # Registrations with different names are ok.
+    registration.register_serializable(package="duplicate")(Duplicate)
+    # Registrations are checked in reverse order.
+    self.assertEqual(
+        registration.get_registered_class_name(dup), "duplicate.Duplicate")
+    # Both names should resolve to the same class.
+    self.assertIs(
+        registration.get_registered_class("Custom.Duplicate"), Duplicate)
+    self.assertIs(
+        registration.get_registered_class("duplicate.Duplicate"), Duplicate)
+
+    # Registrations of the same name fails
+    with self.assertRaisesRegex(ValueError, "already been registered"):
+      registration.register_serializable(
+          package="testing", name="CustomPackage")(
+              Duplicate)
+
+  def test_register_non_class_fails(self):
+    obj = RegisteredClass()
+    with self.assertRaisesRegex(TypeError, "must be a class"):
+      registration.register_serializable()(obj)
+
+  def test_register_bad_predicate_fails(self):
+    with self.assertRaisesRegex(TypeError, "must be callable"):
+      registration.register_serializable(predicate=0)(RegisteredClass)
+
+  def test_predicate(self):
+
+    class Predicate(base.Trackable):
+
+      def __init__(self, register_this):
+        self.register_this = register_this
+
+    registration.register_serializable(
+        name="RegisterThisOnlyTrue",
+        predicate=lambda x: isinstance(x, Predicate) and x.register_this)(
+            Predicate)
+
+    a = Predicate(True)
+    b = Predicate(False)
+    self.assertEqual(
+        registration.get_registered_class_name(a),
+        "Custom.RegisterThisOnlyTrue")
+    self.assertIsNone(registration.get_registered_class_name(b))
+
+    registration.register_serializable(
+        name="RegisterAllPredicate",
+        predicate=lambda x: isinstance(x, Predicate))(
+            Predicate)
+
+    self.assertEqual(
+        registration.get_registered_class_name(a),
+        "Custom.RegisterAllPredicate")
+    self.assertEqual(
+        registration.get_registered_class_name(b),
+        "Custom.RegisterAllPredicate")
+
+
+class CheckpointSaverRegistrationTest(test.TestCase):
+
+  def test_invalid_registration(self):
+    with self.assertRaisesRegex(TypeError, "must be string"):
+      registration.register_checkpoint_saver(
+          package=None,
+          name="test",
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(TypeError, "must be string"):
+      registration.register_checkpoint_saver(
+          name=None,
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(ValueError,
+                                "Invalid registered checkpoint saver."):
+      registration.register_checkpoint_saver(
+          package="package",
+          name="t/est",
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(ValueError,
+                                "Invalid registered checkpoint saver."):
+      registration.register_checkpoint_saver(
+          package="package",
+          name="t/est",
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(
+        TypeError,
+        "The predicate registered to a checkpoint saver must be callable"
+    ):
+      registration.register_checkpoint_saver(
+          name="test",
+          predicate=None,
+          save_fn=lambda: None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(TypeError, "The save_fn must be callable"):
+      registration.register_checkpoint_saver(
+          name="test",
+          predicate=lambda: None,
+          save_fn=None,
+          restore_fn=lambda: None)
+    with self.assertRaisesRegex(TypeError, "The restore_fn must be callable"):
+      registration.register_checkpoint_saver(
+          name="test",
+          predicate=lambda: None,
+          save_fn=lambda: None,
+          restore_fn=None)
+
+  def test_registration(self):
+    registration.register_checkpoint_saver(
+        package="Testing",
+        name="test_predicate",
+        predicate=lambda x: hasattr(x, "check_attr"),
+        save_fn=lambda: "save",
+        restore_fn=lambda: "restore")
+    x = base.Trackable()
+    self.assertIsNone(registration.get_registered_saver_name(x))
+
+    x.check_attr = 1
+    saver_name = registration.get_registered_saver_name(x)
+    self.assertEqual(saver_name, "Testing.test_predicate")
+
+    self.assertEqual(registration.get_save_function(saver_name)(), "save")
+    self.assertEqual(registration.get_restore_function(saver_name)(), "restore")
+
+    registration.validate_restore_function(x, "Testing.test_predicate")
+    with self.assertRaisesRegex(ValueError, "saver cannot be found"):
+      registration.validate_restore_function(x, "Invalid.name")
+    x2 = base.Trackable()
+    with self.assertRaisesRegex(ValueError, "saver cannot be used"):
+      registration.validate_restore_function(x2, "Testing.test_predicate")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/registration/test_util.py

@@ -0,0 +1,25 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utils for testing registered objects."""
+from tensorflow.python.saved_model.registration import registration as registration_lib
+# pylint: disable=protected-access
+
+
+def get_all_registered_serializables():
+  return registration_lib._class_registry.get_registrations()
+
+
+def get_all_registered_checkpoint_savers():
+  return registration_lib._saver_registry.get_registrations()

tensorflow/python/saved_model/registration/tf_checkpoint_saver_allowlist.txt

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

tensorflow/python/saved_model/registration/tf_registration_test.py

@@ -0,0 +1,110 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests TensorFlow registered objects.
+
+New or deleted registrations must be approved by the Saved Model team.
+"""
+import os
+
+import tensorflow as tf
+
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.platform import resource_loader
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import registration
+from tensorflow.python.saved_model.registration import test_util
+
+SERIALIZABLE_ALLOWLIST = os.path.join(resource_loader.get_data_files_path(),
+                                      "tf_serializable_allowlist.txt")
+CHECKPOINT_SAVER_ALLOWLIST = os.path.join(resource_loader.get_data_files_path(),
+                                          "tf_checkpoint_saver_allowlist.txt")
+
+# call TPUEmbedding to load the file and run its registration
+_ = tf.tpu.experimental.embedding.TPUEmbedding
+# call save to load the file and run its registration
+_ = tf.saved_model.save
+
+
+@registration.register_tf_serializable()
+class NotInAllowlistExample(tf.Variable):
+  pass
+
+
+registration.register_tf_checkpoint_saver(
+    name="TestDummySaver",
+    predicate=lambda: False,
+    save_fn=lambda: None,
+    restore_fn=lambda: None)
+
+registration.register_tf_checkpoint_saver(
+    name="NotInAllowlistExample",
+    predicate=lambda: False,
+    save_fn=lambda: None,
+    restore_fn=lambda: None)
+
+
+class AllowlistTest(test.TestCase):
+
+  def _test_allowlist(self, allowlist_file, registrations):
+    allowlist = set([
+        s.strip()
+        for s in file_io.read_file_to_string(allowlist_file).splitlines()
+        if s.strip() and not s.startswith("#")
+    ])
+    registered_names = set(registrations)
+
+    missing_from_allowlist = registered_names - allowlist
+    self.assertIn("tf.NotInAllowlistExample", missing_from_allowlist)
+    missing_from_allowlist.remove("tf.NotInAllowlistExample")
+
+    if missing_from_allowlist:
+      msg = ("[NEEDS ATTENTION] Registered names found that were not added to "
+             "the allowlist. Add the following names to the list:\n\t" +
+             "\n\t".join(missing_from_allowlist))
+    else:
+      msg = "[OK] All registered names have been added to the allowlist.  ✓"
+
+    msg += "\n\n"
+
+    missing_registered_names = allowlist - registered_names
+    if missing_registered_names:
+      msg += ("[NEEDS ATTENTION] Some names were found in the allowlist that "
+              "are not registered in TensorFlow. This could mean that a "
+              "registration was removed from the codebase. If this was "
+              "intended, please remove the following from the allowlist:\n\t" +
+              "\n\t".join(missing_registered_names))
+    else:
+      msg += ("[OK] All allowlisted names are registered in the Tensorflow "
+              "library. ✓")
+
+    if missing_from_allowlist or missing_registered_names:
+      raise AssertionError(
+          "Error found in the registration allowlist.\nPlease update the "
+          "allowlist at .../tensorflow/python/saved_model/registration/"
+          f"{os.path.basename(allowlist_file)}.\n\n" + msg +
+          "\n\nAfter making changes, request approval from "
+          " tf-saved-model-owners@.")
+
+  def test_checkpoint_savers(self):
+    self._test_allowlist(CHECKPOINT_SAVER_ALLOWLIST,
+                         test_util.get_all_registered_checkpoint_savers())
+
+  def test_serializables(self):
+    self._test_allowlist(SERIALIZABLE_ALLOWLIST,
+                         test_util.get_all_registered_serializables())
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/saved_model/registration/tf_serializable_allowlist.txt

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