chore: import upstream snapshot with attribution

tensorflow/python/tpu/tests/BUILD

@@ -0,0 +1,484 @@
+# Description: Tests defined for Cloud TPUs
+
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow/python/tpu:tpu.bzl", "tpu_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "tpu_embedding_base_test",
+    srcs = ["tpu_embedding_base_test.py"],
+    deps = [
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
+        "//tensorflow/python/eager:remote",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:init_ops_v2",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:math_ops_gen",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/tpu:tpu_embedding_v2",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "//tensorflow/python/util:nest",
+        "//third_party/py/numpy",
+        "@absl_py//absl/flags",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_checkpoint_test",
+    srcs = [
+        "tpu_embedding_v2_checkpoint_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/module",
+        "//tensorflow/python/ops:init_ops_v2",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/tpu:tpu_embedding_for_serving",
+        "//tensorflow/python/tpu:tpu_embedding_v2",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "//tensorflow/python/training:checkpoint_utils",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_mp_strategy_test",
+    srcs = [
+        "tpu_embedding_v2_mp_strategy_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/distribute:tpu_strategy",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:init_ops_v2",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:math_ops_gen",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/tpu:device_assignment",
+        "//tensorflow/python/tpu:tpu_embedding_v2",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_enqueue_mode_test",
+    srcs = [
+        "tpu_embedding_v2_enqueue_mode_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:nest",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_invalid_input_test",
+    srcs = [
+        "tpu_embedding_v2_invalid_input_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:config",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/tpu:tpu_embedding_v2",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_valid_input_test",
+    srcs = [
+        "tpu_embedding_v2_valid_input_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:init_ops_v2",
+        "//tensorflow/python/ops/ragged:ragged_tensor",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/tpu:tpu_embedding_v2",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "//tensorflow/python/util:nest",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_hd_valid_input_test",
+    srcs = [
+        "tpu_embedding_v2_hd_valid_input_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_hd_invalid_input_test",
+    srcs = [
+        "tpu_embedding_v2_hd_invalid_input_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_sequence_feature_test",
+    srcs = [
+        "tpu_embedding_v2_sequence_feature_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+pytype_strict_library(
+    name = "tpu_embedding_v2_correctness_base_test",
+    srcs = ["tpu_embedding_v2_correctness_base_test.py"],
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_sparse_training_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_sparse_training_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_sparse_forward_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_sparse_forward_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_ragged_training_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_ragged_training_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_ragged_forward_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_ragged_forward_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_hd_sparse_training_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_hd_sparse_training_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_hd_sparse_forward_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_hd_sparse_forward_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_hd_ragged_training_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_hd_ragged_training_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_hd_ragged_forward_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_hd_ragged_forward_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_dense_lookup_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_dense_lookup_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_correctness_sequence_feature_test",
+    srcs = [
+        "tpu_embedding_v2_correctness_sequence_feature_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_v2_correctness_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/tpu:tpu_embedding_v2",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "//tensorflow/python/util:nest",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v2_initialization_test",
+    srcs = [
+        "tpu_embedding_v2_initialization_test.py",
+    ],
+    disable_experimental = True,
+    disable_mlir_bridge = False,
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:init_ops_v2",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/tpu:tpu_embedding_v2",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "//third_party/py/numpy",
+    ],
+)
+
+### tpu embedding v1 tests
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v1_checkpoint_test",
+    srcs = [
+        "tpu_embedding_v1_checkpoint_test.py",
+    ],
+    disable_mlir_bridge = False,
+    tags = ["no_oss"],
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/checkpoint",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_spec",
+        "//tensorflow/python/ops:init_ops_v2",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/saved_model:load",
+        "//tensorflow/python/saved_model:save",
+        "//tensorflow/python/tpu:tpu_embedding_for_serving",
+        "//tensorflow/python/tpu:tpu_embedding_v1",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "//tensorflow/python/training:checkpoint_utils",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_embedding_v1_correctness_test",
+    srcs = [
+        "tpu_embedding_v1_correctness_test.py",
+    ],
+    disable_mlir_bridge = False,
+    tags = ["no_oss"],
+    deps = [
+        ":tpu_embedding_base_test",
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute:distribute_lib",
+        "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/tpu:tpu_embedding_v1",
+        "//tensorflow/python/tpu:tpu_embedding_v2_utils",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+tpu_py_strict_test(
+    name = "tpu_initialization_test",
+    srcs = [
+        "tpu_initialization_test.py",
+    ],
+    disable_mlir_bridge = False,
+    disable_tfrt = False,
+    disable_v3_4chips = False,
+    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/python/compat:v2_compat",
+        "//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)

tensorflow/python/tpu/tests/tpu_embedding_base_test.py

@@ -0,0 +1,661 @@
+# 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.
+# ==============================================================================
+"""Base Class for TPU Embedding tests."""
+
+import os
+from typing import Tuple
+
+from absl import flags
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import tpu_strategy
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.eager import remote
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v2
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.util import nest
+
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('tpu', '', 'Name of TPU to connect to.')
+flags.DEFINE_string('project', None, 'Name of GCP project with TPU.')
+flags.DEFINE_string('zone', None, 'Name of GCP zone with TPU.')
+flags.DEFINE_string('model_dir', os.environ.get('TEST_TMPDIR'),
+                    'A temporary directory.')
+
+
+class TPUEmbeddingBaseTest(parameterized.TestCase, test.TestCase):
+
+  def skip_if_oss(self):
+    if FLAGS.project is not None or FLAGS.zone is not None:
+      self.skipTest(
+          'Skipping tests for oss as it is slow to run every test in cloud tpu.'
+      )
+
+  def setUp(self):
+    super(TPUEmbeddingBaseTest, self).setUp()
+    self.embedding_values = np.array(list(range(32)), dtype=np.float64)
+    self.initializer = init_ops_v2.Constant(self.embedding_values)
+    # Embedding for video initialized to
+    # 0 1 2 3
+    # 4 5 6 7
+    # ...
+    self.table_video = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=8,
+        dim=4,
+        initializer=self.initializer,
+        combiner='sum',
+        name='video')
+    # Embedding for user initialized to
+    # 0 1
+    # 2 3
+    # 4 5
+    # 6 7
+    # ...
+    self.table_user = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=16,
+        dim=2,
+        initializer=self.initializer,
+        combiner='mean',
+        name='user')
+    self.feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+        table=self.table_video, name='watched'),
+                           tpu_embedding_v2_utils.FeatureConfig(
+                               table=self.table_video, name='favorited'),
+                           tpu_embedding_v2_utils.FeatureConfig(
+                               table=self.table_user, name='friends'))
+
+    self.batch_size = 2
+    self.data_batch_size = 4
+
+    # One (global) batch of inputs
+    # sparse tensor for watched:
+    # row 0: 0
+    # row 1: 0, 1
+    # row 2: 0, 1
+    # row 3: 1
+    self.feature_watched_indices = [[0, 0], [1, 0], [1, 1], [2, 0], [2, 1],
+                                    [3, 0]]
+    self.feature_watched_values = [0, 0, 1, 0, 1, 1]
+    self.feature_watched_row_lengths = [1, 2, 2, 1]
+    # sparse tensor for favorited:
+    # row 0: 0, 1
+    # row 1: 1
+    # row 2: 0
+    # row 3: 0, 1
+    self.feature_favorited_indices = [[0, 0], [0, 1], [1, 0], [2, 0], [3, 0],
+                                      [3, 1]]
+    self.feature_favorited_values = [0, 1, 1, 0, 0, 1]
+    self.feature_favorited_row_lengths = [2, 1, 1, 2]
+    # sparse tensor for friends:
+    # row 0: 3
+    # row 1: 0, 1, 2
+    # row 2: 3
+    # row 3: 0, 1, 2
+    self.feature_friends_indices = [[0, 0], [1, 0], [1, 1], [1, 2], [2, 0],
+                                    [3, 0], [3, 1], [3, 2]]
+    self.feature_friends_values = [3, 0, 1, 2, 3, 0, 1, 2]
+    self.feature_friends_row_lengths = [1, 3, 1, 3]
+    self.resolver = None
+
+    # Basically we are expand the dims of the old feature by 1 and repeat
+    # batch size times for the first dimension.
+    def create_hight_dimensional_indices(indices):
+      indices = np.array(indices, dtype=np.int32)
+      batch_size_index = np.repeat(
+          np.arange(self.data_batch_size), len(indices)).reshape(-1, 1)
+      repeated_indices = np.tile(indices, (self.data_batch_size, 1))
+      return np.concatenate([batch_size_index, repeated_indices], axis=1)
+
+    # Create high dimensional features with shape(4, 4, 2)
+    self.feature_watched_indices_high_dimensional = create_hight_dimensional_indices(
+        self.feature_watched_indices)
+    self.feature_watched_values_high_dimensional = self.feature_watched_values * self.data_batch_size
+    self.feature_watched_row_lengths_high_dimensional = self.feature_watched_row_lengths * self.data_batch_size
+
+    # Create high dimensional features with shape(4, 4, 2)
+    self.feature_favorited_indices_high_dimensional = create_hight_dimensional_indices(
+        self.feature_favorited_indices)
+    self.feature_favorited_values_high_dimensional = self.feature_favorited_values * self.data_batch_size
+    self.feature_favorited_row_lengths_high_dimensional = self.feature_favorited_row_lengths * self.data_batch_size
+
+    # Create high dimensional features with shape(4, 4, 3)
+    self.feature_friends_indices_high_dimensional = create_hight_dimensional_indices(
+        self.feature_friends_indices)
+    self.feature_friends_values_high_dimensional = self.feature_friends_values * self.data_batch_size
+    self.feature_friends_row_lengths_high_dimensional = self.feature_friends_row_lengths * self.data_batch_size
+
+  def _init_tpu_system(self):
+    self.resolver = tpu_cluster_resolver.TPUClusterResolver(
+        tpu=FLAGS.tpu, zone=FLAGS.zone, project=FLAGS.project)
+    if hasattr(self.resolver, '_cloud_tpu_client'):
+      self.resolver._cloud_tpu_client.configure_tpu_version(
+          version='nightly', restart_type='always')
+    remote.connect_to_cluster(self.resolver)
+    return tpu_cluster_resolver.initialize_tpu_system(self.resolver)
+
+  def _get_strategy(self):
+    _ = self._init_tpu_system()
+    return tpu_strategy.TPUStrategy(self.resolver)
+
+  def _create_mid_level(self, optimizer=None):
+    # Create `TPUEmbedding` object.
+    if optimizer is None:
+      optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+
+    return tpu_embedding_v2.TPUEmbedding(
+        feature_config=self.feature_config, optimizer=optimizer)
+
+  def _create_strategy_and_mid_level(self, optimizer_name) -> Tuple[
+      tpu_strategy.TPUStrategy, tpu_embedding_v2.TPUEmbedding,
+      tpu_embedding_v2_utils._Optimizer]:
+    strategy = self._get_strategy()
+
+    with strategy.scope():
+      if optimizer_name == 'sgd':
+        optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      elif optimizer_name == 'adagrad':
+        optimizer = tpu_embedding_v2_utils.Adagrad(learning_rate=0.1)
+      elif optimizer_name == 'adam':
+        optimizer = tpu_embedding_v2_utils.Adam(learning_rate=0.1)
+      elif optimizer_name == 'ftrl':
+        optimizer = tpu_embedding_v2_utils.FTRL(learning_rate=0.1)
+      elif optimizer_name == 'adagrad_momentum':
+        optimizer = tpu_embedding_v2_utils.AdagradMomentum(
+            learning_rate=0.1,
+            momentum=0.9,
+            use_nesterov=True,
+            exponent=3.0,
+            epsilon=0.1,
+            beta2=0.9)
+      else:
+        raise ValueError('optimizer is not recognized: ', optimizer_name)
+      mid_level_api = self._create_mid_level(optimizer=optimizer)
+
+    return strategy, mid_level_api, optimizer
+
+  def _create_sparse_data(self, include_weights, weight=0.5):
+    sparse_features = (sparse_tensor.SparseTensor(
+        indices=self.feature_watched_indices,
+        values=self.feature_watched_values,
+        dense_shape=[self.data_batch_size, 2]),
+                       sparse_tensor.SparseTensor(
+                           indices=self.feature_favorited_indices,
+                           values=self.feature_favorited_values,
+                           dense_shape=[self.data_batch_size, 2]),
+                       sparse_tensor.SparseTensor(
+                           indices=self.feature_friends_indices,
+                           values=self.feature_friends_values,
+                           dense_shape=[self.data_batch_size, 3]))
+    if include_weights:
+      weights = []
+      for sparse in sparse_features:
+        values = (
+            array_ops.ones_like(sparse.values, dtype=dtypes.float32) * weight)
+        weights.append(
+            sparse_tensor.SparseTensor(
+                indices=sparse.indices,
+                values=values,
+                dense_shape=sparse.dense_shape))
+      sparse_features = (sparse_features, tuple(weights))
+    return sparse_features
+
+  def _create_sparse_dataset(self, strategy, include_weights=False, weight=0.5):
+    # Create dataset for enqueue operation
+    sparse_features = self._create_sparse_data(include_weights, weight)
+
+    dataset = dataset_ops.DatasetV2.from_tensors(sparse_features)
+
+    # Data is batched to self.data_batch_size, rebatch to global batch size.
+    return dataset.unbatch().repeat().batch(
+        self.batch_size * strategy.num_replicas_in_sync, drop_remainder=True)
+
+  def _create_high_dimensional_sparse_dataset(self,
+                                              strategy,
+                                              include_weights=False,
+                                              weight=0.5):
+    sparse_features = (
+        sparse_tensor.SparseTensor(
+            indices=self.feature_watched_indices_high_dimensional,
+            values=self.feature_watched_values_high_dimensional,
+            dense_shape=[self.data_batch_size, self.data_batch_size, 2]),
+        sparse_tensor.SparseTensor(
+            indices=self.feature_favorited_indices_high_dimensional,
+            values=self.feature_favorited_values_high_dimensional,
+            dense_shape=[self.data_batch_size, self.data_batch_size, 2]),
+        sparse_tensor.SparseTensor(
+            indices=self.feature_friends_indices_high_dimensional,
+            values=self.feature_friends_values_high_dimensional,
+            dense_shape=[self.data_batch_size, self.data_batch_size, 3]))
+    if include_weights:
+      weights = []
+      for sparse in sparse_features:
+        values = (
+            array_ops.ones_like(sparse.values, dtype=dtypes.float32) * weight)
+        weights.append(
+            sparse_tensor.SparseTensor(
+                indices=sparse.indices,
+                values=values,
+                dense_shape=sparse.dense_shape))
+      sparse_features = (sparse_features, tuple(weights))
+
+    dataset = dataset_ops.DatasetV2.from_tensors(sparse_features)
+    # Data is batched to self.data_batch_size, rebatch to global batch size.
+    return dataset.unbatch().repeat().batch(
+        self.batch_size * strategy.num_replicas_in_sync, drop_remainder=True)
+
+  def _create_high_dimensional_ragged_dataset(self,
+                                              strategy,
+                                              include_weights=False,
+                                              weight=0.5):
+    ragged_features = (
+        ragged_tensor.RaggedTensor.from_row_lengths(
+            row_lengths=self.feature_watched_row_lengths_high_dimensional,
+            values=self.feature_watched_values_high_dimensional),
+        ragged_tensor.RaggedTensor.from_row_lengths(
+            row_lengths=self.feature_favorited_row_lengths_high_dimensional,
+            values=self.feature_favorited_values_high_dimensional),
+        ragged_tensor.RaggedTensor.from_row_lengths(
+            row_lengths=self.feature_friends_row_lengths_high_dimensional,
+            values=self.feature_friends_values_high_dimensional))
+    if include_weights:
+      weights = []
+      for ragged in ragged_features:
+        values = (
+            array_ops.ones_like(ragged.values, dtype=dtypes.float32) * weight)
+        weights.append(
+            ragged_tensor.RaggedTensor(
+                row_lengths=ragged.row_lengths(), values=values))
+      ragged_features = (ragged_features, tuple(weights))
+
+    dataset = dataset_ops.DatasetV2.from_tensors(ragged_features)
+    # Data is batched to self.data_batch_size, rebatch to global batch size.
+    return dataset.unbatch().repeat().batch(
+        self.batch_size * strategy.num_replicas_in_sync, drop_remainder=True)
+
+  def _create_ragged_dataset(self, strategy, include_weights=False, weight=0.5):
+    # Create dataset for enqueue operation
+    sparse_features = self._create_sparse_data(include_weights, weight)
+    ragged_features = nest.map_structure(ragged_tensor.RaggedTensor.from_sparse,
+                                         sparse_features)
+
+    dataset = dataset_ops.DatasetV2.from_tensors(ragged_features)
+
+    # Data is batched to self.data_batch_size, rebatch to global batch size.
+    return dataset.unbatch().repeat().batch(
+        self.batch_size * strategy.num_replicas_in_sync, drop_remainder=True)
+
+  def _create_dense_dataset(self, strategy, include_weights=False, weight=0.5):
+
+    features = (constant_op.constant(
+        self.feature_watched_values[:self.data_batch_size], dtype=dtypes.int32),
+                constant_op.constant(
+                    self.feature_favorited_values[:self.data_batch_size],
+                    dtype=dtypes.int32),
+                constant_op.constant(
+                    self.feature_friends_values[:self.data_batch_size],
+                    dtype=dtypes.int32))
+    if include_weights:
+      weights = [
+          array_ops.ones_like(t, dtype=dtypes.float32) * weight
+          for t in features
+      ]
+      features = (features, tuple(weights))
+
+    dataset = dataset_ops.DatasetV2.from_tensors(features)
+    return dataset.unbatch().repeat().batch(
+        self.batch_size * strategy.num_replicas_in_sync, drop_remainder=True)
+
+  def _create_high_dimensional_dense_dataset(self,
+                                             strategy,
+                                             include_weights=False,
+                                             weight=0.5):
+
+    dense_size = self.data_batch_size * self.data_batch_size
+    features = (constant_op.constant(
+        self.feature_watched_values_high_dimensional[:dense_size],
+        shape=(self.data_batch_size, self.data_batch_size, 1),
+        dtype=dtypes.int32),
+                constant_op.constant(
+                    self.feature_favorited_values_high_dimensional[:dense_size],
+                    shape=(self.data_batch_size, self.data_batch_size, 1),
+                    dtype=dtypes.int32),
+                constant_op.constant(
+                    self.feature_friends_values_high_dimensional[:dense_size],
+                    shape=(self.data_batch_size, self.data_batch_size, 1),
+                    dtype=dtypes.int32))
+    if include_weights:
+      weights = [
+          array_ops.ones_like(t, dtype=dtypes.float32) * weight
+          for t in features
+      ]
+      features = (features, tuple(weights))
+    dataset = dataset_ops.DatasetV2.from_tensors(features)
+    return dataset.unbatch().repeat().batch(
+        self.batch_size * strategy.num_replicas_in_sync, drop_remainder=True)
+
+  def _check_results(self, strategy, shard_out_val, training, input_data,
+                     table_to_variable, optimizer, is_high_dimensional):
+    num_replicas = strategy.num_replicas_in_sync
+
+    # Unpack the values `strategy.run()` returns.
+    loss = self._unpack(strategy, shard_out_val[0])
+    activation_watched = self._unpack(strategy, shard_out_val[1])
+    activation_favorited = self._unpack(strategy, shard_out_val[2])
+    activation_friends = self._unpack(strategy, shard_out_val[3])
+
+    # Core 0:
+    # Calculate the values of embedding activations.
+    activation_watched_gold0 = np.array([[0, 1, 2, 3], [4, 6, 8, 10]])
+    activation_favorited_gold0 = np.array([[4, 6, 8, 10], [4, 5, 6, 7]])
+    # Second row of `activation_friends_gold0` is the mean of the following.
+    # row 0: 0 1
+    # row 1: 2 3
+    # row 2: 4 5
+    activation_friends_gold0 = np.array([[6, 7], [2, 3]])
+
+    loss_gold0 = self._compute_loss(activation_watched_gold0,
+                                    activation_favorited_gold0,
+                                    activation_friends_gold0)
+
+    # Add on values from other cores:
+    # Activations for watched are an alternating sequence of
+    # activation_watched_gold0 and activation_favorited_gold0.
+    # For favorited it is the same but in the opposite order.
+    activation_watched_gold = np.concatenate(
+        (activation_watched_gold0, activation_favorited_gold0))
+    activation_favorited_gold = np.concatenate(
+        (activation_favorited_gold0, activation_watched_gold0))
+    activation_friends_gold = np.concatenate(
+        (activation_friends_gold0, activation_friends_gold0))
+
+    if is_high_dimensional:
+      activation_watched_gold = np.stack([activation_watched_gold] *
+                                         self.batch_size * num_replicas)
+
+      activation_favorited_gold = np.stack([activation_favorited_gold] *
+                                           self.batch_size * num_replicas)
+
+      activation_friends_gold = np.stack([activation_friends_gold] *
+                                         self.batch_size * num_replicas)
+    else:
+      if num_replicas == 1:
+        activation_watched_gold = activation_watched_gold0
+        activation_favorited_gold = activation_favorited_gold0
+        activation_friends_gold = activation_friends_gold0
+      else:
+        activation_watched_gold = np.concatenate(
+            [activation_watched_gold] * (num_replicas // self.batch_size))
+        activation_favorited_gold = np.concatenate(
+            [activation_favorited_gold] * (num_replicas // self.batch_size))
+        activation_friends_gold = np.concatenate(
+            [activation_friends_gold] * (num_replicas // self.batch_size))
+
+    loss_gold = [loss_gold0] * num_replicas
+
+    # Test values.
+    self.assertAllClose(activation_watched_gold, activation_watched)
+    self.assertAllClose(activation_favorited_gold, activation_favorited)
+    self.assertAllClose(activation_friends_gold, activation_friends)
+
+    self.assertAllClose(loss_gold, loss)
+
+    embedding_table_video_before = np.copy(
+        np.reshape(self.embedding_values, [8, 4]))
+    embedding_table_user_before = np.copy(
+        np.reshape(self.embedding_values, [16, 2]))
+    if is_high_dimensional:
+      global_batch_size = self.batch_size * self.data_batch_size * num_replicas
+    else:
+      global_batch_size = self.batch_size * num_replicas
+    if training:
+      gradient_wrt_watched_gold = (2 * activation_watched_gold /
+                                   global_batch_size)
+      gradient_wrt_favorited_gold = (2 * activation_favorited_gold /
+                                     global_batch_size)
+      gradient_wrt_friends_gold = (2 * activation_friends_gold /
+                                   global_batch_size)
+
+      # Calculate gradients wrt embedding tables.
+      gradients_wrt_user = (
+          self._compute_gradients_wrt_embedding_table(
+              gradient_wrt_friends_gold, embedding_table_user_before,
+              input_data[2].indices.numpy(), input_data[2].values.numpy(),
+              self.table_user.combiner))
+      gradients_wrt_video = (
+          self._compute_gradients_wrt_embedding_table(
+              gradient_wrt_favorited_gold, embedding_table_video_before,
+              input_data[1].indices.numpy(), input_data[1].values.numpy(),
+              self.table_video.combiner) +
+          self._compute_gradients_wrt_embedding_table(
+              gradient_wrt_watched_gold, embedding_table_video_before,
+              input_data[0].indices.numpy(), input_data[0].values.numpy(),
+              self.table_video.combiner))
+
+      self._check_embedding_and_slot_variables(embedding_table_user_before,
+                                               gradients_wrt_user,
+                                               embedding_table_video_before,
+                                               gradients_wrt_video, optimizer,
+                                               table_to_variable)
+
+  def _check_embedding_and_slot_variables(self, embedding_table_user_before,
+                                          gradients_wrt_user,
+                                          embedding_table_video_before,
+                                          gradients_wrt_video, optimizer,
+                                          table_to_variable):
+    if isinstance(optimizer, tpu_embedding_v2_utils.SGD):
+      check_fn = self._check_embedding_and_slot_variables_for_sgd
+    elif isinstance(optimizer, tpu_embedding_v2_utils.Adagrad):
+      check_fn = self._check_embedding_and_slot_variables_for_adagrad
+    elif isinstance(optimizer, tpu_embedding_v2_utils.AdagradMomentum):
+      check_fn = self._check_embedding_and_slot_variables_for_adagrad_momentum
+    elif isinstance(optimizer, tpu_embedding_v2_utils.Adam):
+      check_fn = self._check_embedding_and_slot_variables_for_adam
+    elif isinstance(optimizer, tpu_embedding_v2_utils.FTRL):
+      check_fn = self._check_embedding_and_slot_variables_for_ftrl
+    else:
+      raise ValueError('optimizer is not recognized: ', type(optimizer))
+    check_fn(embedding_table_user_before, gradients_wrt_user, optimizer,
+             table_to_variable[self.table_user.name])
+    check_fn(embedding_table_video_before, gradients_wrt_video, optimizer,
+             table_to_variable[self.table_video.name])
+
+  def _check_embedding_and_slot_variables_for_sgd(self, embedding_table_before,
+                                                  gradients, optimizer,
+                                                  variables):
+    embedding_table = np.copy(embedding_table_before)
+    embedding_table -= optimizer.learning_rate * np.sum(gradients, axis=0)
+    self.assertAllClose(
+        self._get_variable(variables['parameters']).numpy(), embedding_table)
+
+  def _check_embedding_and_slot_variables_for_adagrad(self,
+                                                      embedding_table_before,
+                                                      gradients, optimizer,
+                                                      variable):
+    embedding_table = np.copy(embedding_table_before)
+    accumulator = (
+        optimizer.initial_accumulator_value + np.sum(gradients, axis=0)**2)
+    embedding_table -= (
+        optimizer.learning_rate * np.sum(gradients, axis=0) /
+        np.sqrt(accumulator))
+    self.assertAllClose(
+        self._get_variable(variable['parameters']).numpy(), embedding_table)
+    self.assertAllClose(
+        self._get_variable(variable['accumulators']).numpy(), accumulator)
+
+  def _check_embedding_and_slot_variables_for_adagrad_momentum(
+      self, embedding_table_before, gradients, optimizer, variable):
+    embedding_table = np.copy(embedding_table_before)
+    accumulator = np.zeros(self._get_variable(variable['accumulators']).shape)
+    momenta = np.zeros(self._get_variable(variable['momenta']).shape)
+    gradients = np.sum(gradients, axis=0)
+    if optimizer.beta2 == 1.0:
+      accumulator += gradients**2
+    else:
+      accumulator = optimizer.beta2 * accumulator + (
+          1 - optimizer.beta2) * gradients**2
+    accumulator_power = np.power(accumulator + optimizer.epsilon,
+                                 -1.0 / optimizer.exponent)
+    momenta = optimizer.momentum * momenta + gradients * accumulator_power
+    if optimizer.use_nesterov:
+      update = optimizer.momentum * momenta + gradients * accumulator_power
+    else:
+      update = momenta
+    embedding_table -= optimizer.learning_rate * update
+    self.assertAllClose(
+        self._get_variable(variable['parameters']).numpy(),
+        embedding_table,
+        rtol=1e-3)
+    self.assertAllClose(
+        self._get_variable(variable['accumulators']).numpy(),
+        accumulator,
+        rtol=1e-3)
+    self.assertAllClose(
+        self._get_variable(variable['momenta']).numpy(), momenta, rtol=1e-3)
+
+  def _check_embedding_and_slot_variables_for_adam(self, embedding_table_before,
+                                                   gradients, optimizer,
+                                                   variable):
+    embedding_table = np.copy(embedding_table_before)
+    g = np.sum(gradients, axis=0)
+    v = g**2 * (1 - optimizer.beta_2)
+    m = g * (1 - optimizer.beta_1)
+    epsilon = optimizer.epsilon
+    # TPU Embeddings don't have the LR decay factor for Adam.
+    lr_modifier = 1
+    embedding_table -= (
+        m * optimizer.learning_rate * lr_modifier / (np.sqrt(v) + epsilon))
+    self.assertAllClose(
+        self._get_variable(variable['parameters']).numpy(),
+        embedding_table,
+        rtol=1e-4)
+    self.assertAllClose(
+        self._get_variable(variable['momenta']).numpy(), m, rtol=1e-4)
+    self.assertAllClose(
+        self._get_variable(variable['velocities']).numpy(), v, rtol=1e-4)
+
+  def _check_embedding_and_slot_variables_for_ftrl(self, embedding_table_before,
+                                                   gradients, optimizer,
+                                                   variable):
+    embedding_table = np.copy(embedding_table_before)
+    neg_lr_p = -optimizer.learning_rate_power
+    accumulator = (
+        optimizer.initial_accumulator_value + np.sum(gradients, axis=0)**2)
+    sigma = (accumulator**neg_lr_p - optimizer.initial_accumulator_value**
+             neg_lr_p) / optimizer.learning_rate
+    linear = np.sum(gradients, axis=0) - sigma * embedding_table
+    quadratic = accumulator**neg_lr_p / optimizer.learning_rate
+    embedding_table = -linear / quadratic
+    actual_parameters = self._get_variable(variable['parameters']).numpy()
+    # For entries where `linear` == 0, it is not worth comparing since the
+    # initial values have not been touched yet and they will not agree with what
+    # the actual values should be.
+    actual_parameters *= (linear != 0.0)
+    # FTRL has a bit more precision diff on parameters.
+    self.assertAllClose(actual_parameters, embedding_table, rtol=5e-5)
+    self.assertAllClose(
+        self._get_variable(variable['linears']).numpy(), linear, rtol=5e-4)
+    self.assertAllClose(
+        self._get_variable(variable['accumulators']).numpy(), accumulator)
+
+  def _get_replica_numpy(self, structured, strategy, replica_id):
+
+    def select_replica(x):
+      x = strategy.experimental_local_results(x)
+      if len(x) == 1:
+        return x.numpy()
+      return x[replica_id].numpy()
+
+    return nest.map_structure(select_replica, structured)
+
+  def _compute_gradients_wrt_embedding_table(self, gradient_wrt_activation,
+                                             embedding_table, feature_indices,
+                                             feature_values, combiner):
+    """Compute gradients wrt embedding_table.
+
+    Args:
+      gradient_wrt_activation: `np.array` with shape `batch_size` by embedding
+        `dimension`.
+      embedding_table: `np.array` with shape `vocabulary_size` by embedding
+        `dimension`.
+      feature_indices: `indices` as used to construct `SparseTensor`.
+      feature_values: `values` as used to construct `SparseTensor`.
+      combiner: `String`, 'mean' or 'sum'.
+
+    Returns:
+      Gradients wrt `embedding_table`, an `np.array`s with shape
+        `batch_size` by `vocabulary_size` by
+        embedding `dimension`.
+
+    Raises:
+      ValueError: if `combiner` is not one of 'mean' or 'sum'.
+    """
+    if combiner not in ('mean', 'sum'):
+      raise ValueError(
+          '`combiner` must be mean or sum; got {}.'.format(combiner))
+    grads_shape = gradient_wrt_activation.shape[:-1] + embedding_table.shape
+    grads = np.zeros(shape=grads_shape)
+    count = np.zeros(shape=grads_shape)
+    for feature_indice, vocabulary_id in zip(feature_indices, feature_values):
+      batch_index = tuple(feature_indice[:-1])
+      grads[batch_index][vocabulary_id] += gradient_wrt_activation[batch_index]
+      count[batch_index] += 1
+    count[count == 0] = 1
+    if combiner == 'mean':
+      grads = grads / count
+    return np.reshape(grads, (-1, *embedding_table.shape))
+
+  def _unpack(self, strategy, per_replica_output):
+    per_replica_output = strategy.experimental_local_results(per_replica_output)
+    per_replica_output = array_ops.concat(per_replica_output, axis=0).numpy()
+    return per_replica_output
+
+  def _get_total_loss_tensor(self, activations):
+    losses = []
+    for activation in activations:
+      losses.append(
+          math_ops.reduce_mean(
+              math_ops.reduce_sum(
+                  gen_math_ops.squared_difference(activation, 0), axis=-1)))
+    total_loss = array_ops.expand_dims_v2(sum(losses), 0)
+    return total_loss
+
+  def _compute_loss(self, activation_watched, activation_favorited,
+                    activation_friends):
+    watched_loss = np.mean(np.sum(activation_watched**2, axis=-1))
+    favorited_loss = np.mean(np.sum(activation_favorited**2, axis=-1))
+    friends_loss = np.mean(np.sum(activation_friends**2, axis=-1))
+    loss = watched_loss + favorited_loss + friends_loss
+    return loss
+
+  def _get_variable(self, variable):
+    if isinstance(variable, tpu_embedding_v2.TPUEmbeddingVariable):
+      return variable.variables[0]
+    return variable
+
+  def _get_tmpdir(self, name, subdir=''):
+    segments = [FLAGS.model_dir, name] + ([subdir] if subdir else [])
+    return os.path.join(*segments)

tensorflow/python/tpu/tests/tpu_embedding_v1_checkpoint_test.py

@@ -0,0 +1,214 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPUEmbeddingV0 mid level API on TPU."""
+
+from absl.testing import parameterized
+import numpy as np
+from tensorflow.python.checkpoint import checkpoint as util
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import save
+from tensorflow.python.tpu import tpu_embedding_for_serving
+from tensorflow.python.tpu import tpu_embedding_v1
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+from tensorflow.python.training import checkpoint_utils
+
+
+class TPUEmbeddingCheckpointTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  def _get_strategy(self):
+    # We can cache the strategy as TPUEmbeddingV0 doesn't require
+    # reconfiguration to the tpu.
+    if hasattr(self, 'strategy'):
+      return self.strategy
+    return super(TPUEmbeddingCheckpointTest, self)._get_strategy()
+
+  def _create_mid_level(self, optimizer=None):
+    # Create `TPUEmbedding` object.
+    if optimizer is None:
+      optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+
+    return tpu_embedding_v1.TPUEmbeddingV0(
+        feature_config=self.feature_config, optimizer=optimizer)
+
+  def test_checkpoint_save_and_restore(self):
+    strategy = self._get_strategy()
+    with strategy.scope():
+      first_mid_level_contents = np.ones((4, 4))
+      first_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(first_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=4,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+
+      first_mid_level = tpu_embedding_v1.TPUEmbeddingV0(
+          feature_config, first_mid_level_optimizer)
+      first_mid_level.build()
+
+    first_checkpoint = util.Checkpoint(model=first_mid_level)
+    first_checkpoint.save(self._get_tmpdir('restore', 'save'))
+
+    with strategy.scope():
+      second_mid_level_contents = np.ones((4, 4)) * 2
+      second_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(second_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=4,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+      second_mid_level = tpu_embedding_v1.TPUEmbeddingV0(
+          feature_config, second_mid_level_optimizer)
+      second_mid_level.build()
+    # We restore the checkpoint of our first model into our second model.
+    second_checkpoint = util.Checkpoint(model=second_mid_level)
+    second_checkpoint.restore(self._get_tmpdir('restore', 'save-1'))
+
+    self.assertAllClose(
+        first_mid_level_contents,
+        second_mid_level._variables['table']['parameters'].numpy(),
+        msg='Second mid level api should have restored the first model values.')
+
+  def test_model_export_cpu(self):
+    strategy = self._get_strategy()
+
+    with strategy.scope():
+      first_mid_level_contents = np.ones((4, 4))
+      first_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(first_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=4,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+
+      first_mid_level = tpu_embedding_v1.TPUEmbeddingV0(
+          feature_config, first_mid_level_optimizer)
+
+      first_mid_level.build()
+
+    cpu_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+    cpu_mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config, cpu_mid_level_optimizer)
+
+    cpu_mid_level.build()
+
+    tpu_checkpoint = util.Checkpoint(model=first_mid_level)
+    tpu_checkpoint.save(self._get_tmpdir('export_cpu', 'save'))
+
+    # We restore the checkpoint of our tpu mid level onto our cpu mid level.
+    cpu_checkpoint = util.Checkpoint(model=cpu_mid_level)
+    cpu_checkpoint.restore(self._get_tmpdir('export_cpu', 'save-1'))
+
+    @def_function.function
+    def serve_tensors(features):
+      features = tpu_embedding_for_serving.cpu_embedding_lookup(
+          features, None, cpu_mid_level.embedding_tables,
+          cpu_mid_level._feature_config)
+      return features[0]
+
+    signatures = {
+        'serving_default':
+            serve_tensors.get_concrete_function((tensor_spec.TensorSpec(
+                shape=(2,), dtype=dtypes.int32, name='feature'),))
+    }
+    save.save(
+        cpu_mid_level,
+        export_dir=self._get_tmpdir('export_cpu', 'exported_model'),
+        signatures=signatures)
+
+    imported = load.load(self._get_tmpdir('export_cpu', 'exported_model'))
+    predict_fn = imported.signatures['serving_default']
+
+    input_feature_value = np.array([1, 0])
+    input_batch = (constant_op.constant(
+        input_feature_value, dtype=dtypes.int32),)
+    prediction = predict_fn(*input_batch)['output_0']
+    self.assertAllClose(prediction.numpy(),
+                        first_mid_level_contents[input_feature_value])
+
+  @parameterized.parameters(tpu_embedding_v2_utils.SGD,
+                            tpu_embedding_v2_utils.Adagrad,
+                            tpu_embedding_v2_utils.Adam,
+                            tpu_embedding_v2_utils.FTRL)
+  def test_check_checkpoint_variable_names_are_same_on_cpu_and_tpu(
+      self, optimizer):
+    # Reinitialize the TPU so that we can re-initialize the embeddings with the
+    # given optimizer.
+    if optimizer != tpu_embedding_v2_utils.SGD:
+      self.skip_if_oss()
+    strategy = self._get_strategy()
+
+    with strategy.scope():
+      first_mid_level_contents = np.ones((4, 4))
+      first_mid_level_optimizer = optimizer(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(first_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=4,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+
+      first_mid_level = tpu_embedding_v1.TPUEmbeddingV0(
+          feature_config, first_mid_level_optimizer)
+
+      first_mid_level.build()
+
+    cpu_mid_level_optimizer = optimizer(learning_rate=0.1)
+    cpu_mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config, cpu_mid_level_optimizer)
+    cpu_mid_level.build()
+
+    tpu_checkpoint = util.Checkpoint(model=first_mid_level)
+    tpu_checkpoint.save(self._get_tmpdir('save-tpu', 'save'))
+    tpu_variables = checkpoint_utils.list_variables(
+        self._get_tmpdir('save-tpu'))
+
+    cpu_checkpoint = util.Checkpoint(model=cpu_mid_level)
+    cpu_checkpoint.save(self._get_tmpdir('save-cpu', 'save'))
+    cpu_variables = checkpoint_utils.list_variables(
+        self._get_tmpdir('save-cpu'))
+
+    self.assertAllEqual(tpu_variables, cpu_variables)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v1_correctness_test.py

@@ -0,0 +1,118 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v1
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+
+
+class TPUEmbeddingV0CorrectnessTest(tpu_embedding_base_test.TPUEmbeddingBaseTest
+                                   ):
+
+  def _get_strategy(self):
+    if hasattr(self, 'strategy'):
+      return self.strategy
+    return super(TPUEmbeddingV0CorrectnessTest, self)._get_strategy()
+
+  def _create_mid_level(self, optimizer=None):
+    # Create `TPUEmbedding` object.
+    if optimizer is None:
+      optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+
+    return tpu_embedding_v1.TPUEmbeddingV0(
+        feature_config=self.feature_config, optimizer=optimizer)
+
+  def _create_strategy_and_mid_level(self, optimizer_name):
+    strategy = self._get_strategy()
+
+    with strategy.scope():
+      if optimizer_name == 'sgd':
+        embedding_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      elif optimizer_name == 'adagrad':
+        embedding_optimizer = tpu_embedding_v2_utils.Adagrad(
+            learning_rate=0.1)
+      elif optimizer_name == 'adam':
+        embedding_optimizer = tpu_embedding_v2_utils.Adam(
+            learning_rate=0.1)
+      elif optimizer_name == 'ftrl':
+        embedding_optimizer = tpu_embedding_v2_utils.FTRL(
+            learning_rate=0.1)
+      else:
+        raise ValueError('optimizer is not recognized: ', optimizer_name)
+
+      mid_level_api = self._create_mid_level(optimizer=embedding_optimizer)
+
+    return strategy, mid_level_api
+
+  @parameterized.parameters(True, False)
+  def test_enqueue_with_weights(self, ragged):
+    strategy, mid_level_api = self._create_strategy_and_mid_level('sgd')
+    weight = 0.5
+    if ragged:
+      dataset = self._create_ragged_dataset(
+          strategy, include_weights=True, weight=weight)
+    else:
+      dataset = self._create_sparse_dataset(
+          strategy, include_weights=True, weight=weight)
+
+    dataset_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def embedding_lookup(features, weights):
+
+      def step(features, weights):
+        return mid_level_api(features, weights)
+
+      return strategy.run(step, args=(features, weights))
+
+    features, weights = next(dataset_iter)
+    # Replace the weight for the second feature by None to test.
+    weights = (weights[0], None, weights[2])
+
+    no_weights_activations = embedding_lookup(features, weights=None)
+    weights_activations = embedding_lookup(features, weights=weights)
+
+    no_weights0 = (self._unpack(strategy, no_weights_activations[0]),
+                   self._unpack(strategy, no_weights_activations[1]),
+                   self._unpack(strategy, no_weights_activations[2]))
+    weights0 = (self._unpack(strategy, weights_activations[0]),
+                self._unpack(strategy, weights_activations[1]),
+                self._unpack(strategy, weights_activations[2]))
+    # videos table has sum combiner and users table has mean combiner.
+    # i.e. users table lookups isn't affected by the weights as all the weights
+    # are the same.
+    # Tuple entry 0 and 1 are the watched and favorited features from the videos
+    # table and entry 2 is the friends feature from the users table.
+    # Note that None was passed as a weight for entry 1 so weight should have no
+    # effect.
+    weight = (0.5, 1.0, 1.0)
+    golden = tuple([no_weight * w for no_weight, w in zip(no_weights0, weight)])
+
+    self.assertAllClose(golden, weights0)
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_checkpoint_test.py

@@ -0,0 +1,372 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+
+from absl.testing import parameterized
+import numpy as np
+from tensorflow.python.checkpoint import checkpoint as util
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.module import module
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import load
+from tensorflow.python.saved_model import save
+from tensorflow.python.tpu import tpu_embedding_for_serving
+from tensorflow.python.tpu import tpu_embedding_v2
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+from tensorflow.python.training import checkpoint_utils
+
+
+class TPUEmbeddingCheckpointTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  def make_checkpoint_and_get_embedding(self, name, model, num_rows):
+    """Saves model to checkpoint name, retrieves embedding variables."""
+    checkpoint = util.Checkpoint(model=model)
+    checkpoint.save(self._get_tmpdir(name, 'save'))
+
+    # Get the name of the table video variable which should be the only
+    # [8, 4] shaped tensor in the checkpoint. Note that we do this
+    # as the key can change.
+    variables = checkpoint_utils.list_variables(self._get_tmpdir(name))
+    variables = [name for name, size in variables if size == [num_rows, 4]]
+    if len(variables) != 1:
+      raise RuntimeError('Found {} copies of the parameter variable in the '
+                         'checkpoint. Exactly one copy exported.'.format(
+                             len(variables)))
+    return checkpoint_utils.load_variable(self._get_tmpdir(name), variables[0])
+
+  def test_checkpoint_save_retrieves(self):
+    strategy = self._get_strategy()
+    num_rows = strategy.num_replicas_in_sync
+
+    with strategy.scope():
+      first_mid_level_contents = np.ones((num_rows, 4))
+      first_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(first_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=num_rows,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+
+      first_mid_level = tpu_embedding_v2.TPUEmbedding(
+          feature_config, first_mid_level_optimizer)
+      first_mid_level.build(64)
+
+    # Ensure that the variables from the first model are loaded.
+    first_mid_level._load_variables()
+
+    self.assertAllClose(
+        first_mid_level_contents,
+        self.make_checkpoint_and_get_embedding('before_load', first_mid_level,
+                                               num_rows),
+        msg='Checkpoint should contain values from the first api object.')
+
+    # Reinitialize the tpu.
+    tpu_cluster_resolver.initialize_tpu_system(self.resolver)
+
+    with strategy.scope():
+      second_mid_level_contents = np.ones((num_rows, 4)) * 2
+      second_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(second_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=num_rows,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+      second_mid_level = tpu_embedding_v2.TPUEmbedding(
+          feature_config, second_mid_level_optimizer)
+      second_mid_level.build(64)
+
+    second_mid_level._load_variables()
+
+    # When we load the variables from the second mid level API object to the TPU
+    # we expect that checkpointing the first mid level API object will now
+    # retrieve the values from the TPU which are now different from the current
+    # variables in the first mid level.
+    self.assertAllClose(
+        second_mid_level_contents,
+        self.make_checkpoint_and_get_embedding('after_load', first_mid_level,
+                                               num_rows),
+        msg='Checkpoint should contain values from the second api object.')
+
+  def test_checkpoint_restore_loads(self):
+    strategy = self._get_strategy()
+    num_rows = strategy.num_replicas_in_sync
+
+    def get_values(mid):
+      return ops.convert_to_tensor(
+          mid._variables['table']['parameters'].variables[0])
+
+    with strategy.scope():
+      first_mid_level_contents = np.ones((num_rows, 4))
+      first_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(first_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=num_rows,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+
+      first_mid_level = tpu_embedding_v2.TPUEmbedding(
+          feature_config, first_mid_level_optimizer)
+      first_mid_level.build(64)
+
+    first_mid_level._load_variables()
+
+    first_checkpoint = util.Checkpoint(model=first_mid_level)
+    first_checkpoint.save(self._get_tmpdir('restore', 'save'))
+
+    tpu_cluster_resolver.initialize_tpu_system(self.resolver)
+
+    with strategy.scope():
+      second_mid_level_contents = np.ones((num_rows, 4)) * 2
+      second_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(second_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=num_rows,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+      second_mid_level = tpu_embedding_v2.TPUEmbedding(
+          feature_config, second_mid_level_optimizer)
+      second_mid_level.build(64)
+
+    second_mid_level._load_variables()
+
+    self.assertAllClose(
+        second_mid_level_contents,
+        get_values(second_mid_level),
+        msg='Second mid level api should contain its initial values.',
+    )
+    # We restore the checkpoint of our first model into our second model.
+    # This should load the first mid level API object onto the TPU.
+    second_checkpoint = util.Checkpoint(model=second_mid_level)
+    second_checkpoint.restore(self._get_tmpdir('restore', 'save-1'))
+
+    # Call retrieve here as a way to check what the TPU contains.
+    # Calling the retrieve ops directly might make for a cleaner separation of
+    # test and module, though.
+    second_mid_level._retrieve_variables()
+
+    self.assertAllClose(
+        first_mid_level_contents,
+        get_values(second_mid_level),
+        msg='Second mid level api should have retrieved the first model values.'
+    )
+
+  def test_checkpoint_restore_before_variable_creation(self):
+    self.skip_if_oss()
+
+    class TestModule(module.Module):
+
+      def __init__(self, initializer, rows):
+        self._initializer = initializer
+        self._rows = rows
+
+        table = tpu_embedding_v2_utils.TableConfig(
+            vocabulary_size=self._rows,
+            dim=4,
+            initializer=self._initializer,
+            combiner='sum',
+            name='table')
+        feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+            table=table, name='feature'),)
+        optimizer = tpu_embedding_v2_utils.SGD()
+
+        self.tpu_embedding = tpu_embedding_v2.TPUEmbedding(
+            feature_config, optimizer)
+
+      def create_embedding(self):
+        # We aren't training so batch_size here doesn't matter.
+        self.tpu_embedding.build(64)
+
+    strategy = self._get_strategy()
+    with strategy.scope():
+      module1 = TestModule(init_ops_v2.Ones(),
+                           strategy.num_replicas_in_sync * 2)
+      module1.create_embedding()
+
+    checkpoint = util.Checkpoint(test_module=module1)
+    checkpoint.save(self._get_tmpdir('restore_before_create', 'save'))
+
+    # Reinitialize the tpu
+    strategy = self._get_strategy()
+
+    with strategy.scope():
+      module2 = TestModule(init_ops_v2.Zeros(),
+                           strategy.num_replicas_in_sync * 2)
+
+    checkpoint = util.Checkpoint(test_module=module2)
+    checkpoint.restore(self._get_tmpdir('restore_before_create', 'save-1'))
+
+    with strategy.scope():
+      module2.create_embedding()
+
+    def get_values(mid):
+      return mid._variables['table']['parameters'].variables[0].numpy()
+
+    self.assertAllClose(
+        np.ones((strategy.num_replicas_in_sync * 2, 4)),
+        get_values(module2.tpu_embedding))
+
+    # Fetch the values from the TPU to check that they are the same.
+    module2.tpu_embedding._retrieve_variables()
+
+    self.assertAllClose(
+        np.ones((strategy.num_replicas_in_sync * 2, 4)),
+        get_values(module2.tpu_embedding))
+
+  def test_model_export_cpu(self):
+    strategy = self._get_strategy()
+    num_rows = strategy.num_replicas_in_sync
+
+    with strategy.scope():
+      first_mid_level_contents = np.ones((num_rows, 4))
+      first_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(first_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=num_rows,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+
+      first_mid_level = tpu_embedding_v2.TPUEmbedding(
+          feature_config, first_mid_level_optimizer)
+
+      first_mid_level.build(64)
+
+    cpu_mid_level_optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+    cpu_mid_level = tpu_embedding_v2.TPUEmbedding(feature_config,
+                                                  cpu_mid_level_optimizer)
+
+    cpu_mid_level.build(64)
+
+    first_mid_level._load_variables()
+
+    tpu_checkpoint = util.Checkpoint(model=first_mid_level)
+    tpu_checkpoint.save(self._get_tmpdir('export_cpu', 'save'))
+
+    # We restore the checkpoint of our tpu mid level onto our cpu mid level.
+    cpu_checkpoint = util.Checkpoint(model=cpu_mid_level)
+    cpu_checkpoint.restore(self._get_tmpdir('export_cpu', 'save-1'))
+
+    @def_function.function
+    def serve_tensors(features):
+      features = tpu_embedding_for_serving.cpu_embedding_lookup(
+          features, None, cpu_mid_level.embedding_tables,
+          cpu_mid_level._feature_config)
+      return features[0]
+
+    signatures = {
+        'serving_default':
+            serve_tensors.get_concrete_function((tensor_spec.TensorSpec(
+                shape=(2,), dtype=dtypes.int32, name='feature'),))
+    }
+    save.save(
+        cpu_mid_level,
+        export_dir=self._get_tmpdir('export_cpu', 'exported_model'),
+        signatures=signatures)
+
+    imported = load.load(self._get_tmpdir('export_cpu', 'exported_model'))
+    predict_fn = imported.signatures['serving_default']
+
+    input_feature_value = np.array([1, 0])
+    input_batch = (constant_op.constant(
+        input_feature_value, dtype=dtypes.int32),)
+    prediction = predict_fn(*input_batch)['output_0']
+    self.assertAllClose(prediction.numpy(),
+                        first_mid_level_contents[input_feature_value])
+
+  @parameterized.parameters(tpu_embedding_v2_utils.SGD,
+                            tpu_embedding_v2_utils.Adagrad,
+                            tpu_embedding_v2_utils.Adam,
+                            tpu_embedding_v2_utils.FTRL)
+  def test_check_checkpoint_variable_names_are_same_on_cpu_and_tpu(
+      self, optimizer):
+    # Reinitialize the TPU so that we can re-initialize the embeddings with the
+    # given optimizer.
+    if optimizer != tpu_embedding_v2_utils.SGD:
+      self.skip_if_oss()
+    strategy = self._get_strategy()
+    num_rows = strategy.num_replicas_in_sync
+
+    with strategy.scope():
+      first_mid_level_contents = np.ones((num_rows, 4))
+      first_mid_level_optimizer = optimizer(learning_rate=0.1)
+      initializer = init_ops_v2.Constant(first_mid_level_contents)
+
+      table = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=num_rows,
+          dim=4,
+          initializer=initializer,
+          combiner='sum',
+          name='table')
+      feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+          table=table, name='feature'),)
+
+      first_mid_level = tpu_embedding_v2.TPUEmbedding(
+          feature_config, first_mid_level_optimizer)
+
+      first_mid_level.build(64)
+
+    cpu_mid_level_optimizer = optimizer(learning_rate=0.1)
+    cpu_mid_level = tpu_embedding_v2.TPUEmbedding(feature_config,
+                                                  cpu_mid_level_optimizer)
+    cpu_mid_level.build(64)
+
+    tpu_checkpoint = util.Checkpoint(model=first_mid_level)
+    tpu_checkpoint.save(self._get_tmpdir('save-tpu', 'save'))
+    tpu_variables = checkpoint_utils.list_variables(
+        self._get_tmpdir('save-tpu'))
+
+    cpu_checkpoint = util.Checkpoint(model=cpu_mid_level)
+    cpu_checkpoint.save(self._get_tmpdir('save-cpu', 'save'))
+    cpu_variables = checkpoint_utils.list_variables(
+        self._get_tmpdir('save-cpu'))
+
+    self.assertAllEqual(tpu_variables, cpu_variables)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_base_test.py

@@ -0,0 +1,108 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework.tensor_shape import TensorShape
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+
+
+class TPUEmbeddingCorrectnessBaseTest(
+    tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  def _test_embedding(self, optimizer_name, training, sparse,
+                      is_high_dimensional):
+    strategy, mid_level_api, optimizer = (
+        self._create_strategy_and_mid_level(optimizer_name))
+
+    if sparse:
+      if is_high_dimensional:
+        dataset = self._create_high_dimensional_sparse_dataset(strategy)
+      else:
+        dataset = self._create_sparse_dataset(strategy)
+    else:
+      if is_high_dimensional:
+        dataset = self._create_high_dimensional_sparse_dataset(strategy)
+      else:
+        dataset = self._create_ragged_dataset(strategy)
+
+    if is_high_dimensional:
+      if sparse:
+        mid_level_api.build([
+            TensorShape([self.batch_size, self.data_batch_size, 2]),
+            TensorShape([self.batch_size, self.data_batch_size, 2]),
+            TensorShape([self.batch_size, self.data_batch_size, 3]),
+        ])
+      else:
+        mid_level_api.build([
+            TensorShape([self.batch_size, self.data_batch_size, None]),
+            TensorShape([self.batch_size, self.data_batch_size, None]),
+            TensorShape([self.batch_size, self.data_batch_size, None]),
+        ])
+
+    dist = strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(experimental_fetch_to_device=False))
+    dist_iter = iter(dist)
+
+    @def_function.function
+    def test_fn():
+
+      def step():
+        """Create and run computation that returns the embedding activations."""
+        if not training:
+          activations = mid_level_api.dequeue()
+          total_loss = self._get_total_loss_tensor(activations)
+          ret_val = [total_loss] + list(activations)
+          return ret_val
+        else:
+          with backprop.GradientTape() as tape:
+            activations = mid_level_api.dequeue()
+            tape.watch(activations)
+            total_loss = self._get_total_loss_tensor(activations)
+            loss_per_replica = total_loss / strategy.num_replicas_in_sync
+          gradients = tape.gradient(loss_per_replica, activations)
+          mid_level_api.apply_gradients(gradients)
+        ret_val = [total_loss] + list(activations)
+        return ret_val
+
+      mid_level_api.enqueue(next(dist_iter), training=training)
+      result = strategy.run(step)
+      return result
+
+    # Run model.
+    shard_out_val = test_fn()
+
+    # Retrieve TPU weights to CPU.
+    mid_level_api._retrieve_variables()
+
+    # Compute sparse tensors for global batch.
+    if is_high_dimensional:
+      input_data = next(
+          iter(self._create_high_dimensional_sparse_dataset(strategy)))
+    else:
+      input_data = next(iter(self._create_sparse_dataset(strategy)))
+
+    # Check results.
+    self._check_results(strategy, shard_out_val, training, input_data,
+                        mid_level_api._variables, optimizer,
+                        is_high_dimensional)
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_dense_lookup_test.py

@@ -0,0 +1,91 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters([True, False])
+  def test_dense_lookup(self, is_high_dimensional):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    if is_high_dimensional:
+      dataset = self._create_high_dimensional_dense_dataset(strategy)
+    else:
+      dataset = self._create_dense_dataset(strategy)
+    dist = strategy.experimental_distribute_dataset(
+        dataset,
+        options=distribute_lib.InputOptions(experimental_fetch_to_device=False))
+    dist_iter = iter(dist)
+
+    @def_function.function
+    def test_fn():
+      def step():
+        return mid_level_api.dequeue()
+
+      mid_level_api.enqueue(next(dist_iter), training=False)
+      return strategy.run(step)
+
+    # Run model.
+    shard_out_val = test_fn()
+
+    shard0 = (self._unpack(strategy, shard_out_val[0]),
+              self._unpack(strategy, shard_out_val[1]),
+              self._unpack(strategy, shard_out_val[2]))
+
+    # embedding_values is a linear list, so we reshape to match the correct
+    # shape of the corresponding table before performing the lookup.
+    numpy_videos = np.reshape(self.embedding_values, (8, 4))
+    numpy_users = np.reshape(self.embedding_values, (16, 2))
+
+    repeat_batch_num = strategy.num_replicas_in_sync // 2
+
+    golden = (
+        (numpy_videos[self.feature_watched_values[:self.data_batch_size] *
+                      repeat_batch_num],
+         numpy_videos[self.feature_favorited_values[:self.data_batch_size] *
+                      repeat_batch_num],
+         numpy_users[self.feature_friends_values[:self.data_batch_size] *
+                     repeat_batch_num]))
+    if is_high_dimensional:
+      dense_size = self.data_batch_size * self.data_batch_size
+      golden = ((
+          numpy_videos[self.feature_watched_values_high_dimensional[:dense_size]
+                       * repeat_batch_num].reshape(
+                           self.data_batch_size * repeat_batch_num,
+                           self.data_batch_size, -1),
+          numpy_videos[
+              self.feature_favorited_values_high_dimensional[:dense_size] *
+              repeat_batch_num].reshape(self.data_batch_size * repeat_batch_num,
+                                        self.data_batch_size, -1),
+          numpy_users[self.feature_friends_values_high_dimensional[:dense_size]
+                      * repeat_batch_num].reshape(
+                          self.data_batch_size * repeat_batch_num,
+                          self.data_batch_size, -1)))
+    self.assertAllClose(shard0, golden)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_hd_ragged_forward_test.py

@@ -0,0 +1,37 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters(
+      ['sgd', 'adagrad', 'adam', 'ftrl', 'adagrad_momentum'])
+  def test_embedding(self, optimizer_name):
+    if optimizer_name != 'sgd':
+      self.skip_if_oss()
+    self._test_embedding(
+        optimizer_name, training=False, sparse=False, is_high_dimensional=True)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_hd_ragged_training_test.py

@@ -0,0 +1,37 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters(
+      ['sgd', 'adagrad', 'adam', 'ftrl', 'adagrad_momentum'])
+  def test_embedding(self, optimizer_name):
+    if optimizer_name != 'sgd':
+      self.skip_if_oss()
+    self._test_embedding(
+        optimizer_name, training=True, sparse=False, is_high_dimensional=True)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_hd_sparse_forward_test.py

@@ -0,0 +1,37 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters(
+      ['sgd', 'adagrad', 'adam', 'ftrl', 'adagrad_momentum'])
+  def test_embedding(self, optimizer_name):
+    if optimizer_name != 'sgd':
+      self.skip_if_oss()
+    self._test_embedding(
+        optimizer_name, training=False, sparse=True, is_high_dimensional=True)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_hd_sparse_training_test.py

@@ -0,0 +1,37 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters(
+      ['sgd', 'adagrad', 'adam', 'ftrl', 'adagrad_momentum'])
+  def test_embedding(self, optimizer_name):
+    if optimizer_name != 'sgd':
+      self.skip_if_oss()
+    self._test_embedding(
+        optimizer_name, training=True, sparse=True, is_high_dimensional=True)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_ragged_forward_test.py

@@ -0,0 +1,37 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters(
+      ['sgd', 'adagrad', 'adam', 'ftrl', 'adagrad_momentum'])
+  def test_embedding(self, optimizer_name):
+    if optimizer_name != 'sgd':
+      self.skip_if_oss()
+    self._test_embedding(
+        optimizer_name, training=False, sparse=False, is_high_dimensional=False)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_ragged_training_test.py

@@ -0,0 +1,37 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters(
+      ['sgd', 'adagrad', 'adam', 'ftrl', 'adagrad_momentum'])
+  def test_embedding(self, optimizer_name):
+    if optimizer_name != 'sgd':
+      self.skip_if_oss()
+    self._test_embedding(
+        optimizer_name, training=True, sparse=False, is_high_dimensional=False)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_sequence_feature_test.py

@@ -0,0 +1,127 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v2
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+from tensorflow.python.util import nest
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters([True, False])
+  def test_sequence_embeddings(self, sparse):
+    feature_config = (
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=self.table_video, name='watched',
+            max_sequence_length=2),
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=self.table_video, name='favorited',
+            max_sequence_length=2),
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=self.table_user, name='friends',
+            max_sequence_length=3))
+    optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+    strategy = self._get_strategy()
+    num_replicas = strategy.num_replicas_in_sync
+    with strategy.scope():
+      mid_level = tpu_embedding_v2.TPUEmbedding(
+          feature_config=feature_config,
+          optimizer=optimizer)
+    # Call build here. We call 'next' outside of the tf.function and this
+    # results in data where the shape of the sparse tensor is a tensor which we
+    # can't tell the shape of at tracing time.
+    mid_level.build(self.batch_size)
+    if sparse:
+      dataset = self._create_sparse_dataset(strategy)
+    else:
+      dataset = self._create_ragged_dataset(strategy)
+    data = next(
+        iter(
+            strategy.experimental_distribute_dataset(
+                dataset,
+                options=distribute_lib.InputOptions(
+                    experimental_fetch_to_device=False))))
+
+    @def_function.function
+    def embedding_and_set_gradients(data):
+      def tpu_fn():
+        activations = mid_level.dequeue()
+        mid_level.apply_gradients(nest.map_structure(array_ops.ones_like,
+                                                     activations))
+        return activations
+      mid_level.enqueue(data)
+      return strategy.run(tpu_fn)
+
+    @def_function.function
+    def embedding_only(data):
+      def tpu_fn():
+        return mid_level.dequeue()
+      mid_level.enqueue(data, training=False)
+      return strategy.run(tpu_fn)
+
+    # Only check core 0.
+    before_update = self._get_replica_numpy(
+        embedding_and_set_gradients(data), strategy, 0)
+    after_update = self._get_replica_numpy(embedding_only(data), strategy, 0)
+
+    # For videos table, row 0 and row 1 are looked up 3*num_replicas times as
+    # they occur 3 times per replica (considering the features 0 and 1 which are
+    # both looked up in the videos table).
+    # Feature 0 has ids [0, 0, 1], [0, 1, 1], ... repeated over num_replicas
+    # Feature 1 has ids [0, 1, 1], [0, 0, 1], ... repeated over num_replicas
+    # This means that both rows 0 and 1 get a -0.1*3*num_replicas update
+    # For users table, each row is looked up twice:
+    # Feature 2 has ids [3, 0, 1, 2], .. repeated over num_replicas
+    # This means that we get a -0.1*num_replicas update to the third feature.
+
+    # In general this means that after the update, if we lookup feature 0 and 1
+    # the values will be 0.3*num_replicas lower per entry and for feature 2 they
+    # will be 0.1*num_replicas lower.
+    # The one issue is that these lookups contain padding values.
+    # For core 0, we get the first 2 elements of the 4 element batch.
+    # For feature 0, the indices are [[0, 0], [1, 0], [1, 1]] with max sequence
+    # length of 2, which means that [0, 1] will be 0s.
+    # For feature 1, the indices are [[0, 0], [0, 1], [1, 0]] with max sequence
+    # length of 2, which means that [1, 1] will be 0s.
+    # For feature 2, the indices are [[0, 0], [1, 0], [1, 1], [1, 2]] with max
+    # sequence length of 3, which means that [0, 1], [0, 2] will be 0s.
+    # The following masks represent that so that we only apply the above updates
+    # to the non-padding rows:
+    masks = (
+        np.array([[[1], [0]], [[1], [1]]]),
+        np.array([[[1], [1]], [[1], [0]]]),
+        np.array([[[1], [0], [0]], [[1], [1], [1]]]))
+
+    per_row_update = (0.3 * num_replicas,
+                      0.3 * num_replicas,
+                      0.1 * num_replicas)
+    golden = tuple([before - update * mask for before, update, mask in
+                    zip(before_update, per_row_update, masks)])
+    self.assertAllClose(golden, after_update)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_sparse_forward_test.py

@@ -0,0 +1,37 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters(
+      ['sgd', 'adagrad', 'adam', 'ftrl', 'adagrad_momentum'])
+  def test_embedding(self, optimizer_name):
+    if optimizer_name != 'sgd':
+      self.skip_if_oss()
+    self._test_embedding(
+        optimizer_name, training=False, sparse=True, is_high_dimensional=False)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_correctness_sparse_training_test.py

@@ -0,0 +1,37 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_v2_correctness_base_test
+
+
+class TPUEmbeddingCorrectnessTest(
+    tpu_embedding_v2_correctness_base_test.TPUEmbeddingCorrectnessBaseTest):
+
+  @parameterized.parameters(
+      ['sgd', 'adagrad', 'adam', 'ftrl', 'adagrad_momentum'])
+  def test_embedding(self, optimizer_name):
+    if optimizer_name != 'sgd':
+      self.skip_if_oss()
+    self._test_embedding(
+        optimizer_name, training=True, sparse=True, is_high_dimensional=False)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_enqueue_mode_test.py

@@ -0,0 +1,205 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework.tensor_shape import TensorShape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+from tensorflow.python.util import nest
+
+
+class TPUEmbeddingTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  @parameterized.parameters([True, False])
+  def test_enqueue_with_outside_compilation(self, use_mlir):
+    if use_mlir:
+      config.enable_mlir_bridge()
+
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    mid_level_api.build([
+        TensorShape((self.batch_size, 2)),
+        TensorShape((self.batch_size, 2)),
+        TensorShape((self.batch_size, 3))
+    ])
+    dataset = self._create_sparse_dataset(strategy)
+    dataset_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def enqueue_with_outside_compilation(data):
+      def get_activations(features):
+        mid_level_api.enqueue(features, training=False)
+        return mid_level_api.dequeue()
+      return strategy.run(get_activations, args=(data,))
+
+    @def_function.function
+    def enqueue_without_outside_compilation(data):
+      def get_activations():
+        return mid_level_api.dequeue()
+      mid_level_api.enqueue(data, training=False)
+      return strategy.run(get_activations)
+
+    features = next(dataset_iter)
+
+    activations_oc = enqueue_with_outside_compilation(features)
+    activations = enqueue_without_outside_compilation(features)
+
+    # Extact per core numpy arrays.
+    activations_oc0 = self._get_replica_numpy(activations_oc, strategy, 0)
+    activations0 = self._get_replica_numpy(activations, strategy, 0)
+
+    self.assertAllClose(activations_oc0, activations0)
+
+  @parameterized.parameters(True, False)
+  def test_enqueue_with_outside_compilation_in_control_flow(self, use_mlir):
+    self.skip_if_oss()
+    if use_mlir:
+      config.enable_mlir_bridge()
+
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    dataset = self._create_sparse_dataset(strategy)
+    dataset_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    # This is one way to force the enqueue in some control flow. @tf.functions
+    # aren't inlined in the calling tf.function. An alternative would be to
+    # place the enqueue in a switch_v2 or something similar.
+    @def_function.function
+    def enqueue_fn(features):
+      mid_level_api.enqueue(features, training=False)
+
+    @def_function.function
+    def enqueue_with_outside_compilation():
+      def get_activations(features):
+        enqueue_fn(features)
+        return mid_level_api.dequeue()
+      return strategy.run(get_activations, args=(next(dataset_iter),))
+
+    with self.assertRaisesRegex(
+        RuntimeError,
+        'does not match graph which contains TPUReplicateContext'):
+      enqueue_with_outside_compilation()
+
+  def test_enqueue_with_outside_compilation_non_direct_input(self):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    mid_level_api.build([
+        TensorShape((self.batch_size, 2)),
+        TensorShape((self.batch_size, 2)),
+        TensorShape((self.batch_size, 3))
+    ])
+    dataset = self._create_sparse_dataset(strategy)
+    dataset_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def enqueue_with_outside_compilation():
+      def get_activations(features):
+        # This inserts a mul operation on the TPU to trigger the direct input
+        # error.
+        features = (features[0]*2, features[1]*2, features[2]*2)
+        mid_level_api.enqueue(features, training=False)
+        return mid_level_api.dequeue()
+      return strategy.run(get_activations, args=(next(dataset_iter),))
+
+    with self.assertRaisesRegex(
+        ValueError, 'which does not have the `_tpu_input_identity` attr'):
+      enqueue_with_outside_compilation()
+
+  def test_enqueue_with_outside_compilation_auto_mode(self):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    mid_level_api.build([
+        TensorShape((self.batch_size, 2)),
+        TensorShape((self.batch_size, 2)),
+        TensorShape((self.batch_size, 3))
+    ])
+    dataset = self._create_sparse_dataset(strategy)
+    dataset_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def enqueue_with_no_gradient_apply(data):
+      def get_activations(features):
+        # Note the lack of setting training=False, so training defaults to true
+        # here even though we don't have apply gradients.
+        # We detect the correct mode based on which ops exist that share the
+        # same 'name'.
+        mid_level_api.enqueue(features, name='call1')
+        return mid_level_api.dequeue(name='call1')
+      return strategy.run(get_activations, args=(data,))
+
+    @def_function.function
+    def enqueue_with_gradient_apply(data):
+      def get_activations(features):
+        mid_level_api.enqueue(features, name='call2')
+        activations = mid_level_api.dequeue(name='call2')
+        # Apply an all ones gradient
+        gradients = nest.map_structure(array_ops.ones_like, activations)
+        mid_level_api.apply_gradients(gradients, name='call2')
+        return activations
+      return strategy.run(get_activations, args=(data,))
+
+    data = next(dataset_iter)
+    before_gradient_apply = enqueue_with_gradient_apply(data)
+    after_gradient_apply = enqueue_with_no_gradient_apply(data)
+    before_gradient_apply0 = self._get_replica_numpy(before_gradient_apply,
+                                                     strategy, 0)
+    after_gradient_apply0 = self._get_replica_numpy(after_gradient_apply,
+                                                    strategy, 0)
+
+    num_replicas = strategy.num_replicas_in_sync
+    # We are passing a gradient of 1 for all lookups, optimizer is SGD with a
+    # learning rate of 0.1. Feature 0 and 1 are looked up with a sum combiner
+    # with the following ids:
+    # Feature 0: [0, 0, 1], [0, 1, 1], ... repeated over num_replicas
+    # Feature 1: [0, 1, 1], [0, 0, 1], ... repeated over num_replicas
+    # i.e. Row 0 and 1 were looked up 3*num_replicas times over all cores and as
+    # the gradient is 1, the accumulated gradient is 3*num_replicas for each
+    # position in row 0 and 1 in table.
+    #
+    # See comments in test_pass_none_to_apply_gradients for the update to
+    # Feature 2 and its table.
+    # The *2 in the next tests are because those rows have 2 lookups vs
+    # the 1 lookup in the other row.
+    update = ([[0.3 * num_replicas], [0.3 * num_replicas * 2]],
+              [[0.3 * num_replicas * 2], [0.3 * num_replicas]],
+              [[0.1 * num_replicas], [0.1 / 3 * num_replicas]])
+    golden = tuple([before - np.array(up) for before, up in
+                    zip(before_gradient_apply0, update)])
+
+    self.assertAllClose(golden, after_gradient_apply0)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_hd_invalid_input_test.py

@@ -0,0 +1,83 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework.tensor_shape import TensorShape
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+
+
+class TPUEmbeddingTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  def test_build_incorrect_output_shapes(self):
+    _, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    # Output shapes is set in the mid_level_api, but build with incorrect output
+    # shapes.
+    mid_level_api._output_shapes = [TensorShape((2, 4)) for _ in range(3)]
+
+    with self.assertRaisesRegex(ValueError,
+                                'Inconsistent shape founded for input feature'):
+      mid_level_api.build([TensorShape([1, 1, 1]) for _ in range(3)])
+
+  def test_enqueue_incorrect_shape_feature(self):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    sparse = self._create_high_dimensional_sparse_dataset(strategy)
+    sparse_iter = iter(
+        strategy.experimental_distribute_dataset(
+            sparse,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    mid_level_api._output_shapes = [TensorShape((1, 1)) for _ in range(3)]
+    # The output shape passed to build method is consistent.
+    mid_level_api.build([TensorShape([1, 1, 1]) for _ in range(3)])
+
+    @def_function.function
+    def test_fn():
+
+      def step():
+        return mid_level_api.dequeue()
+
+      mid_level_api.enqueue(next(sparse_iter), training=False)
+      return strategy.run(step)
+
+    # Enqueued tensor has shape inconsistent with the output shape setting.
+    with self.assertRaisesRegex(ValueError,
+                                'Inconsistent shape founded for input feature'):
+      test_fn()
+
+  def test_not_fully_defined_output_shapes_in_feature_config(self):
+    _, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    # Feature config sets undefined output shapes
+    mid_level_api._output_shapes = [TensorShape(None) for _ in range(3)]
+    with self.assertRaisesRegex(ValueError, 'Input Feature'):
+      mid_level_api.build()
+
+  def test_not_fully_defined_output_shapes_for_build(self):
+    _, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    # Build with undefined output shape
+    with self.assertRaisesRegex(ValueError, 'Input Feature'):
+      mid_level_api.build([TensorShape([1, None, None]) for _ in range(3)])
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_hd_valid_input_test.py

@@ -0,0 +1,125 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+
+
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework.tensor_shape import TensorShape
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+
+
+class TPUEmbeddingTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  def test_enqueue_dense_sparse_ragged(self):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    dataset = self._create_high_dimensional_dense_dataset(strategy)
+    dense_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    sparse = self._create_high_dimensional_sparse_dataset(strategy)
+    sparse_iter = iter(
+        strategy.experimental_distribute_dataset(
+            sparse,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    ragged = self._create_high_dimensional_ragged_dataset(strategy)
+    ragged_iter = iter(
+        strategy.experimental_distribute_dataset(
+            ragged,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    mid_level_api.build([
+        TensorShape([self.batch_size, self.data_batch_size, 1]),
+        TensorShape([self.batch_size, self.data_batch_size, 2]),
+        TensorShape([self.batch_size, self.data_batch_size, 3])
+    ])
+
+    @def_function.function
+    def test_fn():
+
+      def step():
+        return mid_level_api.dequeue()
+
+      features = (next(dense_iter)[0], next(sparse_iter)[1],
+                  next(ragged_iter)[2])
+      mid_level_api.enqueue(features, training=False)
+      return strategy.run(step)
+
+    test_fn()
+
+  def test_different_input_shapes(self):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    sparse = self._create_high_dimensional_sparse_dataset(strategy)
+    sparse_iter = iter(
+        strategy.experimental_distribute_dataset(
+            sparse,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+    # Create a feature with shape (1, 3, 1)
+    dense_feature = constant_op.constant(
+        np.zeros(3), shape=(1, 3, 1), dtype=dtypes.int32)
+    dense_dataset = dataset_ops.DatasetV2.from_tensors(
+        dense_feature).unbatch().repeat().batch(
+            1 * strategy.num_replicas_in_sync, drop_remainder=True)
+    dense_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dense_dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def test_fn():
+
+      def step():
+        return mid_level_api.dequeue()
+
+      features = (next(dense_iter), next(sparse_iter)[1], next(sparse_iter)[2])
+      mid_level_api.enqueue(features, training=False)
+      return strategy.run(step)
+
+    test_fn()
+
+    self.assertEqual(mid_level_api._output_shapes, [
+        TensorShape((1, 3)),
+        TensorShape((self.batch_size, self.data_batch_size)),
+        TensorShape((self.batch_size, self.data_batch_size))
+    ])
+
+  def test_output_shapes_priority_over_feature_config_and_build(self):
+    _, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    # The output shapes setting in the feature config has the first priority.
+    mid_level_api._output_shapes = [TensorShape((2, 4)) for _ in range(3)]
+    mid_level_api.build([TensorShape((2, None, None)) for _ in range(3)])
+    self.assertEqual(mid_level_api._output_shapes,
+                     [TensorShape((2, 4)) for _ in range(3)])
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_initialization_test.py

@@ -0,0 +1,129 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework.tensor_shape import TensorShape
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v2
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+
+
+class TPUEmbeddingTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  def test_enqueue_dequeue_apply_gradients_on_cpu(self):
+    # Dequeue on CPU.
+    mid_level_api = self._create_mid_level()
+    with self.assertRaises(RuntimeError):
+      mid_level_api.dequeue()
+    # Enqueue on CPU.
+    features = {
+        'watched': sparse_tensor.SparseTensor(
+            indices=self.feature_watched_indices,
+            values=self.feature_watched_values,
+            dense_shape=[2, 2])}
+    with self.assertRaises(RuntimeError):
+      mid_level_api.enqueue(features)
+    # Apply gradient on CPU.
+    mid_level_api = self._create_mid_level()
+    with self.assertRaises(RuntimeError):
+      mid_level_api.apply_gradients(None)
+
+  def test_multiple_creation(self):
+    feature_config = tpu_embedding_v2_utils.FeatureConfig(
+        table=self.table_user, name='friends', max_sequence_length=2)
+    optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+    strategy = self._get_strategy()
+    with strategy.scope():
+      embedding_one = tpu_embedding_v2.TPUEmbedding(
+          feature_config=feature_config, optimizer=optimizer)
+      embedding_two = tpu_embedding_v2.TPUEmbedding(
+          feature_config=feature_config, optimizer=optimizer)
+
+    # The first TPU embedding should be able to be built.
+    # The second one should fail with a runtime error indicating another TPU
+    # embedding has already been initialized on TPU.
+    embedding_one.build(64)
+    with self.assertRaisesRegex(RuntimeError,
+                                'TPU is already initialized for embeddings.'):
+      embedding_two.build(64)
+
+  def test_same_config_different_instantiations(self):
+    num_tables = 30
+    table_dim = np.random.randint(1, 128, size=[num_tables])
+    table_vocab_size = np.random.randint(100, 1000, size=[num_tables])
+    table_names = ['table{}'.format(i) for i in range(num_tables)]
+    table_data = list(zip(table_dim, table_vocab_size, table_names))
+    strategy = self._get_strategy()
+
+    def tpu_embedding_config():
+      feature_configs = []
+      for dim, vocab, name in table_data:
+        optimizer = None
+        if dim % 2 == 0:
+          optimizer = tpu_embedding_v2_utils.Adagrad(
+              learning_rate=lambda: constant_op.constant(1.0))
+
+        feature_configs.append(
+            tpu_embedding_v2_utils.FeatureConfig(
+                table=tpu_embedding_v2_utils.TableConfig(
+                    vocabulary_size=int(vocab),
+                    dim=int(dim),
+                    initializer=init_ops_v2.Zeros(),
+                    optimizer=optimizer,
+                    name=name)))
+      optimizer = tpu_embedding_v2_utils.Adagrad(learning_rate=0.1)
+      with strategy.scope():
+        mid_level_api = tpu_embedding_v2.TPUEmbedding(
+            feature_config=feature_configs, optimizer=optimizer)
+      mid_level_api._output_shapes = [TensorShape(128)] * len(feature_configs)
+      return mid_level_api._create_config_proto()
+
+    self.assertProtoEquals(tpu_embedding_config(), tpu_embedding_config())
+
+  def test_learning_rate_tag_order(self):
+    num_tables = 30
+    strategy = self._get_strategy()
+
+    feature_configs = []
+    for i in range(num_tables):
+      optimizer = tpu_embedding_v2_utils.Adagrad(
+          learning_rate=lambda: constant_op.constant(1.0))
+
+    feature_configs.append(
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=tpu_embedding_v2_utils.TableConfig(
+                vocabulary_size=100,
+                dim=128,
+                initializer=init_ops_v2.Zeros(),
+                optimizer=optimizer)))
+    with strategy.scope():
+      mid_level_api = tpu_embedding_v2.TPUEmbedding(
+          feature_config=feature_configs, optimizer=optimizer)
+    mid_level_api._output_shapes = [TensorShape(128)] * len(feature_configs)
+    result = mid_level_api._create_config_proto()
+    for i, table in enumerate(result.table_descriptor):
+      self.assertEqual(i,
+                       table.optimization_parameters.learning_rate.dynamic.tag)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_invalid_input_test.py

@@ -0,0 +1,226 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v2
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+
+
+class TPUEmbeddingTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  def test_tables_with_same_name(self):
+    with self.assertRaisesRegex(
+        ValueError, 'Multiple tables with name table found.'):
+      with self._get_strategy().scope():
+        tpu_embedding_v2.TPUEmbedding(
+            (tpu_embedding_v2_utils.FeatureConfig(
+                table=tpu_embedding_v2_utils.TableConfig(
+                    name='table',
+                    vocabulary_size=4,
+                    dim=2,
+                    initializer=self.initializer,),
+                name='watched'),
+             tpu_embedding_v2_utils.FeatureConfig(
+                 table=tpu_embedding_v2_utils.TableConfig(
+                     name='table',
+                     vocabulary_size=4,
+                     dim=2,
+                     initializer=self.initializer),
+                 name='favorited')),
+            tpu_embedding_v2_utils.SGD(learning_rate=0.1))
+
+  def test_pass_non_tensor_to_apply_gradients(self):
+    self.skip_if_oss()
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    # We aren't going to actually run anything, so the batch_size here does not
+    # matter.
+    mid_level_api.build(64)
+
+    # Test pass non tensor to apply_gradients.
+    @def_function.function
+    def test_apply_1():
+      mid_level_api.apply_gradients((1, 2, 3))
+
+    with self.assertRaisesRegex(ValueError, 'found non-tensor type'):
+      strategy.run(test_apply_1)
+
+    # Test pass different structure to apply_gradients.
+    @def_function.function
+    def test_apply_2():
+      # This should be a tuple as feature_config is a tuple of 3 configs.
+      mid_level_api.apply_gradients([1, 2, 3])
+
+    with self.assertRaisesRegex(
+        TypeError, 'The two structures don\'t have the same nested structure.'):
+      strategy.run(test_apply_2)
+
+  def test_enqueue_weight_for_dense_tensor(self):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    dataset = self._create_dense_dataset(strategy, include_weights=True)
+    dense_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def test_fn():
+      def step():
+        return mid_level_api.dequeue()
+
+      features, weights = next(dense_iter)
+      mid_level_api.enqueue(features, weights=weights, training=False)
+      return strategy.run(step)
+
+    with self.assertRaisesRegex(ValueError, 'Weight specified for dense input'):
+      test_fn()
+
+  def test_enqueue_wrong_weight_type_for_sparse_and_ragged_tensor(self):
+    self.skip_if_oss()
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    sparse = self._create_sparse_dataset(strategy, include_weights=True)
+    ragged = self._create_ragged_dataset(strategy, include_weights=True)
+    sparse_iter = iter(
+        strategy.experimental_distribute_dataset(
+            sparse,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+    ragged_iter = iter(
+        strategy.experimental_distribute_dataset(
+            ragged,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def test_sparse_fn():
+      def step():
+        return mid_level_api.dequeue()
+
+      features, _ = next(sparse_iter)
+      _, weights = next(ragged_iter)
+      mid_level_api.enqueue(features, weights=weights, training=False)
+      return strategy.run(step)
+
+    with self.assertRaisesRegex(
+        ValueError, 'which does not match type input which is SparseTensor.'):
+      test_sparse_fn()
+
+    @def_function.function
+    def test_ragged_fn():
+      def step():
+        return mid_level_api.dequeue()
+
+      _, weights = next(sparse_iter)
+      features, _ = next(ragged_iter)
+      mid_level_api.enqueue(features, weights=weights, training=False)
+      return strategy.run(step)
+
+    with self.assertRaisesRegex(
+        ValueError, 'which does not match type input which is RaggedTensor.'):
+      test_ragged_fn()
+
+  def test_enqueue_incorrect_structure_for_features_and_weights(self):
+    self.skip_if_oss()
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    sparse = self._create_sparse_dataset(strategy, include_weights=True)
+    sparse_iter = iter(
+        strategy.experimental_distribute_dataset(
+            sparse,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def test_features_fn():
+      def step():
+        return mid_level_api.dequeue()
+
+      features = next(sparse_iter)
+      features = (features[0],)
+      mid_level_api.enqueue(features, training=False)
+      return strategy.run(step)
+
+    # The error here is raised from nest.assert_same_structure
+    with self.assertRaises(ValueError):
+      test_features_fn()
+
+    @def_function.function
+    def test_weights_fn():
+      def step():
+        return mid_level_api.dequeue()
+
+      features, weights = next(sparse_iter)
+      weights = (weights[0],)
+      mid_level_api.enqueue(features, weights=weights, training=False)
+      return strategy.run(step)
+
+    # The error here is raised from nest.assert_same_structure
+    with self.assertRaises(ValueError):
+      test_weights_fn()
+
+  def test_enqueue_cpu_tensor(self):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    dataset = self._create_dense_dataset(strategy)
+    dense_iter = iter(strategy.experimental_distribute_dataset(dataset))
+
+    @def_function.function
+    def test_fn():
+      def get_activations():
+        return mid_level_api.dequeue()
+
+      features = next(dense_iter)
+      mid_level_api.enqueue(features, training=False)
+      activations = strategy.run(get_activations)
+      return activations
+
+    with self.assertRaisesRegex(ValueError, 'which is on a TPU input device'):
+      test_fn()
+
+  @parameterized.parameters([True, False])
+  def test_enqueue_cpu_tensor_with_outside_compilation(self, use_mlir):
+
+    if use_mlir:
+      config.enable_mlir_bridge()
+
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    dataset = self._create_dense_dataset(strategy)
+    dense_iter = iter(strategy.experimental_distribute_dataset(dataset))
+
+    @def_function.function
+    def test_fn():
+      def get_activations(features):
+        mid_level_api.enqueue(features, training=False)
+        return mid_level_api.dequeue()
+
+      activations = strategy.run(get_activations, args=(next(dense_iter),))
+      return activations
+
+    with self.assertRaisesRegex(ValueError, 'which is on a TPU input device'):
+      test_fn()
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_mp_strategy_test.py

@@ -0,0 +1,199 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import tpu_strategy
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables as tf_variables
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import device_assignment as device_lib
+from tensorflow.python.tpu import tpu_embedding_v2
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+
+
+class TPUEmbeddingTPUStrategyV2Test(
+    tpu_embedding_base_test.TPUEmbeddingBaseTest
+):
+
+  def setUp(self):
+    super().setUp()
+
+    self._num_replicas = 1
+    self._num_cores_per_replica = 2
+
+  def _get_strategy(self) -> tpu_strategy.TPUStrategy:
+    topology = self._init_tpu_system()
+
+    d_assign = device_lib.device_assignment(
+        topology,
+        computation_shape=[1, 1, 1, 2],
+        num_replicas=1,
+    )
+
+    self.strategy = tpu_strategy.TPUStrategyV2(
+        self.resolver,
+        experimental_device_assignment=d_assign,
+        experimental_spmd_xla_partitioning=True,
+    )
+
+    self.embedding_devices = sum(
+        (list(replica) for replica in self.strategy.extended._tpu_devices), []
+    )
+
+    return self.strategy
+
+  def enqueue(self, inp, mid_level_api, use_device, training):
+    if use_device:
+      for emb, device in zip(inp, self.embedding_devices):
+        mid_level_api.enqueue(emb, device=device, training=training)
+    else:
+      mid_level_api.enqueue(inp[0], training=training)
+
+  @parameterized.parameters(False, True)
+  def test_spmd_training(self, use_device):
+    num_steps = 10
+    num_steps_float = float(num_steps)
+    starting_lr = 1.0
+    ending_lr = 0.5
+
+    strategy = self._get_strategy()
+
+    # Create model with Keras.
+    with strategy.scope():
+      step_counter = tf_variables.Variable(0.0, dtypes.float32)
+
+      def lr_function():
+        return gen_math_ops.maximum(
+            ending_lr,
+            starting_lr
+            + ((ending_lr - starting_lr) * step_counter) / num_steps_float,
+        )
+
+      optimizer = tpu_embedding_v2_utils.SGD(learning_rate=lr_function)
+      table_config = tpu_embedding_v2_utils.TableConfig(
+          vocabulary_size=2,
+          dim=4,
+          initializer=init_ops_v2.Constant(np.zeros((2, 4))),
+          combiner='sum',
+          name='table',
+      )
+      mid_level_api = tpu_embedding_v2.TPUEmbedding(
+          feature_config={
+              'feature': tpu_embedding_v2_utils.FeatureConfig(
+                  table=table_config, name='feature'
+              )
+          },
+          optimizer=optimizer,
+      )
+
+    def input_fn(ctx):
+      del ctx
+      feature = {
+          'feature': constant_op.constant(
+              [0, 1], shape=(2, 1), dtype=dtypes.int32
+          )
+      }
+      return dataset_ops.DatasetV2.from_tensors(feature).repeat()
+
+    def create_datasets():
+      """Creates either a per-replica dataset, or multiple per-devices ones.
+
+      This function explicitly creates per-device datasets because the strategy
+      does not produce a distributed dataset in the model-parallel case; there
+      is only one replica. Without this consideration, the embeddings would be
+      read as [0, 0] instead of the expected [0, 1] since all the devices would
+      receive the same value.
+
+      Returns:
+        A list of one or more dataset(s).
+      """
+      if use_device:
+        datasets = []
+        for i in range(len(self.embedding_devices)):
+          datasets.append(
+              dataset_ops.DatasetV2.from_tensor_slices(
+                  {'feature': [[[i % self._num_cores_per_replica]]]}
+              ).repeat()
+          )
+        return datasets
+      else:
+        dataset = strategy.distribute_datasets_from_function(
+            input_fn,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False
+            ),
+        )
+        return [dataset]
+
+    datasets = create_datasets()
+    iterators = [iter(ds) for ds in datasets]
+
+    @def_function.function(jit_compile=True)
+    def test_fn():
+      def step():
+        with backprop.GradientTape() as tape:
+          activations = mid_level_api.dequeue()
+          tape.watch(activations)
+          result = math_ops.reduce_sum(activations['feature'])
+          loss = result / self._num_replicas
+        grads = tape.gradient(loss, activations)
+        mid_level_api.apply_gradients(grads)
+        return activations
+
+      inp = [next(it) for it in iterators]
+      self.enqueue(inp, mid_level_api, use_device, training=True)
+      return strategy.run(step)
+
+    # Run model.
+    results = []
+    for _ in range(num_steps):
+      result = test_fn()
+      results.append(self._unpack(strategy, result['feature']))
+      step_counter.assign_add(1.0)
+
+    # Table is 2 elements wide, per-replica batch size of 1, with id 0.
+    # Loss for the gradient is the sum of the entries divided by the number of
+    # replicas. Thus the per replica gradient is 1/#of replicas for row 0 and no
+    # other updates. The reduced gradient is therefore 1.
+    # Learning rate schedule over num_steps steps:
+    # 1.0 0.95 0.9 0.85 0.8 ...
+    # Since use SGD and the gradient is one, the first row of the table is
+    # [0, 0] [-1.0, -1.0] [-1.95, -1.95] [-2.85, -2.85] ... (the negative
+    # partial sums of the above).
+
+    learning_rates = [starting_lr - (starting_lr - ending_lr) / num_steps * j
+                      for j in range(num_steps)]
+    cumsum = [sum(learning_rates[0:j]) for j in range(num_steps)]
+    goldens = [[[-cumsum[i]] * table_config.dim] * self._num_cores_per_replica
+               for i in range(10)]
+    self.assertAllClose(results, goldens)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_sequence_feature_test.py

@@ -0,0 +1,107 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework.tensor_shape import TensorShape
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+
+
+class TPUEmbeddingTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  @parameterized.parameters([True, False])
+  def test_sequence_feature(self, is_sparse):
+    seq_length = 3
+    # Set the max_seq_length in feature config
+    for feature in self.feature_config:
+      feature.max_sequence_length = seq_length
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    if is_sparse:
+      dataset = self._create_sparse_dataset(strategy)
+    else:
+      dataset = self._create_ragged_dataset(strategy)
+    feature_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def test_fn():
+
+      def step():
+        return mid_level_api.dequeue()
+
+      mid_level_api.enqueue(next(feature_iter), training=False)
+      return strategy.run(step)
+
+    output = test_fn()
+    self.assertEqual(
+        self._get_replica_numpy(output[0], strategy, 0).shape, (2, 3, 4))
+    self.assertEqual(
+        self._get_replica_numpy(output[1], strategy, 0).shape, (2, 3, 4))
+    self.assertEqual(
+        self._get_replica_numpy(output[2], strategy, 0).shape, (2, 3, 2))
+
+  @parameterized.parameters([True, False])
+  def test_sequence_feature_with_build(self, is_updated_shape):
+    seq_length = 3
+    # Set the max_seq_length in feature config
+    for feature in self.feature_config:
+      feature.max_sequence_length = seq_length
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    dataset = self._create_sparse_dataset(strategy)
+    feature_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+    if is_updated_shape:
+      mid_level_api.build([
+          TensorShape([self.batch_size, seq_length, 2]),
+          TensorShape([self.batch_size, seq_length, 2]),
+          TensorShape([self.batch_size, seq_length, 3])
+      ])
+    else:
+      mid_level_api.build([
+          TensorShape([self.batch_size, 2]),
+          TensorShape([self.batch_size, 2]),
+          TensorShape([self.batch_size, 3])
+      ])
+
+    @def_function.function
+    def test_fn():
+
+      def step():
+        return mid_level_api.dequeue()
+
+      mid_level_api.enqueue(next(feature_iter), training=False)
+      return strategy.run(step)
+
+    output = test_fn()
+    self.assertEqual(
+        self._get_replica_numpy(output[0], strategy, 0).shape, (2, 3, 4))
+    self.assertEqual(
+        self._get_replica_numpy(output[1], strategy, 0).shape, (2, 3, 4))
+    self.assertEqual(
+        self._get_replica_numpy(output[2], strategy, 0).shape, (2, 3, 2))
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_embedding_v2_valid_input_test.py

@@ -0,0 +1,288 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Embeddings mid level API on TPU."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework.tensor_shape import TensorShape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v2
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu.tests import tpu_embedding_base_test
+from tensorflow.python.util import nest
+
+
+class TPUEmbeddingTest(tpu_embedding_base_test.TPUEmbeddingBaseTest):
+
+  def test_pass_none_to_apply_gradients(self):
+    self.skip_if_oss()
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    mid_level_api.build([
+        TensorShape((self.batch_size, 2)),
+        TensorShape((self.batch_size, 2)),
+        TensorShape((self.batch_size, 3))
+    ])
+    dataset = self._create_sparse_dataset(strategy)
+    data = next(
+        iter(
+            strategy.experimental_distribute_dataset(
+                dataset,
+                options=distribute_lib.InputOptions(
+                    experimental_fetch_to_device=False))))
+
+    @def_function.function
+    def embedding_and_set_gradients(data):
+      mid_level_api.enqueue(data)
+      def tpu_fn():
+        results = mid_level_api.dequeue()
+        mid_level_api.apply_gradients((None, None,
+                                       array_ops.ones_like(results[2])))
+        return results
+      return strategy.run(tpu_fn)
+
+    @def_function.function
+    def embedding_only(data):
+      mid_level_api.enqueue(data, training=False)
+      def tpu_fn():
+        return mid_level_api.dequeue()
+      return strategy.run(tpu_fn)
+
+    first = self._get_replica_numpy(
+        embedding_and_set_gradients(data), strategy, 0)
+    second = self._get_replica_numpy(embedding_only(data), strategy, 0)
+
+    # First two features should be the same as None gradient was applied.
+    # Third feature had gradient of 1 passed in from each core.
+    # Each core received the same ids per core and returned the following batch:
+    # [ row 3, row 0 + row 1 + row 2 ]
+    # so gradient update was (learning rate = 0.1):
+    #   row 0: -1/3*0.1
+    #   row 1: -1/3*0.1
+    #   row 2: -1/3*0.1
+    #   row 3: -1*0.1
+    # There is a factor of num_replicas because each replica gave an update.
+
+    num_replicas = strategy.num_replicas_in_sync
+    update = ([[0.0]], [[0.0]],
+              [[0.1 * num_replicas], [0.1 / 3 * num_replicas]])
+    golden = tuple([feature-np.array(up) for feature, up in zip(first, update)])
+
+    self.assertAllClose(golden, second)
+
+  def test_enqueue_sparse_and_ragged(self):
+    self.skip_if_oss()
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    sparse = self._create_sparse_dataset(strategy)
+    ragged = self._create_ragged_dataset(strategy)
+    sparse_iter = iter(
+        strategy.experimental_distribute_dataset(
+            sparse,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+    ragged_iter = iter(
+        strategy.experimental_distribute_dataset(
+            ragged,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def test_fn():
+      def step():
+        return mid_level_api.dequeue()
+
+      sparse_features = next(sparse_iter)
+      ragged_features = next(ragged_iter)
+      features = (sparse_features[0], ragged_features[1], sparse_features[2])
+      mid_level_api.enqueue(features, training=False)
+      return strategy.run(step)
+
+    test_fn()
+
+  def test_enqueue_per_device(self):
+    self.skip_if_oss()
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+
+    sparse = self._create_sparse_dataset(strategy)
+    sparse_iter = iter(
+        strategy.experimental_distribute_dataset(
+            sparse,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def test_fn():
+      def get_activations(dense_value):
+        return mid_level_api.dequeue(), dense_value
+
+      sparse_features = next(sparse_iter)
+      mid_level_api.enqueue(sparse_features, training=False)
+      activations, dense_value1 = strategy.run(get_activations, args=(0.0,))
+
+      def enqueue_fn(ctx):
+        core_id = ctx.replica_id_in_sync_group
+        device = strategy.extended.worker_devices[core_id]
+        sparse_features_local = nest.map_structure(
+            lambda x: strategy.experimental_local_results(x)[core_id],
+            sparse_features)
+        mid_level_api.enqueue(sparse_features_local, training=False,
+                              device=device)
+        return 0.0
+
+      data = strategy.experimental_distribute_values_from_function(
+          enqueue_fn)
+      per_device_activations, dense_value2 = strategy.run(get_activations,
+                                                          args=(data,))
+      return activations, per_device_activations, dense_value1, dense_value2
+
+    activations, per_device_activations, _, _ = test_fn()
+
+    # Extact per core numpy arrays and check that both sparse and ragged have
+    # the same results.
+    activations0 = self._get_replica_numpy(activations, strategy, 0)
+    per_device_activations0 = self._get_replica_numpy(
+        per_device_activations, strategy, 0)
+    self.assertAllClose(activations0, per_device_activations0)
+    test_fn()
+
+  @parameterized.parameters(True, False)
+  def test_enqueue_with_weights(self, ragged):
+    strategy, mid_level_api, _ = self._create_strategy_and_mid_level('sgd')
+    weight = 0.5
+    if ragged:
+      dataset = self._create_ragged_dataset(strategy, include_weights=True,
+                                            weight=weight)
+    else:
+      dataset = self._create_sparse_dataset(strategy, include_weights=True,
+                                            weight=weight)
+      mid_level_api.build([
+          TensorShape((self.batch_size, 2)),
+          TensorShape((self.batch_size, 2)),
+          TensorShape((self.batch_size, 3))
+      ])
+
+    dataset_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def enqueue_and_get(features, weights):
+      def get_activations():
+        return mid_level_api.dequeue()
+      mid_level_api.enqueue(features, weights=weights, training=False)
+      return strategy.run(get_activations)
+
+    features, weights = next(dataset_iter)
+    # Replace the weight for the second feature by None to test.
+    weights = (weights[0], None, weights[2])
+
+    no_weights_activations = enqueue_and_get(features, weights=None)
+    weights_activations = enqueue_and_get(features, weights=weights)
+
+    # Extact per core numpy arrays.
+    no_weights0 = self._get_replica_numpy(no_weights_activations, strategy, 0)
+    weights0 = self._get_replica_numpy(weights_activations, strategy, 0)
+    # videos table has sum combiner and users table has mean combiner.
+    # i.e. users table lookups isn't affected by the weights as all the weights
+    # are the same.
+    # Tuple entry 0 and 1 are the watched and favorited features from the videos
+    # table and entry 2 is the friends feature from the users table.
+    # Note that None was passed as a weight for entry 1 so weight should have no
+    # effect.
+    weight = (0.5, 1.0, 1.0)
+    golden = tuple([no_weight * w for no_weight, w in zip(no_weights0, weight)])
+
+    self.assertAllClose(golden, weights0)
+
+  def test_same_config_different_instantiations(self):
+    self.skip_if_oss()
+    num_tables = 30
+    table_dim = np.random.randint(1, 128, size=[num_tables])
+    table_vocab_size = np.random.randint(100, 1000, size=[num_tables])
+    table_names = ['table{}'.format(i) for i in range(num_tables)]
+    table_data = list(zip(table_dim, table_vocab_size, table_names))
+    strategy = self._get_strategy()
+
+    def tpu_embedding_config():
+      feature_configs = []
+      for dim, vocab, name in table_data:
+        feature_configs.append(tpu_embedding_v2_utils.FeatureConfig(
+            table=tpu_embedding_v2_utils.TableConfig(
+                vocabulary_size=int(vocab), dim=int(dim),
+                initializer=init_ops_v2.Zeros(), name=name)))
+      optimizer = tpu_embedding_v2_utils.Adagrad(
+          learning_rate=0.1)
+      with strategy.scope():
+        mid_level_api = tpu_embedding_v2.TPUEmbedding(
+            feature_config=feature_configs,
+            optimizer=optimizer)
+      mid_level_api._output_shapes = [TensorShape(128)] * len(feature_configs)
+      return mid_level_api._create_config_proto()
+
+    self.assertProtoEquals(tpu_embedding_config(), tpu_embedding_config())
+
+  @parameterized.parameters([True, False])
+  def test_missing_feature(self, is_sparse):
+    strategy = self._get_strategy()
+    with strategy.scope():
+      optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+      mid_level_api = tpu_embedding_v2.TPUEmbedding(
+          feature_config=tpu_embedding_v2_utils.FeatureConfig(
+              table=self.table_video, name='watched'),
+          optimizer=optimizer)
+    # Create sparse or ragged feature with last sample missing.
+    if is_sparse:
+      features = sparse_tensor.SparseTensor(
+          indices=self.feature_watched_indices[:-1],
+          values=self.feature_watched_values[:-1],
+          dense_shape=[self.data_batch_size, 2])
+    else:
+      features = ragged_tensor.RaggedTensor.from_row_lengths(
+          row_lengths=[1, 2, 2, 0], values=self.feature_watched_values[:-1])
+
+    dataset = dataset_ops.DatasetV2.from_tensors(features)
+
+    dataset = dataset.unbatch().repeat().batch(
+        self.batch_size * strategy.num_replicas_in_sync, drop_remainder=True)
+    dataset_iter = iter(
+        strategy.experimental_distribute_dataset(
+            dataset,
+            options=distribute_lib.InputOptions(
+                experimental_fetch_to_device=False)))
+
+    @def_function.function
+    def test_fn():
+
+      def get_activations():
+        return mid_level_api.dequeue()
+
+      mid_level_api.enqueue(next(dataset_iter), training=False)
+      return strategy.run(get_activations)
+
+    test_fn()
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tests/tpu_initialization_test.py

@@ -0,0 +1,33 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for TPU Initialization."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.platform import test
+
+
+class TPUInitializationTest(parameterized.TestCase, test.TestCase):
+
+  def test_tpu_initialization(self):
+    resolver = tpu_cluster_resolver.TPUClusterResolver('')
+    tpu_cluster_resolver.initialize_tpu_system(resolver)
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()