chore: import upstream snapshot with attribution

tensorflow/python/tpu/__init__.py

@@ -0,0 +1,22 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Ops related to Tensor Processing Units."""
+
+import os
+from tensorflow.python.framework import versions
+
+os.environ['TPU_ML_PLATFORM'] = 'Tensorflow'
+os.environ['TPU_ML_PLATFORM_VERSION'] = versions.__version__

tensorflow/python/tpu/_pywrap_sparse_core_layout.pyi

@@ -0,0 +1,22 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+class SparseCoreLayoutStacker:
+    def __init__(self, num_partitions: int, disable_table_stacking: bool, sparse_cores_per_partition: int) -> None: ...
+    def AddTable(self, table_name: str, table_height: int, table_width: int, group: str, output_samples: int, num_features: int) -> None: ...
+    def GetLayouts(self, *args, **kwargs): ...
+    def SetActivationMemoryBytesLimit(self, arg0: int) -> None: ...
+    def SetStackingEnabled(self, arg0: bool) -> None: ...
+    def SetVariableShardBytesLimit(self, arg0: int) -> None: ...

tensorflow/python/tpu/_pywrap_tpu_embedding.pyi

@@ -0,0 +1,16 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+def stack_tables(arg0: list[int], arg1: list[int], arg2: list[int], arg3: list[int], arg4: list[str], arg5: int) -> list[list[str]]: ...

tensorflow/python/tpu/api.py

@@ -0,0 +1,33 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Modules that need to be exported to the API.
+
+List TPU modules that aren't included elsewhere here so that they can be scanned
+for tf_export decorations.
+"""
+
+# pylint: disable=unused-import
+from tensorflow.python.tpu import bfloat16
+from tensorflow.python.tpu import feature_column_v2
+from tensorflow.python.tpu import tpu
+
+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
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu import tpu_embedding_v3
+from tensorflow.python.tpu import tpu_hardware_feature
+from tensorflow.python.tpu import tpu_optimizer
+# pylint: enable=unused-import

tensorflow/python/tpu/async_checkpoint.py

@@ -0,0 +1,269 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the 'License');
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an 'AS IS' BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ======================================
+"""Hook for asynchronous checkpointing.
+
+This hook dispatches checkpoint writing operations in a separate thread to
+allow execution to continue on the main thread.
+"""
+
+import os
+import threading
+import time
+from typing import Any, List, Optional, Text
+
+from tensorflow.core.util import event_pb2
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import meta_graph
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model.pywrap_saved_model import metrics
+from tensorflow.python.training import basic_session_run_hooks
+from tensorflow.python.training import monitored_session
+from tensorflow.python.training import saver as saver_lib
+from tensorflow.python.training import session_run_hook
+from tensorflow.python.training import training_util
+from tensorflow.python.training.summary_io import SummaryWriterCache
+
+
+# Captures the timestamp of the first Saver object instantiation or end of a
+# save operation. Can be accessed by multiple Saver instances.
+_END_TIME_OF_LAST_WRITE = None
+_END_TIME_OF_LAST_WRITE_LOCK = threading.Lock()
+
+# API label for cell names used in TF1 async checkpoint metrics.
+_ASYNC_CHECKPOINT_V1 = "async_checkpoint_v1"
+
+
+def _get_duration_microseconds(start_time_seconds, end_time_seconds) -> int:
+  """Returns the duration between start and end time in microseconds."""
+  return max(int((end_time_seconds - start_time_seconds) * 1000000), 0)
+
+
+class AsyncCheckpointSaverHook(basic_session_run_hooks.CheckpointSaverHook):
+  """Saves checkpoints every N steps or seconds."""
+
+  def __init__(self,
+               checkpoint_dir: Text,
+               save_secs: Optional[int] = None,
+               save_steps: Optional[int] = None,
+               saver: Optional[saver_lib.Saver] = None,
+               checkpoint_basename: Text = "model.ckpt",
+               scaffold: Optional[monitored_session.Scaffold] = None,
+               listeners: Optional[List[
+                   basic_session_run_hooks.CheckpointSaverListener]] = None):
+    """Initializes a `CheckpointSaverHook`.
+
+    Args:
+      checkpoint_dir: `str`, base directory for the checkpoint files.
+      save_secs: `int`, save every N secs.
+      save_steps: `int`, save every N steps.
+      saver: `Saver` object, used for saving.
+      checkpoint_basename: `str`, base name for the checkpoint files.
+      scaffold: `Scaffold`, use to get saver object.
+      listeners: List of `CheckpointSaverListener` subclass instances. Used for
+        callbacks that run immediately before or after this hook saves the
+        checkpoint.
+
+    Raises:
+      ValueError: One of `save_steps` or `save_secs` should be set.
+      ValueError: At most one of `saver` or `scaffold` should be set.
+    """
+    save_path = os.path.join(checkpoint_dir, checkpoint_basename)
+    logging.info("Create AsyncCheckpointSaverHook saving to path\n%s",
+                 save_path)
+    if listeners:
+      logging.info(" with %d listener(s).", len(listeners))
+    if saver is not None and scaffold is not None:
+      raise ValueError("You cannot provide both saver and scaffold.")
+    self._saver = saver
+    self._save_thread = None
+    self._write_graph_thread = None
+    self._checkpoint_dir = checkpoint_dir
+    self._save_path = save_path
+    self._scaffold = scaffold
+    self._timer = basic_session_run_hooks.SecondOrStepTimer(
+        every_secs=save_secs, every_steps=save_steps)
+    self._listeners = listeners or []
+    self._steps_per_run = 1
+    self._summary_writer = None
+    self._global_step_tensor = None
+
+    self._last_checkpoint_step = None
+
+    # Initialize the first timestamp for _END_TIME_OF_LAST_WRITE.
+    global _END_TIME_OF_LAST_WRITE
+    with _END_TIME_OF_LAST_WRITE_LOCK:
+      if _END_TIME_OF_LAST_WRITE is None:
+        _END_TIME_OF_LAST_WRITE = time.time()
+
+  def _set_steps_per_run(self, steps_per_run):
+    self._steps_per_run = steps_per_run
+
+  def begin(self):
+    self._summary_writer = SummaryWriterCache.get(self._checkpoint_dir)
+    self._global_step_tensor = training_util._get_or_create_global_step_read()  # pylint: disable=protected-access
+    if self._global_step_tensor is None:
+      raise RuntimeError(
+          "Global step should be created to use CheckpointSaverHook.")
+    for l in self._listeners:
+      l.begin()
+
+  def after_create_session(self, session: session_lib.Session, coord: Any):
+    global_step = session.run(self._global_step_tensor)
+
+    # We do write graph and saver_def at the first call of before_run.
+    # We cannot do this in begin, since we let other hooks to change graph and
+    # add variables in begin. Graph is finalized after all begin calls.
+    def _write_graph_fn(self):
+      training_util.write_graph(
+          ops.get_default_graph().as_graph_def(add_shapes=True),
+          self._checkpoint_dir, "graph.pbtxt")
+    self._write_graph_thread = threading.Thread(target=_write_graph_fn,
+                                                args=[self])
+    self._write_graph_thread.start()
+
+    saver_def = self._get_saver().saver_def if self._get_saver() else None
+    graph = ops.get_default_graph()
+    meta_graph_def = meta_graph.create_meta_graph_def(
+        graph_def=graph.as_graph_def(add_shapes=True), saver_def=saver_def)
+    if self._summary_writer is None:
+      raise ValueError("Summary writer is not initialised")
+    self._summary_writer.add_graph(graph)
+    self._summary_writer.add_meta_graph(meta_graph_def)
+    # The checkpoint saved here is the state at step "global_step".
+    self._save(session, global_step)
+    self._timer.update_last_triggered_step(global_step)
+
+  def before_run(self, run_context: Any):  # pylint: disable=unused-argument
+    return session_run_hook.SessionRunArgs(self._global_step_tensor)
+
+  def after_run(self, run_context: session_run_hook.SessionRunContext,
+                run_values: Any):
+    global_step = run_context.session.run(self._global_step_tensor)
+    if self._timer.should_trigger_for_step(global_step):
+      self._timer.update_last_triggered_step(global_step)
+      logging.info("Triggering checkpoint. %s", global_step)
+      if self._save(run_context.session, global_step):
+        run_context.request_stop()
+
+  def end(self, session: session_lib.Session):
+    if self._save_thread:
+      logging.info("Waiting for any pending checkpoints to finish.")
+      self._save_thread.join()
+    if self._write_graph_thread:
+      logging.info("Waiting for any pending write_graph to finish.")
+      self._write_graph_thread.join()
+
+    last_step = session.run(self._global_step_tensor)
+
+    if self._last_checkpoint_step != last_step:
+      self._save(session, last_step, asynchronous=False)
+
+    for l in self._listeners:
+      l.end(session, last_step)
+
+  def _save(self, session, step, asynchronous=True):
+    """Saves the latest checkpoint, returns should_stop."""
+
+    def _save_fn():
+      """Run the saver process."""
+      logging.info("Saving checkpoints for %d into %s.", step, self._save_path)
+
+      start_time = time.time()
+      for l in self._listeners:
+        l.before_save(session, step)
+
+      self._get_saver().save(session, self._save_path, global_step=step)
+      if self._summary_writer is None:
+        raise ValueError("Summary writer is not initialised")
+      self._summary_writer.add_session_log(
+          event_pb2.SessionLog(
+              status=event_pb2.SessionLog.CHECKPOINT,
+              checkpoint_path=self._save_path), step)
+
+      for l in self._listeners:
+        l.after_save(session, step)
+
+      # Measure the async checkpoint write duration, i.e., non-blocking time.
+      end_time = time.time()
+      metrics.AddAsyncCheckpointWriteDuration(
+          api_label=_ASYNC_CHECKPOINT_V1,
+          microseconds=_get_duration_microseconds(start_time, end_time))
+
+      # Measure the elapsed time since the last checkpoint.
+      # Due to the nature of async checkpoint, here it actually captures the
+      # duration between the start_time of the previous checkpoint and the start
+      # time of this checkpoint. As a result, the duration of the final async
+      # checkpoint is excluded, which is fine since it does not take much time.
+      global _END_TIME_OF_LAST_WRITE
+      with _END_TIME_OF_LAST_WRITE_LOCK:
+        metrics.AddTrainingTimeSaved(
+            api_label=_ASYNC_CHECKPOINT_V1,
+            microseconds=_get_duration_microseconds(_END_TIME_OF_LAST_WRITE,
+                                                    start_time))
+      _END_TIME_OF_LAST_WRITE = start_time
+
+      logging.info("Checkpoint actual writing time: (%.3f sec)",
+                   end_time - start_time)
+      logging.info("Checkpoint finished for %d into %s.", step, self._save_path)
+
+    # Measure the checkpoint write duration that is blocking the main thread.
+    blocking_start_time = time.time()
+    def end_of_blocking_time():
+      blocking_end_time = time.time()
+      metrics.AddCheckpointWriteDuration(
+          api_label=_ASYNC_CHECKPOINT_V1,
+          microseconds=_get_duration_microseconds(blocking_start_time,
+                                                  blocking_end_time))
+
+    if not asynchronous:
+      self._last_checkpoint_step = step
+      _save_fn()
+      end_of_blocking_time()
+      return
+
+    if self._save_thread is not None:
+      self._save_thread.join(timeout=0.1)
+      if self._save_thread.is_alive():
+        logging.info("Saver thread still in progress, skipping checkpoint.")
+        end_of_blocking_time()
+        return
+
+    self._last_checkpoint_step = step
+    self._save_thread = threading.Thread(target=_save_fn)
+    self._save_thread.start()
+    end_of_blocking_time()
+
+  def _get_saver(self):
+    if self._saver is not None:
+      return self._saver
+    elif self._scaffold is not None:
+      return self._scaffold.saver
+
+    # Get saver from the SAVERS collection if present.
+    collection_key = ops.GraphKeys.SAVERS
+    savers = ops.get_collection(collection_key)
+    if not savers:
+      raise RuntimeError(
+          "No items in collection {}. Please add a saver to the collection "
+          "or provide a saver or scaffold.".format(collection_key))
+    elif len(savers) > 1:
+      raise RuntimeError(
+          "More than one item in collection {}. "
+          "Please indicate which one to use by passing it to the constructor."
+          .format(collection_key))
+
+    self._saver = savers[0]
+    return savers[0]

tensorflow/python/tpu/async_checkpoint_test.py

@@ -0,0 +1,227 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Test async checkpointing."""
+
+import os
+
+import numpy as np
+
+from tensorflow.core.framework import summary_pb2
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.lib.io import file_io
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import metrics as metrics_lib
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops.losses import losses
+from tensorflow.python.platform import flags
+from tensorflow.python.platform import test
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.saved_model.pywrap_saved_model import metrics
+from tensorflow.python.tpu import async_checkpoint
+from tensorflow.python.tpu import tpu_optimizer
+from tensorflow.python.training import basic_session_run_hooks
+from tensorflow.python.training import training
+from tensorflow_estimator.python.estimator import estimator as estimator_lib
+from tensorflow_estimator.python.estimator import model_fn as model_fn_lib
+from tensorflow_estimator.python.estimator.tpu import tpu_config
+from tensorflow_estimator.python.estimator.tpu import tpu_estimator
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string('tpu', '', 'TPU to use in this test.')
+flags.DEFINE_string('zone', None, 'Name of GCP zone with TPU.')
+flags.DEFINE_string('project', None, 'Name of GCP project with TPU.')
+flags.DEFINE_string(
+    'model_dir',
+    os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR'),
+    'GCS path to store model and checkpoints.')
+
+
+def _get_checkpoint_metrics_counts() -> (int, int):
+  """Get the count for recorded sync and async checkpoint write durations."""
+  def get_count(method):
+    proto_bytes = method(api_label=async_checkpoint._ASYNC_CHECKPOINT_V1)
+    histogram_proto = summary_pb2.HistogramProto()
+    histogram_proto.ParseFromString(proto_bytes)
+    return int(histogram_proto.num)
+  return get_count(metrics.GetCheckpointWriteDurations), get_count(
+      metrics.GetAsyncCheckpointWriteDurations)
+
+
+def input_fn(params):
+  """Return a dataset of source and target sequences for training."""
+  return (constant_op.constant(
+      np.random.randn(params['batch_size'], 1000), dtype=dtypes.float32),
+          constant_op.constant(
+              np.random.randint(0, 10, params['batch_size']),
+              dtype=dtypes.int32))
+
+
+def model_fn(features, labels, mode, params):
+  del params  # unused
+  with variable_scope.variable_scope('m', reuse=variable_scope.AUTO_REUSE):
+    w = variable_scope.get_variable('W', shape=[1000, 10])
+  logits = math_ops.matmul(features, w)
+  loss = losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
+
+  if mode == model_fn_lib.ModeKeys.TRAIN:
+    optimizer = training.RMSPropOptimizer(learning_rate=0.01)
+    optimizer = tpu_optimizer.CrossShardOptimizer(optimizer)
+    train_op = optimizer.minimize(loss, training.get_global_step())
+    return tpu_estimator.TPUEstimatorSpec(
+        mode=model_fn_lib.ModeKeys.TRAIN,
+        loss=loss,
+        train_op=train_op,
+    )
+  elif mode == model_fn_lib.ModeKeys.EVAL:
+
+    def metric_fn(labels, logits):
+      labels = math_ops.cast(labels, dtypes.int64)
+      logging.info('LABELS %s %s', labels, logits)
+      return {
+          'recall@1': metrics_lib.recall_at_k(labels, logits, 1),
+          'recall@5': metrics_lib.recall_at_k(labels, logits, 5),
+      }
+
+    loss = losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
+    eval_metrics = (metric_fn, [labels, logits])
+    return tpu_estimator.TPUEstimatorSpec(
+        mode=model_fn_lib.ModeKeys.EVAL, loss=loss, eval_metrics=eval_metrics)
+
+
+class AsyncCheckpointingTest(test.TestCase):
+
+  def testAsyncCheckpointHookEnabled(self):
+    resolver = tpu_cluster_resolver.TPUClusterResolver(
+        tpu=FLAGS.tpu, zone=FLAGS.zone, project=FLAGS.project)
+
+    checkpoint_interval = 5
+    config = tpu_config.RunConfig(
+        master=resolver.master(),
+        model_dir=os.path.join(FLAGS.model_dir, 'runconfig'),
+        save_checkpoints_steps=1000,
+        keep_checkpoint_max=11,  # off by one
+        tpu_config=tpu_config.TPUConfig(
+            iterations_per_loop=checkpoint_interval,))
+
+    estimator = tpu_estimator.TPUEstimator(
+        use_tpu=True,
+        model_fn=model_fn,
+        config=config,
+        train_batch_size=32,
+        eval_batch_size=32,
+        predict_batch_size=1,
+        params={},
+    )
+
+    max_steps = 100
+    mock_listener = test.mock.create_autospec(
+        basic_session_run_hooks.CheckpointSaverListener)
+    estimator.train(
+        input_fn=input_fn,
+        max_steps=max_steps,
+        hooks=[
+            async_checkpoint.AsyncCheckpointSaverHook(
+                FLAGS.model_dir,
+                save_steps=checkpoint_interval,
+                listeners=[mock_listener])
+        ])
+
+    current_step = estimator_lib._load_global_step_from_checkpoint_dir(
+        FLAGS.model_dir)  # pylint: disable=protected-access
+
+    # TODO(power) -- identify a better way to count the number of checkpoints.
+    checkpoints = file_io.get_matching_files(
+        FLAGS.model_dir + '/model.ckpt*.meta')
+    checkpoint_count = len(checkpoints)
+    logging.info('Found %d checkpoints: %s', checkpoint_count, checkpoints)
+    self.assertLessEqual(checkpoint_count, 10)
+    self.assertEqual(current_step, max_steps)
+    mock_listener.before_save.assert_called()
+    mock_listener.after_save.assert_called()
+
+    # save called by hook in `after_create_session` and every `after_run`
+    num_save_calls = 1 + max_steps // checkpoint_interval
+    sync_count_1, async_count_1 = _get_checkpoint_metrics_counts()
+    # save might be called one extra time in `end` hook based on timing of
+    # `_last_checkpoint_step` update in the final `after_run` call
+    self.assertIn(sync_count_1, [num_save_calls, num_save_calls + 1])
+    self.assertLessEqual(async_count_1, num_save_calls)
+    training_time_saved = metrics.GetTrainingTimeSaved(
+        api_label=async_checkpoint._ASYNC_CHECKPOINT_V1)
+    self.assertGreater(training_time_saved, 0)
+
+  def testAsyncCheckpointHookWithoutListeners(self):
+    resolver = tpu_cluster_resolver.TPUClusterResolver(
+        tpu=FLAGS.tpu, zone=FLAGS.zone, project=FLAGS.project)
+
+    checkpoint_interval = 5
+    keep_checkpoint_max = 10
+    config = tpu_config.RunConfig(
+        master=resolver.master(),
+        model_dir=os.path.join(FLAGS.model_dir, 'runconfig'),
+        save_checkpoints_steps=1000,
+        keep_checkpoint_max=keep_checkpoint_max+1,  # off by one
+        tpu_config=tpu_config.TPUConfig(
+            iterations_per_loop=checkpoint_interval,))
+
+    estimator = tpu_estimator.TPUEstimator(
+        use_tpu=True,
+        model_fn=model_fn,
+        config=config,
+        train_batch_size=32,
+        eval_batch_size=32,
+        predict_batch_size=1,
+        params={},
+    )
+
+    max_steps = 100
+    estimator.train(
+        input_fn=input_fn,
+        max_steps=max_steps,
+        hooks=[
+            async_checkpoint.AsyncCheckpointSaverHook(
+                FLAGS.model_dir,
+                save_steps=checkpoint_interval)
+        ])
+
+    current_step = estimator_lib._load_global_step_from_checkpoint_dir(
+        FLAGS.model_dir)  # pylint: disable=protected-access
+
+    # TODO(power) -- identify a better way to count the number of checkpoints.
+    checkpoints = file_io.get_matching_files(
+        FLAGS.model_dir + '/model.ckpt*.meta')
+    checkpoint_count = len(checkpoints)
+    logging.info('Found %d checkpoints: %s', checkpoint_count, checkpoints)
+    self.assertLessEqual(checkpoint_count, keep_checkpoint_max)
+    self.assertEqual(current_step, max_steps)
+
+    # save called by hook in `after_create_session` and every `after_run`
+    num_save_calls = 1 + max_steps // checkpoint_interval
+    sync_count_1, async_count_1 = _get_checkpoint_metrics_counts()
+    # save might be called one extra time in `end` hook based on timing of
+    # `_last_checkpoint_step` update in the final `after_run` call
+    self.assertIn(sync_count_1, [num_save_calls, num_save_calls + 1])
+    self.assertLessEqual(async_count_1, num_save_calls)
+    training_time_saved = metrics.GetTrainingTimeSaved(
+        api_label=async_checkpoint._ASYNC_CHECKPOINT_V1)
+    self.assertGreater(training_time_saved, 0)
+
+
+if __name__ == '__main__':
+  v2_compat.disable_v2_behavior()
+  test.main()

tensorflow/python/tpu/bfloat16.py

@@ -0,0 +1,88 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Helper context for running models with bfloat16."""
+
+from typing import Generator, Optional, Text
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.util import tf_contextlib
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _get_custom_getter():
+  """Returns a custom getter that this class's methods must be called under.
+
+  All methods of this class must be called under a variable scope that was
+  passed this custom getter. Example:
+
+  ```python
+  network = ConvNetBuilder(...)
+  with tf.compat.v1.variable_scope('cg',
+                                   custom_getter=network.get_custom_getter()):
+    network.conv(...)
+    # Call more methods of network here
+  ```
+
+  Currently, this custom getter only does anything if self.use_tf_layers is
+  True. In that case, it causes variables to be stored as dtype
+  self.variable_type, then casted to the requested dtype, instead of directly
+  storing the variable as the requested dtype.
+  """
+
+  def inner_custom_getter(getter, *args, **kwargs):
+    """Custom getter that forces variables to have type self.variable_type."""
+    cast_to_bfloat16 = False
+    requested_dtype = kwargs['dtype']
+    if requested_dtype == dtypes.bfloat16:
+      # Only change the variable dtype if doing so does not decrease variable
+      # precision.
+      kwargs['dtype'] = dtypes.float32
+      cast_to_bfloat16 = True
+    var = getter(*args, **kwargs)
+    # This if statement is needed to guard the cast, because batch norm
+    # assigns directly to the return value of this custom getter. The cast
+    # makes the return value not a variable so it cannot be assigned. Batch
+    # norm variables are always in fp32 so this if statement is never
+    # triggered for them.
+    if cast_to_bfloat16:
+      var = math_ops.cast(var, dtypes.bfloat16)
+    return var
+
+  return inner_custom_getter
+
+
+@tf_export(v1=['tpu.bfloat16_scope'])
+@tf_contextlib.contextmanager
+def bfloat16_scope(
+    name: Optional[Text] = None
+) -> Generator[variable_scope.variable_scope, None, None]:
+  """Scope class for bfloat16 variables so that the model uses custom getter.
+
+  This enables variables to be read as bfloat16 type when using get_variable.
+
+  Arguments:
+    name: Name to use for scope.
+
+  Yields:
+    a variable scope.
+  """
+  if name is None:
+    name = ''
+  with variable_scope.variable_scope(
+      name, custom_getter=_get_custom_getter()) as varscope:
+    yield varscope

tensorflow/python/tpu/bfloat16_test.py

@@ -0,0 +1,81 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Tests for bfloat16 helper."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import bfloat16
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+
+
+class BFloat16ScopeTest(test.TestCase):
+
+  def testDefaultScopeName(self):
+    """Test if name for the variable scope is propagated correctly."""
+    with bfloat16.bfloat16_scope() as bf:
+      self.assertEqual(bf.name, "")
+
+  def testCustomScopeName(self):
+    """Test if custom name for the variable scope is propagated correctly."""
+    name = 'bfloat16'
+    with bfloat16.bfloat16_scope('bfloat16') as bf:
+      self.assertEqual(bf.name, name)
+
+  def testVariableName(self):
+    """Test if custom name for the variable scope is propagated correctly."""
+    g = ops.Graph()
+    with g.as_default():
+      a = variables.Variable(2.2, name='var_a')
+      b = variables.Variable(3.3, name='var_b')
+      d = variables.Variable(4.4, name='var_b')
+      with g.name_scope('scope1'):
+        with bfloat16.bfloat16_scope('bf16'):
+          a = math_ops.cast(a, dtypes.bfloat16)
+          b = math_ops.cast(b, dtypes.bfloat16)
+          c = math_ops.add(a, b, name='addition')
+        with bfloat16.bfloat16_scope():
+          d = math_ops.cast(d, dtypes.bfloat16)
+          math_ops.add(c, d, name='addition')
+
+    g_ops = g.get_operations()
+    ops_name = []
+    for op in g_ops:
+      ops_name.append(str(op.name))
+
+    self.assertIn('scope1/bf16/addition', ops_name)
+    self.assertIn('scope1/bf16/Cast', ops_name)
+    self.assertIn('scope1/addition', ops_name)
+    self.assertIn('scope1/Cast', ops_name)
+
+  @test_util.run_deprecated_v1
+  def testRequestedDType(self):
+    """Test if requested dtype is honored in the getter.
+    """
+    with bfloat16.bfloat16_scope() as scope:
+      v1 = variable_scope.get_variable("v1", [])
+      self.assertEqual(v1.dtype.base_dtype, dtypes.float32)
+      v2 = variable_scope.get_variable("v2", [], dtype=dtypes.bfloat16)
+      self.assertEqual(v2.dtype.base_dtype, dtypes.bfloat16)
+      self.assertEqual([dtypes.float32, dtypes.float32],
+                       [v.dtype.base_dtype for v in scope.global_variables()])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/tpu/client/BUILD

@@ -0,0 +1,44 @@
+# Cloud TPU Client.
+
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:internal",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "client",
+    srcs = [
+        "client.py",
+        "version.py",
+    ],
+    strict_deps = True,
+    deps = ["@absl_py//absl/flags"],
+)
+
+py_library(
+    name = "client_lib",
+    srcs = [
+        "__init__.py",
+    ],
+    strict_deps = True,
+    deps = [":client"],
+)
+
+tf_py_strict_test(
+    name = "client_py_test",
+    size = "small",
+    srcs = ["client_test.py"],
+    grpc_enabled = True,
+    main = "client_test.py",
+    deps = [
+        ":client",
+        "//tensorflow/python/platform:client_testlib",
+        "@absl_py//absl/flags",
+    ],
+)

tensorflow/python/tpu/client/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Cloud TPU Client."""
+
+from .version import __version__
+
+from tensorflow.python.tpu.client.client import Client

tensorflow/python/tpu/client/client.py

@@ -0,0 +1,438 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Cloud TPU Client."""
+
+from concurrent import futures
+import datetime
+import json
+import logging
+import os
+import time
+import urllib
+
+from absl import flags
+
+_GOOGLE_API_CLIENT_INSTALLED = True
+try:
+  from googleapiclient import discovery  # pylint: disable=g-import-not-at-top
+  from oauth2client import client  # pylint: disable=g-import-not-at-top
+except ImportError:
+  _GOOGLE_API_CLIENT_INSTALLED = False
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_bool('runtime_oom_exit', True,
+                  'Exit the script when the TPU runtime is OOM.')
+flags.DEFINE_bool('hbm_oom_exit', True,
+                  'Exit the script when the TPU HBM is OOM.')
+
+_GKE_ENV_VARIABLE = 'KUBE_GOOGLE_CLOUD_TPU_ENDPOINTS'
+_DEFAULT_TPUCONFIG_VARIABLE = 'TPU_CONFIG'
+_ENDPOINTS_SEPARATOR = ','
+_DEFAULT_ENV_VARIABLE = 'TPU_NAME'
+_DISCOVERY_SERVICE_URL_ENV_VARIABLE = 'TPU_API_DISCOVERY_URL'
+_GCE_METADATA_URL_ENV_VARIABLE = 'GCE_METADATA_IP'
+_GCE_METADATA_ENDPOINT_ENV_VARIABLE = 'GCE_METADATA_HOST'
+_DEFAULT_ENDPOINT_PORT = '8470'
+_OOM_EVENT_COOL_TIME_SEC = 90
+_VERSION_SWITCHER_ENDPOINT = 'http://{}:8475/requestversion'
+
+
+def _utcnow():
+  """A wrapper function around datetime.datetime.utcnow.
+
+  This function is created for unit testing purpose. It's not easy to do
+  StubOutWithMock with datetime.datetime package.
+
+  Returns:
+    datetime.datetime
+  """
+  return datetime.datetime.utcnow()
+
+
+def _environment_discovery_url():
+  return os.environ.get(_DISCOVERY_SERVICE_URL_ENV_VARIABLE)
+
+
+def _gce_metadata_endpoint():
+  endpoint = os.environ.get(_GCE_METADATA_ENDPOINT_ENV_VARIABLE)
+  if not endpoint:
+    endpoint = os.environ.get(
+        _GCE_METADATA_URL_ENV_VARIABLE, 'metadata.google.internal'
+    )
+  return 'http://' + endpoint
+
+
+def _request_compute_metadata(path):
+  req = urllib.request.Request(
+      '%s/computeMetadata/v1/%s' % (_gce_metadata_endpoint(), path),
+      headers={'Metadata-Flavor': 'Google'})
+  resp = urllib.request.urlopen(req)
+  return _as_text(resp.read())
+
+
+def _environment_var_to_network_endpoints(endpoints):
+  """Yields a dict with ip address and port."""
+  for endpoint in endpoints.split(','):
+    grpc_prefix = 'grpc://'
+    if endpoint.startswith(grpc_prefix):
+      endpoint = endpoint.split(grpc_prefix)[1]
+    parts = endpoint.split(':')
+    ip_address = parts[0]
+    port = _DEFAULT_ENDPOINT_PORT
+    if len(parts) > 1:
+      port = parts[1]
+    yield {
+        'ipAddress': ip_address,
+        'port': port
+    }
+
+
+def _get_tpu_node_config():
+  tpu_config_env = os.environ.get(_DEFAULT_TPUCONFIG_VARIABLE)
+  if tpu_config_env:
+    return json.loads(tpu_config_env)
+  return None
+
+
+def _get_tpu_name(tpu):
+  if tpu:
+    return tpu
+
+  for e in [_GKE_ENV_VARIABLE, _DEFAULT_ENV_VARIABLE]:
+    if e in os.environ:
+      return os.environ[e]
+  return None
+
+
+def _as_text(s):
+  if isinstance(s, bytes):
+    return s.decode('utf-8')
+  return s
+
+
+class Client:
+  """Client for working with the Cloud TPU API.
+
+  This client is intended to be used for resolving tpu name to ip addresses.
+
+  It's recommended to use this library as a contextlib to utilize all
+  functionality.
+  """
+
+  def __init__(self,
+               tpu=None,
+               zone=None,
+               project=None,
+               credentials='default',
+               service=None,
+               discovery_url=None):
+    if isinstance(tpu, list):
+      if not tpu:
+        raise ValueError('At least one TPU must be specified.')
+      if len(tpu) != 1:
+        raise NotImplementedError(
+            'Using multiple TPUs in a single session is not yet implemented')
+      tpu = tpu[0]
+
+    tpu = _get_tpu_name(tpu)
+
+    if tpu is None:
+      tpu_node_config = _get_tpu_node_config()
+      if tpu_node_config:
+        tpu = tpu_node_config.get('tpu_node_name')
+        project = project or tpu_node_config.get('project')
+        zone = zone or tpu_node_config.get('zone')
+      else:
+        raise ValueError('Please provide a TPU Name to connect to.')
+
+    self._tpu = _as_text(tpu)
+
+    self._use_api = not self._tpu.startswith('grpc://')
+    self._service = service
+
+    self._credentials = None
+    self._project = None
+    self._zone = None
+    self._discovery_url = None
+    if self._use_api:
+      if credentials != 'default':
+        self._credentials = credentials
+      # Automatically detect project and zone if unspecified.
+      if project:
+        self._project = _as_text(project)
+      else:
+        self._project = _request_compute_metadata('project/project-id')
+      if zone:
+        self._zone = _as_text(zone)
+      else:
+        zone_path = _request_compute_metadata('instance/zone')
+        self._zone = zone_path.split('/')[-1]
+      self._discovery_url = _environment_discovery_url() or discovery_url
+
+  def _symptom_msg(self, msg):
+    """Return the structured Symptom message."""
+    return 'Symptom: ' + msg
+
+  def _oom_event(self, symptoms):
+    """Check if a runtime OOM event is reported."""
+    if not symptoms:
+      return False
+    for symptom in reversed(symptoms):
+      if symptom['symptomType'] != 'OUT_OF_MEMORY':
+        continue
+      oom_datetime_str = symptom['createTime'].split('.')[0]
+      oom_datetime = datetime.datetime.strptime(oom_datetime_str,
+                                                '%Y-%m-%dT%H:%M:%S')
+      time_diff = _utcnow() - oom_datetime
+      if time_diff < datetime.timedelta(seconds=_OOM_EVENT_COOL_TIME_SEC):
+        logging.warning(
+            self._symptom_msg(
+                'a recent runtime OOM has occurred ~{} seconds ago. The model '
+                'script will terminate automatically. To prevent future OOM '
+                'events, please consider reducing the model size. To disable this '
+                'behavior, set flag --runtime_oom_exit=false when starting the '
+                'script.'.format(time_diff.seconds)))
+        return True
+    return False
+
+  def _hbm_oom_event(self, symptoms):
+    """Check if a HBM OOM event is reported."""
+    if not symptoms:
+      return False
+    for symptom in reversed(symptoms):
+      if symptom['symptomType'] != 'HBM_OUT_OF_MEMORY':
+        continue
+      oom_datetime_str = symptom['createTime'].split('.')[0]
+      oom_datetime = datetime.datetime.strptime(oom_datetime_str,
+                                                '%Y-%m-%dT%H:%M:%S')
+      time_diff = _utcnow() - oom_datetime
+      if time_diff < datetime.timedelta(seconds=_OOM_EVENT_COOL_TIME_SEC):
+        logging.warning(
+            self._symptom_msg(
+                'a recent HBM OOM has occurred ~{} seconds ago. The model '
+                'script will terminate automatically. To prevent future HBM OOM '
+                'events, please consider reducing the model size. To disable this '
+                'behavior, set flag --hbm_oom_exit=false when starting the '
+                'script.'.format(time_diff.seconds)))
+        return True
+    return False
+
+  def _tpu_service(self):
+    """Creates a new Cloud TPU API object.
+
+    This works around an issue where the underlying HTTP connection sometimes
+    times out when the script has been running for too long. Other methods in
+    this object call this method to get a new API object whenever they need
+    to communicate with the Cloud API.
+
+    Raises:
+      RuntimeError: If the dependent Python packages are missing.
+
+    Returns:
+      A Google Cloud TPU API object.
+    """
+    if self._service:
+      return self._service
+
+    if not _GOOGLE_API_CLIENT_INSTALLED:
+      raise RuntimeError('Missing runtime dependency on the Google API client. '
+                         'Run `pip install cloud-tpu-client` to fix.')
+
+    credentials = self._credentials
+    if credentials is None or credentials == 'default':
+      credentials = client.GoogleCredentials.get_application_default()
+
+    if self._discovery_url:
+      return discovery.build(
+          'tpu',
+          'v1',
+          credentials=credentials,
+          discoveryServiceUrl=self._discovery_url,
+          cache_discovery=False)
+    else:
+      return discovery.build(
+          'tpu', 'v1', credentials=credentials, cache_discovery=False)
+
+  def _full_name(self):
+    """Returns the full Cloud name for this TPU."""
+    return 'projects/%s/locations/%s/nodes/%s' % (
+        self._project, self._zone, self._tpu)
+
+  def _fetch_cloud_tpu_metadata(self):
+    """Returns the TPU metadata object from the TPU Get API call."""
+    service = self._tpu_service()
+    try:
+      r = service.projects().locations().nodes().get(name=self._full_name())
+      return r.execute()
+    except Exception as e:
+      raise ValueError("Could not lookup TPU metadata from name '%s'. Please "
+                       'doublecheck the tpu argument in the TPUClusterResolver '
+                       'constructor. Exception: %s' % (self._tpu, e))
+
+  def _get_tpu_property(self, key):
+    if self._use_api:
+      metadata = self._fetch_cloud_tpu_metadata()
+      return metadata.get(key)
+
+    return None
+
+  def __enter__(self):
+    self._open = True
+
+  def __exit__(self, type, value, traceback):  # pylint: disable=redefined-builtin
+    del type, value, traceback
+
+  def recoverable(self):
+    """Returns true if the TPU is in a state where training should eventually resume.
+
+    If false the TPU is in a unrecoverable state and should be recreated.
+    """
+    state = self.state()
+    symptoms = self.symptoms()
+    if state and state in ['TERMINATED', 'PREEMPTED']:
+      return False
+    elif FLAGS.runtime_oom_exit and self._oom_event(symptoms):
+      return False
+    elif FLAGS.hbm_oom_exit and self._hbm_oom_event(symptoms):
+      return False
+    return True
+
+  def symptoms(self):
+    """Return Cloud TPU Symptoms of the TPU."""
+    return self._get_tpu_property('symptoms')
+
+  def state(self):
+    """Return state of the TPU."""
+    return self._get_tpu_property('state')
+
+  def health(self):
+    """Return health of the TPU."""
+    return self._get_tpu_property('health')
+
+  def runtime_version(self):
+    """Return runtime version of the TPU."""
+
+    if not self._use_api:
+      # Fallback on getting version directly from TPU.
+      url = _VERSION_SWITCHER_ENDPOINT.format(
+          self.network_endpoints()[0]['ipAddress'])
+      try:
+        req = urllib.request.Request(url)
+        resp = urllib.request.urlopen(req)
+        version_details = json.loads(resp.read())
+        return version_details.get('currentVersion')
+      except urllib.error.HTTPError as e:
+        status_code = e.code
+        if status_code == 404:
+          return None
+        else:
+          raise e
+    return self._get_tpu_property('tensorflowVersion')
+
+  def accelerator_type(self):
+    """Return accelerator type of the TPU."""
+    return self._get_tpu_property('acceleratorType')
+
+  def api_available(self):
+    """Return if the Cloud TPU API is available, if not certain features will not work."""
+    return self._use_api
+
+  def name(self):
+    """Return the name of the tpu, or the ip address if name is not provided."""
+    return self._tpu
+
+  def get_local_ip(self):
+    """Return the local ip address of the Google Cloud VM the workload is running on."""
+    return _request_compute_metadata('instance/network-interfaces/0/ip')
+
+  def network_endpoints(self):
+    """Return a list of tpu endpoints."""
+    if not self._use_api:
+      return list(_environment_var_to_network_endpoints(self._tpu))
+    response = self._fetch_cloud_tpu_metadata()
+
+    if response.get('state') != 'READY':
+      raise RuntimeError('TPU "%s" is not yet ready; state: "%s"' %
+                         (self._tpu, response.get('state')))
+    if 'networkEndpoints' in response:
+      return response['networkEndpoints']
+    else:
+      return [{'ipAddress': response['ipAddress'], 'port': response['port']}]
+
+  def wait_for_healthy(self, timeout_s=1200, interval=30):
+    """Wait for TPU to become healthy or raise error if timeout reached.
+
+    Args:
+      timeout_s (int): The timeout in seconds for waiting TPU to become healthy.
+      interval (int): The interval in seconds to poll the TPU for health.
+
+    Raises:
+      RuntimeError: If the TPU doesn't become healthy by the timeout.
+    """
+    timeout = time.time() + timeout_s
+    while self.health() != 'HEALTHY':
+      logging.warning(
+          ('Waiting for TPU "%s" with state "%s" '
+           'and health "%s" to become healthy'),
+          self.name(), self.state(), self.health())
+      if time.time() + interval > timeout:
+        raise RuntimeError(
+            'Timed out waiting for TPU "%s" to become healthy' % self.name())
+      time.sleep(interval)
+
+    logging.warning('TPU "%s" is healthy.', self.name())
+
+  def configure_tpu_version(self, version, restart_type='always'):
+    """Configure TPU software version.
+
+    Args:
+      version (string): Version of software to configure the TPU with.
+      restart_type (string): Restart behaviour when switching versions,
+        defaults to always restart. Options are 'always', 'ifNeeded'.
+
+    """
+
+    def configure_worker(worker):
+      """Configure individual TPU worker.
+
+      Args:
+        worker: A dict with the field ipAddress where the configure request will
+          be sent.
+      """
+      ip_address = worker['ipAddress']
+      url = (_VERSION_SWITCHER_ENDPOINT + '/{}?restartType={}').format(
+          ip_address, version, restart_type)
+      req = urllib.request.Request(url, data=b'')
+      try:
+        urllib.request.urlopen(req)
+      except urllib.error.HTTPError as e:
+        status_code = e.code
+        if status_code == 404:
+          raise Exception(
+              'Tensorflow version {} is not available on Cloud TPU, '
+              'try a previous nightly version or refer to '
+              'https://cloud.google.com/tpu/docs/release-notes for '
+              'the latest official version.'.format(version))
+        else:
+          raise Exception('Failed to configure worker {}'.format(ip_address))
+
+    workers = self.network_endpoints()
+
+    with futures.ThreadPoolExecutor(max_workers=len(workers)) as executor:
+      results = executor.map(configure_worker, workers)
+      for result in results:
+        if result:
+          result.result()

tensorflow/python/tpu/client/client_test.py

@@ -0,0 +1,863 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for cloud tpu client."""
+
+import datetime
+import json
+import os
+import time
+import urllib
+
+from absl import flags
+
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.client import client
+
+FLAGS = flags.FLAGS
+
+mock = test.mock
+
+_UTCNOW_STR = '2000-01-01T00:30:00'
+
+
+def mock_utcnow():
+  return datetime.datetime.strptime(_UTCNOW_STR, '%Y-%m-%dT%H:%M:%S')
+
+
+def mock_request_compute_metadata(path):
+  if path == 'project/project-id':
+    return 'test-project'
+  elif path == 'instance/zone':
+    return 'projects/test-project/locations/us-central1-c'
+  elif path == 'instance/network-interfaces/0/ip':
+    return '10.128.1.2'
+  return ''
+
+
+class MockRequestClass:
+
+  def __init__(self, name, tpu_map):
+    self._name = name
+    self._tpu_map = tpu_map
+
+  def execute(self):
+    if self._name in self._tpu_map:
+      tpu_dict = self._tpu_map[self._name].copy()
+      if isinstance(tpu_dict.get('health'), list):
+        # Do extraction of health list to a single health string based on time.
+        time_now = time.time()
+        health_now = tpu_dict.get('health')[time_now]
+        tpu_dict['health'] = health_now
+      return tpu_dict
+    else:
+      raise KeyError('Resource %s was not found' % self._name)
+
+
+class MockNodeClass:
+
+  def __init__(self, tpu_map):
+    self._tpu_map = tpu_map
+
+  def get(self, name):
+    return MockRequestClass(name, self._tpu_map)
+
+
+class CloudTpuClientTest(test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    if 'TPU_API_DISCOVERY_URL' in os.environ:
+      del os.environ['TPU_API_DISCOVERY_URL']
+    if 'TPU_NAME' in os.environ:
+      del os.environ['TPU_NAME']
+    self._time_now = 0
+    self.addCleanup(mock.patch.stopall)
+
+  def _mock_time(self, *args, **kwargs):
+    return self._time_now
+
+  def _mock_sleep(self, secs):
+    self._time_now += secs
+
+  def mock_service_client(self, tpu_map=None):
+    if tpu_map is None:
+      tpu_map = {}
+
+    mock_locations = mock.MagicMock()
+    mock_locations.nodes.return_value = MockNodeClass(tpu_map)
+
+    mock_project = mock.MagicMock()
+    mock_project.locations.return_value = mock_locations
+
+    mock_client = mock.MagicMock()
+    mock_client.projects.return_value = mock_project
+    return mock_client
+
+  def testEnvironmentDiscoveryUrl(self):
+    os.environ['TPU_API_DISCOVERY_URL'] = 'https://{api}.internal/{apiVersion}'
+    self.assertEqual('https://{api}.internal/{apiVersion}',
+                     (client._environment_discovery_url()))
+
+  def testEnvironmentGCEDefault(self):
+    self.assertEqual(
+        'http://metadata.google.internal', client._gce_metadata_endpoint()
+    )
+
+  @mock.patch.dict(os.environ, {'GCE_METADATA_IP': '1.2.3.4'})
+  def testEnvironmentGCEIPOverride(self):
+    self.assertEqual('http://1.2.3.4', client._gce_metadata_endpoint())
+
+  @mock.patch.dict(os.environ, {'GCE_METADATA_HOST': 'foo.bar'})
+  def testEnvironmentGCEHostOverride(self):
+    self.assertEqual('http://foo.bar', client._gce_metadata_endpoint())
+
+  def testEnvironmentVarToNetworkEndpointsSingleIp(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '1234'}],
+        list(client._environment_var_to_network_endpoints(
+            '1.2.3.4:1234')))
+
+  def testEnvironmentVarToNetworkEndpointsSingleGrpcAddress(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '2000'}],
+        list(
+            client._environment_var_to_network_endpoints(
+                'grpc://1.2.3.4:2000')))
+
+  def testEnvironmentVarToNetworkEndpointsMultipleIps(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '2000'},
+         {'ipAddress': '5.6.7.8', 'port': '1234'}],
+        list(
+            client._environment_var_to_network_endpoints(
+                '1.2.3.4:2000,5.6.7.8:1234')))
+
+  def testEnvironmentVarToNetworkEndpointsMultipleGrpcAddresses(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '2000'},
+         {'ipAddress': '5.6.7.8', 'port': '1234'}],
+        list(client._environment_var_to_network_endpoints(
+            'grpc://1.2.3.4:2000,grpc://5.6.7.8:1234')))
+
+  def testEnvironmentVarToNetworkEndpointsMissingPortAndMixed(self):
+    self.assertEqual(
+        [{'ipAddress': '1.2.3.4', 'port': '2000'},
+         {'ipAddress': '5.6.7.8', 'port': '8470'}],
+        list(client._environment_var_to_network_endpoints(
+            '1.2.3.4:2000,grpc://5.6.7.8')))
+
+  def testInitializeNoArguments(self):
+    with self.assertRaisesRegex(
+        ValueError, 'Please provide a TPU Name to connect to.'):
+      client.Client()
+
+  def testInitializeMultiElementTpuArray(self):
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        'Using multiple TPUs in a single session is not yet implemented'):
+      client.Client(tpu=['multiple', 'elements'])
+
+  def assertClientContains(self, c):
+    self.assertEqual('tpu_name', c._tpu)
+    self.assertEqual(True, c._use_api)
+    self.assertIsNone(c._credentials)
+    self.assertEqual('test-project', c._project)
+    self.assertEqual('us-central1-c', c._zone)
+    self.assertIsNone(c._discovery_url)
+    self.assertEqual([{
+        'ipAddress': '10.1.2.3',
+        'port': '8470'
+    }], c.network_endpoints())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testNetworkEndpointsNotReadyWithApi(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertRaisesRegex(
+        RuntimeError, 'TPU .* is not yet ready; state: "None"',
+        c.network_endpoints)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testInitializeNoArgumentsWithEnvironmentVariable(self):
+    os.environ['TPU_NAME'] = 'tpu_name'
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY',
+        }
+    }
+    c = client.Client(
+        service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertClientContains(c)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testInitializeNoArgumentsWithTPUEnvironmentVariableTPUConfig(self):
+    os.environ['TPU_CONFIG'] = json.dumps({
+        'project': 'test-project',
+        'zone': 'us-central1-c',
+        'tpu_node_name': 'tpu_name',
+    })
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY',
+        }
+    }
+    c = client.Client(service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertClientContains(c)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testInitializeTpuName(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': 'HEALTHY',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertClientContains(c)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testInitializeIpAddress(self):
+    c = client.Client(tpu='grpc://1.2.3.4:8470')
+    self.assertEqual('grpc://1.2.3.4:8470', c._tpu)
+    self.assertEqual(False, c._use_api)
+    self.assertIsNone(c._service)
+    self.assertIsNone(c._credentials)
+    self.assertIsNone(c._project)
+    self.assertIsNone(c._zone)
+    self.assertIsNone(c._discovery_url)
+    self.assertEqual([{
+        'ipAddress': '1.2.3.4',
+        'port': '8470'
+    }], c.network_endpoints())
+
+  def testInitializeWithoutMetadata(self):
+    c = client.Client(
+        tpu='tpu_name', project='project', zone='zone')
+    self.assertEqual('tpu_name', c._tpu)
+    self.assertEqual(True, c._use_api)
+    self.assertIsNone(c._service)
+    self.assertIsNone(c._credentials)
+    self.assertEqual('project', c._project)
+    self.assertEqual('zone', c._zone)
+    self.assertIsNone(c._discovery_url)
+
+  def testRecoverableNoApiAccess(self):
+    c = client.Client(tpu='grpc://1.2.3.4:8470')
+    self.assertEqual(True, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRecoverableNoState(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual(True, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRecoverableReady(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual(True, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRecoverablePreempted(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'PREEMPTED',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual(False, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableOOM(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, False),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:20.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, False),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:20.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:00.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:10.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:20.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True)
+    ]
+
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='tpu_name',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableOOMDisabled(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+    ]
+
+    FLAGS.runtime_oom_exit = False
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='tpu_name',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+    FLAGS.runtime_oom_exit = True
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableOOMNoAPI(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'OUT_OF_MEMORY',
+                    'details': 'The TPU runtime has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+    ]
+
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='grpc://1.2.3.4:8470',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableHBMOOM(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, False),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:20.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, False),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:28:20.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:00.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:10.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:20.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }, {
+                    'createTime': '2000-01-01T00:29:40.123456Z',
+                    'symptomType': 'LOW_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True)
+    ]
+
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='tpu_name',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableHBMOOMDisabled(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+    ]
+
+    FLAGS.hbm_oom_exit = False
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='tpu_name',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+    FLAGS.hbm_oom_exit = True
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  @mock.patch.object(client, '_utcnow', mock_utcnow)
+  def testRecoverableHBMOOMNoAPI(self):
+    test_cases = [
+        ({
+            'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+                'state':
+                    'READY',
+                'symptoms': [{
+                    'createTime': '2000-01-01T00:29:30.123456Z',
+                    'symptomType': 'HBM_OUT_OF_MEMORY',
+                    'details': 'The TPU HBM has run OOM at timestamp '
+                               '2020-05-29T04:51:32.038721+00:00',
+                    'workerId': '0'
+                }]
+            }
+        }, True),
+    ]
+
+    for tpu_map, want in test_cases:
+      c = client.Client(tpu='grpc://1.2.3.4:8470',
+                        service=self.mock_service_client(tpu_map=tpu_map))
+      self.assertEqual(want, c.recoverable())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testHealthApi(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'PREEMPTED',
+            'health': 'HEALTHY',
+            'acceleratorType': 'v3-8',
+            'tensorflowVersion': 'nightly',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual('HEALTHY', c.health())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testRuntimeVersionApi(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'PREEMPTED',
+            'health': 'HEALTHY',
+            'acceleratorType': 'v3-8',
+            'tensorflowVersion': 'nightly',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual('nightly', c.runtime_version())
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testAcceleratorTypeApi(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'PREEMPTED',
+            'health': 'HEALTHY',
+            'acceleratorType': 'v3-8',
+            'tensorflowVersion': 'nightly',
+        }
+    }
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+    self.assertEqual('v3-8', c.accelerator_type())
+
+  def testHandlesByteStrings(self):
+    self.assertEqual(
+        client.Client(
+            tpu='tpu_name', zone='zone', project='project')._full_name(),
+        client.Client(
+            tpu=b'tpu_name', zone=b'zone', project=b'project')._full_name(),
+    )
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testWaitForHealthy(self):
+    time_mock = mock.patch.object(time, 'time', autospec=True).start()
+    time_mock.side_effect = self._mock_time
+    sleep_mock = mock.patch.object(time, 'sleep', autospec=True).start()
+    sleep_mock.side_effect = self._mock_sleep
+
+    health_timeseries = (['UNHEALTHY_MAINTENANCE']*30 + ['TIMEOUT']*10
+                         + [None]*20 + ['HEALTHY']*30)
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': health_timeseries,
+        },
+    }
+
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+
+    # Doesn't throw RuntimeError as TPU becomes HEALTHY before timeout
+    timeout = 80
+    interval = 5
+    return_time = 60
+    c.wait_for_healthy(timeout_s=timeout, interval=interval)
+    self.assertEqual(time.time(), return_time)
+    self.assertEqual(sleep_mock.call_count, return_time/interval)
+
+  @mock.patch.object(client, '_request_compute_metadata',
+                     mock_request_compute_metadata)
+  def testWaitForHealthyRaisesError(self):
+    time_mock = mock.patch.object(time, 'time', autospec=True).start()
+    time_mock.side_effect = self._mock_time
+    sleep_mock = mock.patch.object(time, 'sleep', autospec=True).start()
+    sleep_mock.side_effect = self._mock_sleep
+
+    # Mock timeseries where takes longer than timeout.
+    health_timeseries = ['UNHEALTHY_MAINTENANCE']*50 + ['TIMEOUT']*50
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'ipAddress': '10.1.2.3',
+            'port': '8470',
+            'state': 'READY',
+            'health': health_timeseries,
+        },
+    }
+
+    c = client.Client(
+        tpu='tpu_name', service=self.mock_service_client(tpu_map=tpu_map))
+
+    # Doesn't throw RuntimeError as TPU becomes HEALTHY before timeout
+    with self.assertRaisesRegex(
+        RuntimeError,
+        'Timed out waiting for TPU .* to become healthy'):
+      c.wait_for_healthy(timeout_s=80, interval=5)
+
+  def baseConfigureTpuVersion(self):
+    tpu_map = {
+        'projects/test-project/locations/us-central1-c/nodes/tpu_name': {
+            'state':
+                'READY',
+            'networkEndpoints': [
+                {
+                    'ipAddress': '1.2.3.4'
+                },
+                {
+                    'ipAddress': '5.6.7.8'
+                },
+            ]
+        }
+    }
+    return client.Client(
+        tpu='tpu_name',
+        project='test-project',
+        zone='us-central1-c',
+        service=self.mock_service_client(tpu_map=tpu_map))
+
+  @mock.patch.object(urllib.request, 'urlopen')
+  def testConfigureTpuVersion(self, urlopen):
+    c = self.baseConfigureTpuVersion()
+    c.configure_tpu_version('1.15')
+    paths = [call[0][0].full_url for call in urlopen.call_args_list]
+    self.assertCountEqual([
+        'http://1.2.3.4:8475/requestversion/1.15?restartType=always',
+        'http://5.6.7.8:8475/requestversion/1.15?restartType=always'
+    ], sorted(paths))
+
+  @mock.patch.object(urllib.request, 'urlopen')
+  def testConfigureTpuVersionRestartIfneeded(self, urlopen):
+    c = self.baseConfigureTpuVersion()
+    c.configure_tpu_version('1.15', restart_type='ifNeeded')
+    paths = [call[0][0].full_url for call in urlopen.call_args_list]
+    self.assertCountEqual([
+        'http://1.2.3.4:8475/requestversion/1.15?restartType=ifNeeded',
+        'http://5.6.7.8:8475/requestversion/1.15?restartType=ifNeeded'
+    ], sorted(paths))
+
+  @mock.patch.object(urllib.request, 'urlopen')
+  def testGetTpuVersion(self, urlopen):
+    c = client.Client(
+        tpu='grpc://1.2.3.4:8470')
+    resp = mock.Mock()
+    resp.read.side_effect = ['{}', '{"currentVersion": "someVersion"}']
+    urlopen.return_value = resp
+    self.assertIsNone(c.runtime_version(), 'Missing key should be handled.')
+    self.assertEqual(
+        'someVersion', c.runtime_version(), 'Should return configured version.')
+    paths = [call[0][0].full_url for call in urlopen.call_args_list]
+    self.assertCountEqual([
+        'http://1.2.3.4:8475/requestversion',
+        'http://1.2.3.4:8475/requestversion',
+    ], sorted(paths))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/tpu/client/pip_package/BUILD

@@ -0,0 +1,20 @@
+# Description:
+#  Tools for building the Cloud TPU Client pip package.
+
+load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:private"],
+)
+
+licenses(["notice"])
+
+sh_binary(
+    name = "build_pip_package",
+    srcs = ["build_pip_package.sh"],
+    data = [
+        "setup.py",
+        "//tensorflow/python/tpu/client:client_lib",
+    ],
+)

tensorflow/python/tpu/client/pip_package/README

@@ -0,0 +1,3 @@
+Client responsible for communicating the Cloud TPU API. Released separately from tensorflow.
+
+https://pypi.org/project/cloud-tpu-client/

tensorflow/python/tpu/client/pip_package/build_pip_package.sh

@@ -0,0 +1,64 @@
+#!/usr/bin/env bash
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+
+set -e
+
+if [ "$(uname)" = "Darwin" ]; then
+  sedi="sed -i ''"
+else
+  sedi="sed -i"
+fi
+
+PACKAGE_NAME="cloud_tpu_client"
+PIP_PACKAGE="tensorflow/python/tpu/client/pip_package"
+RUNFILES="bazel-bin/tensorflow/python/tpu/client/pip_package/build_pip_package.runfiles/org_tensorflow/tensorflow/python/tpu/client"
+
+function main() {
+  if [ $# -lt 1 ] ; then
+    echo "No destination dir provided"
+    exit 1
+  fi
+
+  DEST=$1
+  TMPDIR=$(mktemp -d -t tmp.XXXXXXXXXX)
+
+  echo $(date) : "=== Using tmpdir: ${TMPDIR}"
+
+  cp ${PIP_PACKAGE}/README ${TMPDIR}
+  cp ${PIP_PACKAGE}/setup.py ${TMPDIR}
+  mkdir ${TMPDIR}/${PACKAGE_NAME}
+  cp -a ${RUNFILES}/. ${TMPDIR}/${PACKAGE_NAME}/
+
+  # Fix the import statements to reflect the copied over path.
+  find  ${TMPDIR}/${PACKAGE_NAME} -name \*.py |
+    xargs $sedi -e '
+      s/^from tensorflow.python.tpu.client/from '${PACKAGE_NAME}'/
+  '
+  echo $(ls $TMPDIR)
+
+  pushd ${TMPDIR}
+  echo $(date) : "=== Building wheel"
+  echo $(pwd)
+  python3 setup.py bdist_wheel >/dev/null
+  mkdir -p ${DEST}
+  cp dist/* ${DEST}
+  popd
+  rm -rf ${TMPDIR}
+  echo $(date) : "=== Output wheel file is in: ${DEST}"
+}
+
+main "$@"

tensorflow/python/tpu/client/pip_package/setup.py

@@ -0,0 +1,50 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Cloud TPU Client package."""
+
+from cloud_tpu_client.version import __version__
+from setuptools import find_packages
+from setuptools import setup
+
+setup(
+    name='cloud-tpu-client',
+    version=__version__.replace('-', ''),
+    description='Client for using Cloud TPUs',
+    long_description='Client for using Cloud TPUs',
+    url='https://cloud.google.com/tpu/',
+    author='Google Inc.',
+    author_email='packages@tensorflow.org',
+    packages=find_packages(),
+    classifiers=[
+        'Development Status :: 5 - Production/Stable',
+        'Intended Audience :: Developers',
+        'Intended Audience :: Education',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.4',
+        'Programming Language :: Python :: 3.5',
+        'Programming Language :: Python :: 3.6',
+        'Topic :: Scientific/Engineering',
+        'Topic :: Scientific/Engineering :: Mathematics',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development',
+        'Topic :: Software Development :: Libraries',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+    ],
+    license='Apache 2.0',
+    keywords='tensorflow tpu',
+    install_requires=['google-api-python-client==1.8.0', 'oauth2client']
+)

tensorflow/python/tpu/client/version.py

@@ -0,0 +1,17 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Cloud TPU Client version information."""
+
+__version__ = "0.11"

tensorflow/python/tpu/datasets.py

@@ -0,0 +1,202 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ======================================
+"""Library of Cloud TPU helper functions for data loading."""
+
+from typing import Callable, Optional, Text, Union
+
+from tensorflow.python.data.experimental.ops import interleave_ops
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import iterator_ops
+from tensorflow.python.data.ops import readers
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import functional_ops
+from tensorflow.python.types import data as data_types
+
+
+def _TextLineDataset(filename: Text) -> dataset_ops.Dataset:
+  buffer_size = 8 * 1024 * 1024  # 8 MiB per file
+  dataset = readers.TextLineDataset(filename, buffer_size=buffer_size)
+  return dataset
+
+
+def _TFRecordDataset(filename: Text) -> dataset_ops.Dataset:
+  buffer_size = 8 * 1024 * 1024  # 8 MiB per file
+  dataset = readers.TFRecordDataset(filename, buffer_size=buffer_size)
+  return dataset
+
+
+_FILETYPE_MAP = {
+    'tfrecord': _TFRecordDataset,
+    'textline': _TextLineDataset,
+    'text': _TextLineDataset,
+}
+
+
+def StreamingFilesDataset(
+    files: Union[Text, dataset_ops.Dataset],
+    filetype: Optional[Union[Text, Callable[[Text],
+                                            dataset_ops.Dataset]]] = None,
+    file_reader_job: Optional[Text] = None,
+    worker_job: Optional[Text] = None,
+    num_epochs: Optional[int] = None,
+    filename_shuffle_buffer_size: Optional[Union[int, bool]] = None,
+    num_parallel_reads: Optional[int] = None,
+    batch_transfer_size: Optional[Union[int, bool]] = None,
+    sloppy: bool = True) -> dataset_ops.Dataset:
+  """StreamingFilesDataset constructs a dataset to stream from workers (GCE VM).
+
+  Because Cloud TPUs are allocated over the network, a Cloud TPU cannot read
+  files local to your GCE VM. In order to train using files stored on your local
+  VM (e.g. on local SSD for extreme performance), use the StreamingFilesDataset
+  helper to generate a dataset to feed your Cloud TPU with files from your GCE
+  VM.
+
+  The resulting dataset may return an OutOfRangeError if there are no files
+  found as a result of the fileglob expansion.
+
+  Note: StreamingFilesDataset assumes that the session is using a
+  TPUClusterResolver and has therefore a worker and a coordinator job. File
+  loading will be done on the coordinator job.
+
+  Args:
+    files: A string glob to match files, or a `tf.data.Dataset` generating file
+      names.
+    filetype: A string (one of 'tfrecord', or 'textline') or a single-argument
+      TensorFlow function that when given a filename returns a dataset.
+    file_reader_job: An optional string that corresponds to the job that should
+      perform the file reads.
+    worker_job: An optional string that corresponds to the job that should
+      process the tensors (i.e. your GPU or TPU worker).
+    num_epochs: The number of epochs through the training set that should be
+      generated. By default, it will repeat infinitely.
+    filename_shuffle_buffer_size: An optional integer whose value controls the
+      shuffling of the file names. If you would like to read from the files in
+      the same order, set to 0 or False.
+    num_parallel_reads: An optional integer controlling the number of files to
+      read from concurrently. (Set to 1 for no parallelism.)
+    batch_transfer_size: An optional integer controlling the batching used to
+      amortize the remote function invocation overhead. Set to a very large
+      number to increase throughput. Set to a very small number to reduce memory
+      consumption. Set to False to skip batching.
+    sloppy: (Optional.) If `False`, read input data while maintaining a
+      deterministic order. (This may have significant performance impacts.)
+      sloppy defaults to: True.
+  Returns:
+    A `tf.data.Dataset` with an infinite stream of elements generated by a
+    parallel interleaving of the set of files matched (or generated) by `files`
+    with a type is the output of the dataset specified by `filetype`.
+
+  Raises:
+    ValueError: if any argument is not of the expected type.
+  """
+  if filetype is None:
+    filetype = 'tfrecord'
+
+  if isinstance(filetype, str):
+    if filetype not in _FILETYPE_MAP:
+      raise ValueError(
+          f'Unexpected filetype. Received: {filetype}. Expected one of '
+          f'{list(_FILETYPE_MAP.keys())}')
+    reader_fn = _FILETYPE_MAP[filetype]
+  elif callable(filetype):
+    reader_fn = filetype
+  else:
+    raise ValueError(f'Argument `filetype` should be a string or a callable. '
+                     f'Received: {filetype} of type {type(filetype)}.')
+
+  file_reader_job = file_reader_job or 'coordinator'
+
+  worker_job = worker_job or 'worker'
+
+  if filename_shuffle_buffer_size is None:
+    filename_shuffle_buffer_size = 4096
+
+  num_parallel_reads = num_parallel_reads or 8
+
+  if batch_transfer_size is None:
+    batch_transfer_size = 256
+
+  if file_reader_job == 'coordinator':
+    file_reader_device = '/job:coordinator/task:0'
+  else:
+    file_reader_device = '/job:%s' % file_reader_job
+
+  with ops.device(file_reader_device):
+    if isinstance(files, str):
+      source_dataset = dataset_ops.Dataset.list_files(files)
+    elif isinstance(files, data_types.DatasetV2):
+      source_dataset = files
+    else:
+      raise ValueError(
+          'Argument `files` should be a string or a `tf.data.Dataset` '
+          f'instance. Received: {files}')
+
+    if filename_shuffle_buffer_size:
+      source_dataset = source_dataset.shuffle(
+          buffer_size=filename_shuffle_buffer_size)
+
+    source_dataset = source_dataset.apply(
+        interleave_ops.parallel_interleave(
+            reader_fn, cycle_length=num_parallel_reads, sloppy=sloppy))
+
+    source_dataset = source_dataset.repeat(num_epochs)
+
+    if batch_transfer_size:
+      source_dataset = source_dataset.batch(batch_transfer_size)
+
+    source_dataset = source_dataset.prefetch(1)
+
+    source_iterator = dataset_ops.make_one_shot_iterator(source_dataset)
+    source_handle = source_iterator.string_handle()
+
+  @function.Defun(dtypes.string)
+  def LoadingFunc(h):
+    remote_iterator = iterator_ops.Iterator.from_string_handle(
+        h, dataset_ops.get_legacy_output_types(source_dataset),
+        dataset_ops.get_legacy_output_shapes(source_dataset))
+    return remote_iterator.get_next()
+
+  def MapFn(unused_input):
+    source_dataset_output_types = dataset_ops.get_legacy_output_types(
+        source_dataset)
+    if isinstance(source_dataset_output_types, dtypes.DType):
+      output_types = [source_dataset_output_types]
+    elif isinstance(source_dataset_output_types, (list, tuple)):
+      output_types = source_dataset_output_types
+    else:
+      raise ValueError('Source dataset has invalid output types. Only '
+                       'list/tuples or TensorFlow tensor types are accepted.')
+    remote_calls = functional_ops.remote_call(
+        args=[source_handle],
+        Tout=output_types,
+        f=LoadingFunc,
+        target='/job:%s/replica:0/task:0/cpu:0' % file_reader_job)
+    if len(remote_calls) == 1:
+      return remote_calls[0]
+    else:
+      return remote_calls
+
+  with ops.device('/job:%s' % worker_job):
+    output_dataset = dataset_ops.Dataset.range(2).repeat().map(
+        MapFn, num_parallel_calls=4 if sloppy else None)
+    output_dataset = output_dataset.prefetch(1)
+
+    if batch_transfer_size:
+      # Undo the batching used during the transfer.
+      output_dataset = output_dataset.unbatch().prefetch(1)
+
+  return output_dataset

tensorflow/python/tpu/datasets_test.py

@@ -0,0 +1,210 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""TPU datasets tests."""
+
+import os
+
+from tensorflow.core.protobuf import cluster_pb2
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import readers
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.lib.io import python_io
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import datasets
+from tensorflow.python.training import server_lib
+from tensorflow.python.util import compat
+
+_NUM_FILES = 10
+_NUM_ENTRIES = 20
+
+
+class DatasetsTest(test.TestCase):
+
+  def setUp(self):
+    super(DatasetsTest, self).setUp()
+    self._coord = server_lib.Server.create_local_server()
+    self._worker = server_lib.Server.create_local_server()
+
+    self._cluster_def = cluster_pb2.ClusterDef()
+    worker_job = self._cluster_def.job.add()
+    worker_job.name = 'worker'
+    worker_job.tasks[0] = self._worker.target[len('grpc://'):]
+    coord_job = self._cluster_def.job.add()
+    coord_job.name = 'coordinator'
+    coord_job.tasks[0] = self._coord.target[len('grpc://'):]
+
+    session_config = config_pb2.ConfigProto(cluster_def=self._cluster_def)
+
+    self._sess = session.Session(self._worker.target, config=session_config)
+    self._worker_device = '/job:' + worker_job.name
+
+  def testTextLineDataset(self):
+    all_contents = []
+    for i in range(_NUM_FILES):
+      filename = os.path.join(self.get_temp_dir(), 'text_line.%d.txt' % i)
+      contents = []
+      for j in range(_NUM_ENTRIES):
+        contents.append(compat.as_bytes('%d: %d' % (i, j)))
+      with open(filename, 'wb') as f:
+        f.write(b'\n'.join(contents))
+      all_contents.extend(contents)
+
+    dataset = datasets.StreamingFilesDataset(
+        os.path.join(self.get_temp_dir(), 'text_line.*.txt'), filetype='text')
+
+    with ops.device(self._worker_device):
+      iterator = dataset_ops.make_initializable_iterator(dataset)
+    self._sess.run(iterator.initializer)
+    get_next = iterator.get_next()
+
+    retrieved_values = []
+    for _ in range(4 * len(all_contents)):
+      retrieved_values.append(compat.as_bytes(self._sess.run(get_next)))
+
+    self.assertEqual(set(all_contents), set(retrieved_values))
+
+  def testTFRecordDataset(self):
+    all_contents = []
+    for i in range(_NUM_FILES):
+      filename = os.path.join(self.get_temp_dir(), 'tf_record.%d' % i)
+      writer = python_io.TFRecordWriter(filename)
+      for j in range(_NUM_ENTRIES):
+        record = compat.as_bytes('Record %d of file %d' % (j, i))
+        writer.write(record)
+        all_contents.append(record)
+      writer.close()
+
+    dataset = datasets.StreamingFilesDataset(
+        os.path.join(self.get_temp_dir(), 'tf_record*'), filetype='tfrecord')
+
+    with ops.device(self._worker_device):
+      iterator = dataset_ops.make_initializable_iterator(dataset)
+    self._sess.run(iterator.initializer)
+    get_next = iterator.get_next()
+
+    retrieved_values = []
+    for _ in range(4 * len(all_contents)):
+      retrieved_values.append(compat.as_bytes(self._sess.run(get_next)))
+
+    self.assertEqual(set(all_contents), set(retrieved_values))
+
+  def testTFRecordDatasetFromDataset(self):
+    filenames = []
+    all_contents = []
+    for i in range(_NUM_FILES):
+      filename = os.path.join(self.get_temp_dir(), 'tf_record.%d' % i)
+      filenames.append(filename)
+      writer = python_io.TFRecordWriter(filename)
+      for j in range(_NUM_ENTRIES):
+        record = compat.as_bytes('Record %d of file %d' % (j, i))
+        writer.write(record)
+        all_contents.append(record)
+      writer.close()
+
+    filenames = dataset_ops.Dataset.from_tensor_slices(filenames)
+
+    dataset = datasets.StreamingFilesDataset(filenames, filetype='tfrecord')
+
+    with ops.device(self._worker_device):
+      iterator = dataset_ops.make_initializable_iterator(dataset)
+    self._sess.run(iterator.initializer)
+    get_next = iterator.get_next()
+
+    retrieved_values = []
+    for _ in range(4 * len(all_contents)):
+      retrieved_values.append(compat.as_bytes(self._sess.run(get_next)))
+
+    self.assertEqual(set(all_contents), set(retrieved_values))
+
+  def testArbitraryReaderFunc(self):
+
+    def MakeRecord(i, j):
+      return compat.as_bytes('%04d-%04d' % (i, j))
+
+    record_bytes = len(MakeRecord(10, 200))
+
+    all_contents = []
+    for i in range(_NUM_FILES):
+      filename = os.path.join(self.get_temp_dir(), 'fixed_length.%d' % i)
+      with open(filename, 'wb') as f:
+        for j in range(_NUM_ENTRIES):
+          record = MakeRecord(i, j)
+          f.write(record)
+          all_contents.append(record)
+
+    def FixedLengthFile(filename):
+      return readers.FixedLengthRecordDataset(filename, record_bytes)
+
+    dataset = datasets.StreamingFilesDataset(
+        os.path.join(self.get_temp_dir(), 'fixed_length*'),
+        filetype=FixedLengthFile)
+
+    with ops.device(self._worker_device):
+      iterator = dataset_ops.make_initializable_iterator(dataset)
+    self._sess.run(iterator.initializer)
+    get_next = iterator.get_next()
+
+    retrieved_values = []
+    for _ in range(4 * len(all_contents)):
+      retrieved_values.append(compat.as_bytes(self._sess.run(get_next)))
+
+    self.assertEqual(set(all_contents), set(retrieved_values))
+
+  def testArbitraryReaderFuncFromDatasetGenerator(self):
+
+    def my_generator():
+      yield (1, [1] * 10)
+
+    def gen_dataset(dummy):
+      return dataset_ops.Dataset.from_generator(
+          my_generator, (dtypes.int64, dtypes.int64),
+          (tensor_shape.TensorShape([]), tensor_shape.TensorShape([10])))
+
+    dataset = datasets.StreamingFilesDataset(
+        dataset_ops.Dataset.range(10), filetype=gen_dataset)
+
+    with ops.device(self._worker_device):
+      iterator = dataset_ops.make_initializable_iterator(dataset)
+    self._sess.run(iterator.initializer)
+    get_next = iterator.get_next()
+
+    retrieved_values = self._sess.run(get_next)
+
+    self.assertIsInstance(retrieved_values, (list, tuple))
+    self.assertEqual(len(retrieved_values), 2)
+    self.assertEqual(retrieved_values[0], 1)
+    self.assertItemsEqual(retrieved_values[1], [1] * 10)
+
+  def testUnexpectedFiletypeString(self):
+    with self.assertRaises(ValueError):
+      datasets.StreamingFilesDataset(
+          os.path.join(self.get_temp_dir(), '*'), filetype='foo')
+
+  def testUnexpectedFiletypeType(self):
+    with self.assertRaises(ValueError):
+      datasets.StreamingFilesDataset(
+          os.path.join(self.get_temp_dir(), '*'), filetype=3)
+
+  def testUnexpectedFilesType(self):
+    with self.assertRaises(ValueError):
+      datasets.StreamingFilesDataset(123, filetype='tfrecord')
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/tpu/device_assignment.py

@@ -0,0 +1,575 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ======================================
+"""Library of TPU helper functions."""
+
+import enum
+import math
+from typing import List, Optional, Tuple
+
+import numpy as np
+
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.tpu.topology import Topology
+from tensorflow.python.util import numpy_compat
+from tensorflow.python.util.tf_export import tf_export
+
+
+SINGLE_CORE_ASSIGNMENT = [[[0, 0, 0, 0]]]
+
+
+def _compute_task_and_cores_to_replicas(core_assignment, topology):
+  """Computes a nested dict which maps task and logical core to replicas."""
+  task_and_cores_to_replicas = {}
+  for replica in range(core_assignment.shape[0]):
+    for logical_core in range(core_assignment.shape[1]):
+      coordinates = core_assignment[replica, logical_core, :]
+      task_id = topology.task_ordinal_at_coordinates(coordinates)
+      if task_id not in task_and_cores_to_replicas:
+        task_and_cores_to_replicas[task_id] = {}
+      if logical_core not in task_and_cores_to_replicas[task_id]:
+        task_and_cores_to_replicas[task_id][logical_core] = set()
+
+      task_and_cores_to_replicas[task_id][logical_core].add(replica)
+
+  task_to_sorted_replica_id = {}
+
+  for task, core_to_replicas in task_and_cores_to_replicas.items():
+    core_to_sorted_replicas = {}
+    for core, replicas in core_to_replicas.items():
+      core_to_sorted_replicas[core] = sorted(replicas)
+
+    task_to_sorted_replica_id[task] = core_to_sorted_replicas
+  return task_to_sorted_replica_id
+
+
+@tf_export("tpu.experimental.DeviceOrderMode")
+class DeviceOrderMode(enum.IntEnum):
+  """The way of determining device orders when computing device assignment."""
+  # By default the mode is set to AUTO, the library will choose to form rings
+  # when that is possible.
+  AUTO = 0
+  # Form rings for replicas and model-parallel cores.
+  RING = 1
+  # Form meshes for replicas and/or model-parallel cores.
+  MESH = 2
+
+
+@tf_export("tpu.experimental.DeviceAssignment")
+class DeviceAssignment(object):
+  """Mapping from logical cores in a computation to the physical TPU topology.
+
+  Prefer to use the `DeviceAssignment.build()` helper to construct a
+  `DeviceAssignment`; it is easier if less flexible than constructing a
+  `DeviceAssignment` directly.
+  """
+
+  def __init__(self, topology: Topology, core_assignment: np.ndarray):
+    """Constructs a `DeviceAssignment` object.
+
+    Args:
+      topology: A `Topology` object that describes the physical TPU topology.
+      core_assignment: A logical to physical core mapping, represented as a
+        rank 3 numpy array. See the description of the `core_assignment`
+        property for more details.
+
+    Raises:
+      ValueError: If `topology` is not `Topology` object.
+      ValueError: If `core_assignment` is not a rank 3 numpy array.
+    """
+    if not isinstance(topology, Topology):
+      raise ValueError("topology must be a Topology object, got {}".format(
+          type(topology)))
+    core_assignment = numpy_compat.np_asarray(core_assignment, dtype=np.int32)
+
+    self._topology = topology
+
+    if core_assignment.ndim != 3:
+      raise ValueError("core_assignment must be a rank 3 numpy array, "
+                       f"got shape {core_assignment.shape}")
+
+    self._num_replicas = core_assignment.shape[0]
+    self._num_cores_per_replica = core_assignment.shape[1]
+
+    if core_assignment.shape[-1] != topology.mesh_rank:
+      raise ValueError(
+          "core_assignment.shape[-1] must have size equal to topology "
+          f"rank ({topology.mesh_rank}), got "
+          f"core_assignment.shape={core_assignment.shape}")
+
+    self._core_assignment = core_assignment
+    self._task_and_cores_to_replicas = _compute_task_and_cores_to_replicas(
+        self._core_assignment, topology)
+
+  @property
+  def topology(self) -> Topology:
+    """A `Topology` that describes the TPU topology."""
+    return self._topology
+
+  @property
+  def num_cores_per_replica(self) -> int:
+    """The number of cores per replica."""
+    return self._num_cores_per_replica
+
+  @property
+  def num_replicas(self) -> int:
+    """The number of replicas of the computation."""
+    return self._num_replicas
+
+  @property
+  def core_assignment(self) -> np.ndarray:
+    """The logical to physical core mapping.
+
+    Returns:
+      An integer numpy array of rank 3, with shape
+      `[num_replicas, num_cores_per_replica, topology_rank]`. Maps
+      (replica, logical core) pairs to physical topology coordinates.
+    """
+    return self._core_assignment
+
+  def coordinates(self, replica: int, logical_core: int) -> Tuple:  # pylint:disable=g-bare-generic
+    """Returns the physical topology coordinates of a logical core."""
+    return tuple(self.core_assignment[replica, logical_core, :])
+
+  def lookup_replicas(self, task_id: int, logical_core: int) -> List[int]:
+    """Lookup replica ids by task number and logical core.
+
+    Args:
+      task_id: TensorFlow task number.
+      logical_core: An integer, identifying a logical core.
+    Returns:
+      A sorted list of the replicas that are attached to that task and
+      logical_core.
+    Raises:
+      ValueError: If no replica exists in the task which contains the logical
+      core.
+    """
+    try:
+      return self._task_and_cores_to_replicas[task_id][logical_core]
+    except KeyError:
+      raise ValueError(
+          "Can not find any replica in task: {} contains logical_core: {} ".
+          format(task_id, logical_core))
+
+  def tpu_ordinal(self, replica: int = 0, logical_core: int = 0) -> int:
+    """Returns the ordinal of the TPU device assigned to a logical core."""
+    coordinates = self.coordinates(replica, logical_core)
+    return self._topology.tpu_device_ordinal_at_coordinates(coordinates)
+
+  def host_device(self,
+                  replica: int = 0,
+                  logical_core: int = 0,
+                  job: Optional[str] = None) -> str:
+    """Returns the CPU device attached to a logical core."""
+    coordinates = self.coordinates(replica, logical_core)
+    return self._topology.cpu_device_name_at_coordinates(coordinates, job=job)
+
+  def tpu_device(self,
+                 replica: int = 0,
+                 logical_core: int = 0,
+                 job: Optional[str] = None) -> str:
+    """Returns the name of the TPU device assigned to a logical core."""
+    coordinates = self.coordinates(replica, logical_core)
+    return self._topology.tpu_device_name_at_coordinates(coordinates, job=job)
+
+  @classmethod
+  def build(
+      cls,
+      topology: Topology,
+      computation_shape: Optional[np.ndarray] = None,
+      computation_stride: Optional[np.ndarray] = None,
+      num_replicas: int = 1,
+      device_order_mode: DeviceOrderMode = DeviceOrderMode.AUTO,
+  ) -> "DeviceAssignment":
+    return device_assignment(
+        topology=topology,
+        computation_shape=computation_shape,
+        computation_stride=computation_stride,
+        num_replicas=num_replicas,
+        device_order_mode=device_order_mode,
+    )
+
+
+def _open_ring_2d(x_size: int, y_size: int,
+                  z_coord: int) -> List[Tuple[int, int, int]]:
+  """Ring-order of a X by Y mesh, with a fixed Z coordinate.
+
+  For example, in a 4x4 mesh, this returns the following order.
+    0 -- 1 -- 2 -- 3
+    |    |    |    |
+    15-- 6 -- 5 -- 4
+    |    |    |    |
+    14-- 7 -- 8 -- 9
+    |    |    |    |
+    13-- 12-- 11-- 10
+
+  Note that chip 0 is not included in the output.
+
+  Args:
+    x_size: An integer represents the mesh size in the x-dimension. Must be
+      larger than 1.
+    y_size: An integer represents the mesh size in the y-dimension. Must be
+      larger than 1.
+    z_coord: An integer represents the z-coordinate to use for the chips in the
+      ring.
+
+  Returns:
+    A list of (x,y,z) triples in ring order.
+  """
+  ret = []
+  for i in range(y_size // 2):
+    for j in range(1, x_size):
+      ret.append((j, 2 * i, z_coord))
+    for j in range(x_size - 1, 0, -1):
+      ret.append((j, 2 * i + 1, z_coord))
+  for i in range(y_size - 1, 0, -1):
+    ret.append((0, i, z_coord))
+  return ret
+
+
+def _ring_3d(x_size: int, y_size: int,
+             z_size: int) -> List[Tuple[int, int, int]]:
+  """Ring-order of a X by Y by Z mesh.
+
+  Constructs the 3d ring from 2d rings that are stacked in the Z dimension and
+  joined in one corner.
+
+  z == 0:
+    0 -- 1 -- 2 -- 3
+    |    |    |    |
+    15 - 6 -- 5 -- 4
+    |    |    |    |
+    14 - 7 -- 8 -- 9
+    |    |    |    |
+    13 - 12 - 11 - 10
+  z == 1:
+    63 - 30 - 29 - 28
+    |    |    |    |
+    16 - 25 - 26 - 27
+    |    |    |    |
+    17 - 24 - 23 - 22
+    |    |    |    |
+    18 - 19 - 20 - 21
+  z == 2:
+    62 - 31 - 32 - 33
+    |    |    |    |
+    45 - 36 - 35 - 34
+    |    |    |    |
+    44 - 37 - 38 - 39
+    |    |    |    |
+    43 - 42 - 41 - 40
+  z == 3:
+    61 - 60 - 59 - 58
+    |    |    |    |
+    46 - 55 - 56 - 57
+    |    |    |    |
+    47 - 54 - 53 - 52
+    |    |    |    |
+    48 - 49 - 50 - 51
+
+  Args:
+    x_size: An integer represents the mesh size in the x-dimension. Must be
+      larger than 1.
+    y_size: An integer represents the mesh size in the y-dimension. Must be
+      larger than 1.
+    z_size: An integer represents the mesh size in the z-dimension. Must be
+      larger than 1.  For example, in a 4x4x4 mesh, this returns the following
+      order.
+
+  Returns:
+    A list of (x,y,z) triples in ring order.
+  """
+
+  # Handle the case where 2 dimensions are size 1.
+  if x_size == 1 and y_size == 1:
+    return [(0, 0, i) for i in range(z_size)]
+  if x_size == 1 and z_size == 1:
+    return [(0, i, 0) for i in range(y_size)]
+  if y_size == 1 and z_size == 1:
+    return [(i, 0, 0) for i in range(x_size)]
+
+  # Handle odd mesh dimensions.  This never happens in practice, so we don't
+  # bother to try building something optimal.
+  if (x_size > 1 and x_size % 2 != 0) or (y_size > 1 and
+                                          y_size % 2 != 0) or (z_size > 1 and
+                                                               z_size % 2 != 0):
+    logging.warning("Odd dimension")
+    ret = []
+    for z in range(z_size):
+      for y in range(y_size):
+        ret.extend((x, y, z) for x in range(x_size))
+    return ret
+
+  # Always start with chip 0.
+  ret = [(0, 0, 0)]
+  # Handle the case where one dimension is size 1.  We just build a flat, 2d
+  # ring.
+  if z_size == 1:
+    ret.extend(_open_ring_2d(x_size, y_size, 0))
+    return ret
+  if y_size == 1:
+    ret = [(0, 0, 0)]
+    ret.extend((x, y, z) for (x, z, y) in _open_ring_2d(x_size, z_size, 0))
+    return ret
+  if x_size == 1:
+    ret = [(0, 0, 0)]
+    ret.extend((x, y, z) for (y, z, x) in _open_ring_2d(y_size, z_size, 0))
+    return ret
+
+  # Handle the case where all dimensions have size > 1 and even.
+  ret = [(0, 0, 0)]
+  for i in range(0, z_size):
+    r = _open_ring_2d(x_size, y_size, i)
+    if i % 2 == 0:
+      ret.extend(r)
+    else:
+      ret.extend(reversed(r))
+  for i in range(z_size - 1, 0, -1):
+    ret.append((0, 0, i))
+  return ret
+
+
+def device_assignment(
+    topology: Topology,
+    computation_shape: Optional[np.ndarray] = None,
+    computation_stride: Optional[np.ndarray] = None,
+    num_replicas: int = 1,
+    device_order_mode: DeviceOrderMode = DeviceOrderMode.AUTO,
+) -> DeviceAssignment:
+  """Computes a device_assignment of a computation across a TPU topology.
+
+  Attempts to choose a compact grid of cores for locality.
+
+  Returns a `DeviceAssignment` that describes the cores in the topology assigned
+  to each core of each replica.
+
+  `computation_shape` and `computation_stride` values should be powers of 2 for
+  optimal packing.
+
+  Args:
+    topology: A `Topology` object that describes the TPU cluster topology. To
+      obtain a TPU topology, evaluate the `Tensor` returned by
+      `initialize_system` using `Session.run`. Either a serialized
+      `TopologyProto` or a `Topology` object may be passed. Note: you must
+        evaluate the `Tensor` first; you cannot pass an unevaluated `Tensor`
+        here.
+    computation_shape: A rank 1 int32 numpy array with size equal to the
+      topology rank, describing the shape of the computation's block of cores.
+      If None, the `computation_shape` is `[1] * topology_rank`.
+    computation_stride: A rank 1 int32 numpy array of size `topology_rank`,
+      describing the inter-core spacing of the `computation_shape` cores in the
+      TPU topology. If None, the `computation_stride` is `[1] * topology_rank`.
+    num_replicas: The number of computation replicas to run. The replicas will
+      be packed into the free spaces of the topology.
+    device_order_mode: An enum of `DeviceOrderMode` class which indicates
+      whether to assign devices to form rings or meshes, or let the library to
+      choose.
+
+  Returns:
+    A DeviceAssignment object, which describes the mapping between the logical
+    cores in each computation replica and the physical cores in the TPU
+    topology.
+
+  Raises:
+    ValueError: If `topology` is not a valid `Topology` object.
+    ValueError: If `computation_shape` or `computation_stride` are not 1D int32
+      numpy arrays with shape [3] where all values are positive.
+    ValueError: If computation's replicas cannot fit into the TPU topology.
+  """
+  # Deserialize the Topology proto, if it is a string.
+  if isinstance(topology, bytes):
+    topology = Topology(serialized=topology)
+
+  if not isinstance(topology, Topology):
+    raise ValueError(
+        f"`topology` is not a Topology object; got {type(topology)}"
+    )
+
+  topology_rank = len(topology.mesh_shape)
+  mesh_shape = topology.mesh_shape
+  if computation_shape is None:
+    computation_shape = np.array([1] * topology_rank, dtype=np.int32)
+  else:
+    computation_shape = numpy_compat.np_asarray(
+        computation_shape, dtype=np.int32
+    )
+
+  if computation_stride is None:
+    computation_stride = np.array([1] * topology_rank, dtype=np.int32)
+  else:
+    computation_stride = numpy_compat.np_asarray(
+        computation_stride, dtype=np.int32
+    )
+
+  if computation_shape.shape != (topology_rank,):
+    raise ValueError(
+        f"computation_shape must have shape [{topology_rank}]; "
+        f"got {computation_shape.shape}"
+    )
+  if computation_stride.shape != (topology_rank,):
+    raise ValueError(
+        f"computation_stride must have shape [{topology_rank}]; "
+        f"got {computation_stride.shape}"
+    )
+
+  if any(computation_shape < 1):
+    raise ValueError(
+        "computation_shape must be positive; got computation_shape={}".format(
+            computation_shape))
+  if any(computation_stride < 1):
+    raise ValueError(
+        "computation_stride must be positive; got computation_stride={}".format(
+            computation_stride))
+
+  # Computes the physical size of one computation instance.
+  computation_footprint = computation_shape * computation_stride
+  if any(computation_footprint > mesh_shape):
+    raise ValueError(
+        "computation footprint {} does not fit in TPU topology shape {}".format(
+            computation_footprint, mesh_shape))
+
+  # Computes how many copies of the computation footprint fit in the mesh.
+  block_counts = mesh_shape // computation_footprint
+
+  replica_counts = block_counts * computation_stride
+  max_replicas = np.prod(replica_counts)
+  if num_replicas > max_replicas:
+    raise ValueError(
+        "requested {} replicas but only {} replicas with shape {} and "
+        "computation_stride {} fit in a TPU mesh of shape {}".format(
+            num_replicas, max_replicas, computation_shape, computation_stride,
+            mesh_shape))
+
+  def ceil_of_ratio(n, m):
+    return (n + m - 1) // m
+
+  if topology.missing_devices.size == 0:
+    replica_shape = [0] * topology_rank
+    if num_replicas > 0:
+      remaining_replicas = num_replicas
+      remaining_dims = topology_rank
+
+      # Choose dimensions as close to an equal cube as possible,
+      # in order of increasing dimension size. By visiting dimensions
+      # in increasing size, we assign the most constrained dimension
+      # first, so we won't make infeasible choices.
+      #
+      # As a secondary sort order, visit the last dimension (core index) first,
+      # then the other dimensions in increasing order. This means we try to use
+      # both cores on the same chip in preference to two cores on different
+      # chips.  We visit the x dimension first, and the z dimension last, so
+      # that we prefer to arrange adjacent replicas on the same machine when
+      # possible.
+      #
+      # For example, if num_replicas == 4, we prefer to use a replica_shape of
+      # (2,1,1,2) over (1,1,2,2).
+
+      for x, ni in sorted(((x, ((i + 1) % topology_rank))
+                           for (i, x) in enumerate(replica_counts))):
+        i = (ni + topology_rank - 1) % topology_rank
+        target_size = int(math.ceil(remaining_replicas**(1.0 / remaining_dims)))
+        replica_shape[i] = min(target_size, x)
+        remaining_replicas = ceil_of_ratio(remaining_replicas, replica_shape[i])
+        remaining_dims -= 1
+
+      assert remaining_replicas == 1 and remaining_dims == 0
+
+    # Assigns an offset to each replica such that no two replicas overlap.
+    replica_offsets = np.full([num_replicas, topology_rank], -1, dtype=np.int32)
+
+    enable_3d_tiling = (
+        topology_rank == 4 and
+        computation_shape[-1] == mesh_shape[-1]  # Only handle 3D case.
+        and np.prod(computation_stride) == 1  # Ensure no stride.
+        and num_replicas == max_replicas)  # Full replication.
+
+    if device_order_mode != DeviceOrderMode.AUTO:
+      if device_order_mode == DeviceOrderMode.RING and not enable_3d_tiling:
+        raise ValueError(
+            "device_order_mode=DeviceOrderMode.RING is not compatible with the "
+            "3D tiling current topology.  Try setting "
+            "device_order_mode=DeviceOrderMode.AUTO"
+        )
+      enable_3d_tiling = device_order_mode == DeviceOrderMode.RING
+
+    if enable_3d_tiling:
+      assignment = []
+      inner_ring = _ring_3d(computation_shape[0], computation_shape[1],
+                            computation_shape[2])
+      outer_ring = _ring_3d(replica_shape[0], replica_shape[1],
+                            replica_shape[2])
+
+      for replica in range(num_replicas):
+        outer_x, outer_y, outer_z = outer_ring[replica]
+        per_replica_assignment = []
+        for index in range(np.prod(computation_shape)):
+          inner_x, inner_y, inner_z = inner_ring[index // mesh_shape[-1]]
+          px = outer_x * computation_shape[0] + inner_x
+          py = outer_y * computation_shape[1] + inner_y
+          pz = outer_z * computation_shape[2] + inner_z
+          pi = index % mesh_shape[-1]
+          per_replica_assignment.append([px, py, pz, pi])
+        assignment.append(per_replica_assignment)
+    else:
+      for replica in range(num_replicas):
+        # Chooses a replica number in each axis.
+        t = replica
+        pos = []
+        # Visit the core number first.
+        for dim in np.concatenate([[replica_shape[-1]], replica_shape[:-1]]):
+          pos.append(t % dim)
+          t //= dim
+        replica_pos = np.concatenate([pos[1:], [pos[0]]])
+
+        # Determines where that replica starts in each axis.
+        outer = replica_pos // computation_stride
+        inner = replica_pos % computation_stride
+        replica_offsets[replica, :] = outer * computation_footprint + inner
+
+      # Computes a logical core -> physical core mapping for each replica.
+      indices = [
+          np.arange(0, computation_shape[i] * computation_stride[i],
+                    computation_stride[i]) for i in range(topology_rank)
+      ]
+      indices = np.concatenate(
+          [i[..., np.newaxis] for i in np.meshgrid(*indices, indexing="ij")],
+          axis=-1)
+      indices = indices.reshape((-1, topology_rank))
+      assignment = indices + replica_offsets[:, np.newaxis, :]
+  else:
+    # We have a slice with missing chips. We define a simple assignment by
+    # ignoring computation stride. This assignment should enable a consistent
+    # and correct device assignment on degraded slices. It is optimal when
+    # weights are not sharded. But this device assignment may be sub-optimal for
+    # other model parallelism scenarios.
+    assert np.prod(computation_stride) == 1
+    # Next, we check if we have sufficient devices.
+    assert num_replicas * np.prod(
+        computation_shape) <= topology.num_tasks * topology.num_tpus_per_task
+    # Map replicas to physical devices in task order.
+    device_coordinates = topology.device_coordinates
+    assignment = []
+    devices_per_replica = np.prod(computation_shape)
+    for rindex in range(num_replicas):
+      replica_assignment = []
+      for index in range(devices_per_replica):
+        logical_id = rindex * devices_per_replica + index
+        # Pick logical cores in task order
+        task = logical_id // topology.num_tpus_per_task
+        device = logical_id % topology.num_tpus_per_task
+        # Append physical cores to the replica assignment
+        replica_assignment.append(device_coordinates[task, device, :])
+      assignment.append(replica_assignment)
+
+  return DeviceAssignment(topology, core_assignment=assignment)

tensorflow/python/tpu/embedding_context_utils.py

@@ -0,0 +1,46 @@
+# Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Helper functions for TPU Embedding Contexts."""
+
+import threading
+from absl import logging
+
+
+class _EmbeddingPipeliningState(threading.local):
+
+  def __init__(self):
+    super().__init__()
+    self.enabled = True
+
+embedding_pipelining_state = _EmbeddingPipeliningState()
+
+
+class SequentialEmbeddingContext:
+  """Disables embedding pipelining for all ops created in the scope."""
+
+  def __init__(self):
+    self._original_embedding_pipelining_state_enabled = (
+        embedding_pipelining_state.enabled
+    )
+
+  def __enter__(self):
+    embedding_pipelining_state.enabled = False
+    logging.info("Entering SequentialEmbeddingContext.")
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    embedding_pipelining_state.enabled = (
+        self._original_embedding_pipelining_state_enabled
+    )
+    logging.info("Exiting SequentialEmbeddingContext.")

tensorflow/python/tpu/experimental/BUILD

@@ -0,0 +1,17 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "experimental",
+    srcs = [
+        "__init__.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/tpu:tpu_strategy_util",
+    ],
+)

tensorflow/python/tpu/experimental/__init__.py

@@ -0,0 +1,19 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Experimental TPU library."""
+
+# pylint: disable=unused-import
+from tensorflow.python.tpu import tpu_strategy_util
+# pylint: enable=unused-import

tensorflow/python/tpu/feature_column.py

@@ -0,0 +1,690 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===================================================================
+"""TPU Feature Column Library."""
+import math
+
+from tensorflow.python.feature_column import feature_column as fc
+from tensorflow.python.feature_column import feature_column_lib as fc_lib
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.tpu import tpu
+from tensorflow.python.tpu import tpu_function
+from tensorflow.python.tpu import tpu_replication
+# pylint: disable=protected-access
+
+
+_TPU_FC_TO_SCOPE = '_tpu_feature_column_scope'
+_SUPPORTED_SEQUENCE_COLUMNS = (fc._SequenceCategoricalColumn,
+                               fc_lib.SequenceCategoricalColumn)
+
+
+# For V2 columns, we support anything that inherits from CategoricalColumn
+# other than those in the denylist. User-provided columns that inherit from
+# CategoricalColumn may or may not be compatible; it is up to the user to
+# manage TPU compatibility for custom columns.
+_SUPPORTED_CATEGORICAL_COLUMNS_V2 = (fc_lib.CategoricalColumn,)
+_DENYLISTED_CATEGORICAL_COLUMNS_V2 = (fc_lib.HashedCategoricalColumn,
+                                      fc_lib.BucketizedColumn,
+                                      fc_lib.CrossedColumn)
+_SUPPORTED_CATEGORICAL_COLUMNS = (fc._IdentityCategoricalColumn,
+                                  fc._VocabularyFileCategoricalColumn,
+                                  fc._VocabularyListCategoricalColumn,
+                                  fc._WeightedCategoricalColumn,
+                                  fc._SequenceCategoricalColumn
+                                 ) + _SUPPORTED_CATEGORICAL_COLUMNS_V2
+_SEQUENCE_FEATURE_LENGTH_POSTFIX = '_seq_length_'
+
+
+def embedding_column(categorical_column,
+                     dimension,
+                     combiner='mean',
+                     initializer=None,
+                     max_sequence_length=0,
+                     learning_rate_fn=None,
+                     use_safe_embedding_lookup=True):
+  """TPU embedding_column for `tf.feature_column.embedding_column`.
+
+  Note that the interface for TPU embedding_column is different from the non-TPU
+  version. The following args available for the non-TPU version are NOT
+  supported: ckpt_to_load_from, tensor_name_in_ckp, max_norm and trainable.
+
+  Args:
+    categorical_column: A categorical_column returned from
+        categorical_column_with_identity, weighted_categorical_column,
+        categorical_column_with_vocabulary_file,
+        categorical_column_with_vocabulary_list,
+        sequence_categorical_column_with_identity,
+        sequence_categorical_column_with_vocabulary_file,
+        sequence_categorical_column_with_vocabulary_list
+    dimension: An integer specifying dimension of the embedding, must be > 0.
+    combiner: A string specifying how to reduce if there are multiple entries
+      in a single row for a non-sequence column. For more information, see
+      `tf.feature_column.embedding_column`.
+    initializer: A variable initializer function to be used in embedding
+      variable initialization. If not specified, defaults to
+      `tf.compat.v1.truncated_normal_initializer` with mean `0.0` and
+      standard deviation `1/sqrt(dimension)`.
+    max_sequence_length: An non-negative integer specifying the max sequence
+      length. Any sequence shorter then this will be padded with 0 embeddings
+      and any sequence longer will be truncated. This must be positive for
+      sequence features and 0 for non-sequence features.
+    learning_rate_fn: A function that takes global step and returns learning
+      rate for the embedding table. If you intend to use the same learning rate
+      for multiple embedding tables, please ensure that you pass the exact same
+      python function to all calls of embedding_column, otherwise performence
+      may suffer.
+    use_safe_embedding_lookup: If true, uses safe_embedding_lookup_sparse
+      instead of embedding_lookup_sparse. safe_embedding_lookup_sparse ensures
+      there are no empty rows and all weights and ids are positive at the
+      expense of extra compute cost. This only applies to rank 2 (NxM) shaped
+      input tensors. Defaults to true, consider turning off if the above checks
+      are not needed. Note that having empty rows will not trigger any error
+      though the output result might be 0 or omitted.
+
+  Returns:
+    A  _TPUEmbeddingColumn.
+
+  Raises:
+    ValueError: if `dimension` not > 0.
+    ValueError: if `initializer` is specified but not callable.
+    TypeError: if categorical_column is not a supported type.
+  """
+  if isinstance(categorical_column, _DENYLISTED_CATEGORICAL_COLUMNS_V2):
+    raise TypeError('categorical_column for tpu '
+                    ' embedding_column was '
+                    f'denylisted type {type(categorical_column)}')
+  if not isinstance(categorical_column, _SUPPORTED_CATEGORICAL_COLUMNS):
+    raise TypeError(
+        'categorical_column for tpu '
+        ' embedding_column must be type {}, got {}.'.format(' or '.join([
+            cc.__name__ for cc in _SUPPORTED_CATEGORICAL_COLUMNS
+        ]), type(categorical_column)))
+  if (dimension is None) or (dimension < 1):
+    raise ValueError('Invalid dimension {}.'.format(dimension))
+
+  if (initializer is not None) and (not callable(initializer)):
+    raise ValueError('initializer must be callable if specified. '
+                     'Embedding of column_name: {}'.format(
+                         categorical_column.name))
+  if initializer is None:
+    initializer = init_ops.truncated_normal_initializer(
+        mean=0.0, stddev=1 / math.sqrt(dimension))
+
+  embedding_shape = categorical_column._num_buckets, dimension  # pylint: disable=protected-access
+
+  def _creator(weight_collections, scope):
+    embedding_column_layer = fc._EmbeddingColumnLayer(
+        embedding_shape=embedding_shape,
+        initializer=initializer,
+        weight_collections=weight_collections,
+        trainable=True,
+        name='embedding_column_layer')
+    return embedding_column_layer(None, scope=scope)  # pylint: disable=not-callable
+
+  column = _TPUEmbeddingColumn(
+      categorical_column=categorical_column,
+      dimension=dimension,
+      combiner=combiner,
+      layer_creator=_creator,
+      ckpt_to_load_from=None,
+      tensor_name_in_ckpt=None,
+      max_norm=None,
+      trainable=True,
+      max_sequence_length=max_sequence_length,
+      learning_rate_fn=learning_rate_fn,
+      use_safe_embedding_lookup=use_safe_embedding_lookup)
+  # For Embedding column, the initializer is hidden inside the creator Fn, which
+  # is not accessible later. So, we attach it to a special field. Also note
+  # that non-TPU Embedding column and non-TPU shared Embedding column handle the
+  # initializer differently. See shared_embedding_columns for details.
+  column._tpu_initializer = initializer
+  return column
+
+
+def shared_embedding_columns(categorical_columns,
+                             dimension,
+                             combiner='mean',
+                             initializer=None,
+                             shared_embedding_collection_name=None,
+                             max_sequence_lengths=None,
+                             learning_rate_fn=None,
+                             use_safe_embedding_lookup=True):
+  """List of dense columns that convert from sparse, categorical input.
+
+  Note that the interface for TPU embedding_column is different from the non-TPU
+  version. The following args available for the non-TPU version are NOT
+  supported: ckpt_to_load_from, tensor_name_in_ckp, max_norm and trainable.
+
+  Args:
+    categorical_columns: A list of categorical_columns returned from
+        categorical_column_with_identity, weighted_categorical_column,
+        categorical_column_with_vocabulary_file,
+        categorical_column_with_vocabulary_list,
+        sequence_categorical_column_with_identity,
+        sequence_categorical_column_with_vocabulary_file,
+        sequence_categorical_column_with_vocabulary_list
+    dimension: An integer specifying dimension of the embedding, must be > 0.
+    combiner: A string specifying how to reduce if there are multiple entries
+      in a single row for a non-sequence column. For more information, see
+      `tf.feature_column.embedding_column`.
+    initializer: A variable initializer function to be used in embedding
+      variable initialization. If not specified, defaults to
+      `tf.truncated_normal_initializer` with mean `0.0` and standard deviation
+      `1/sqrt(dimension)`.
+    shared_embedding_collection_name: Optional name of the collection where
+      shared embedding weights are added. If not given, a reasonable name will
+      be chosen based on the names of `categorical_columns`. This is also used
+      in `variable_scope` when creating shared embedding weights.
+    max_sequence_lengths: An list of non-negative integers, either None or
+      empty or the same length as the argument categorical_columns. Entries
+      corresponding to non-sequence columns must be 0 and entries corresponding
+      to sequence columns specify the max sequence length for the column. Any
+      sequence shorter then this will be padded with 0 embeddings and any
+      sequence longer will be truncated.
+    learning_rate_fn: A function that takes global step and returns learning
+      rate for the embedding table. If you intend to use the same learning rate
+      for multiple embedding tables, please ensure that you pass the exact same
+      python function to all calls of shared_embedding_columns, otherwise
+      performence may suffer.
+    use_safe_embedding_lookup: If true, uses safe_embedding_lookup_sparse
+      instead of embedding_lookup_sparse. safe_embedding_lookup_sparse ensures
+      there are no empty rows and all weights and ids are positive at the
+      expense of extra compute cost. This only applies to rank 2 (NxM) shaped
+      input tensors. Defaults to true, consider turning off if the above checks
+      are not needed. Note that having empty rows will not trigger any error
+      though the output result might be 0 or omitted.
+
+  Returns:
+    A  _TPUEmbeddingColumn.
+
+  Raises:
+    ValueError: if `dimension` not > 0.
+    ValueError: if `initializer` is specified but not callable.
+    ValueError: if `max_sequence_lengths` is specified and not the same length
+      as `categorical_columns`.
+    ValueError: if `max_sequence_lengths` is positive for a non sequence column
+      or 0 for a sequence column.
+  """
+  for categorical_column in categorical_columns:
+    if isinstance(categorical_column, _DENYLISTED_CATEGORICAL_COLUMNS_V2):
+      raise TypeError('categorical_column for tpu '
+                      ' embedding_column was denylisted type '
+                      f'{type(categorical_column)}')
+    if not isinstance(categorical_column, _SUPPORTED_CATEGORICAL_COLUMNS):
+      raise TypeError(
+          'categorical_column for tpu '
+          ' shared_embedding_columns must be type {}, got {}.'.format(
+              ' or '.join(
+                  [cc.__name__ for cc in _SUPPORTED_CATEGORICAL_COLUMNS]),
+              type(categorical_column)))
+
+  if not max_sequence_lengths:
+    max_sequence_lengths = [0] * len(categorical_columns)
+  if len(max_sequence_lengths) != len(categorical_columns):
+    raise ValueError('max_sequence_lengths and categorical_columns must be of '
+                     'the same length. len(max_sequence_lengths)={} '
+                     'len(categorical_columns)={}.'.format(
+                         len(max_sequence_lengths), len(categorical_columns)))
+
+  if (dimension is None) or (dimension < 1):
+    raise ValueError('Invalid dimension {}.'.format(dimension))
+
+  if (initializer is not None) and (not callable(initializer)):
+    raise ValueError('initializer must be callable if specified. ')
+  if initializer is None:
+    initializer = init_ops.truncated_normal_initializer(
+        mean=0.0, stddev=1 / math.sqrt(dimension))
+
+  # Sort the columns so the default collection name is deterministic even if the
+  # user passes columns from an unsorted collection, such as dict.values().
+  sorted_columns = sorted(categorical_columns, key=lambda x: x.name)
+  num_buckets = sorted_columns[0]._num_buckets  # pylint: disable=protected-access
+
+  for c in sorted_columns[1:]:
+    if num_buckets != c._num_buckets:  # pylint: disable=protected-access
+      raise ValueError(
+          'To use shared_embedding_column, all categorical_columns must have '
+          'the same number of buckets. Given column: {} with buckets: {} does  '
+          'not match column: {} with buckets: {}'.format(
+              sorted_columns[0], num_buckets, c, c._num_buckets))  # pylint: disable=protected-access
+
+  if not shared_embedding_collection_name:
+    shared_embedding_collection_name = '_'.join(c.name for c in sorted_columns)
+    shared_embedding_collection_name += '_shared_embedding'
+
+  tpu_columns = []
+
+  # Create the state (_SharedEmbeddingColumnLayer) here.
+  for categorical_column, max_sequence_length in zip(
+      categorical_columns, max_sequence_lengths):
+    column = _TPUSharedEmbeddingColumn(
+        categorical_column=categorical_column,
+        dimension=dimension,
+        combiner=combiner,
+        initializer=initializer,
+        shared_embedding_collection_name=shared_embedding_collection_name,
+        ckpt_to_load_from=None,
+        tensor_name_in_ckpt=None,
+        max_norm=None,
+        trainable=True,
+        max_sequence_length=max_sequence_length,
+        learning_rate_fn=learning_rate_fn,
+        use_safe_embedding_lookup=use_safe_embedding_lookup)
+    tpu_columns.append(column)
+
+  return tpu_columns
+
+
+class _TPUBaseEmbeddingColumn(object):
+  """Base class for TPU Embedding Column."""
+
+  def __init__(self,
+               categorical_column,
+               max_sequence_length=0,
+               learning_rate_fn=None):
+    self._tpu_categorical_column = categorical_column
+    self._max_sequence_length = max_sequence_length
+    self._learning_rate_fn = learning_rate_fn
+    if (self.is_sequence_column() and max_sequence_length < 1):
+      raise ValueError('max_sequence_length must be greater than 0 for '
+                       'sequence columns. Got max_sequence_length={} for '
+                       'sequence column {}.'.format(max_sequence_length,
+                                                    categorical_column.name))
+    if (not self.is_sequence_column() and max_sequence_length != 0):
+      raise ValueError('Non zero max_seq_length={} specified for non '
+                       'sequence column {}.'.format(max_sequence_length,
+                                                    categorical_column.name))
+
+  def get_combiner(self):
+    """Returns the embedding combiner."""
+    raise NotImplementedError('not implemented')
+
+  def get_embedding_table_size(self):
+    """Returns the embedding table size, tuple of vocab size and dimension."""
+    raise NotImplementedError('not implemented')
+
+  def get_feature_key_name(self):
+    """Returns the feature key name in the features dict."""
+    raise NotImplementedError('not impl')
+
+  def get_weight_key_name(self):
+    """Return the key name for weights."""
+    raise NotImplementedError('not impl')
+
+  def get_embedding_var_name(self):
+    """Returns the embedding variable name.
+
+    Feature key name and embedding variable name are usually one-to-one mapping.
+    But for shared embedding columns, it is many-to-one mapping.
+    """
+    raise NotImplementedError('not impl')
+
+  def get_initializer(self):
+    """Returns the initializer."""
+    raise NotImplementedError('not impl')
+
+  def is_categorical_column_weighted(self):
+    """Check if the categorical column of the embedding column is weighted."""
+    raise NotImplementedError('not impl')
+
+  def is_sequence_column(self):
+    return isinstance(self._tpu_categorical_column, _SUPPORTED_SEQUENCE_COLUMNS)
+
+  def get_max_sequence_length(self):
+    return self._max_sequence_length
+
+  def get_learning_rate_fn(self):
+    return self._learning_rate_fn
+
+  def get_sequence_length_feature_key_name(self):
+    """Get the key for the associated sequence length feature."""
+    return get_sequence_length_feature_key_name_from_feature_key_name(
+        self.get_feature_key_name())
+
+
+class _TPUEmbeddingColumn(_TPUBaseEmbeddingColumn, fc._EmbeddingColumn):
+  """Core Embedding Column."""
+
+  def __new__(cls,
+              categorical_column,
+              dimension,
+              combiner='mean',
+              layer_creator=None,
+              ckpt_to_load_from=None,
+              tensor_name_in_ckpt=None,
+              max_norm=None,
+              trainable=True,
+              max_sequence_length=0,
+              learning_rate_fn=None,
+              use_safe_embedding_lookup=True,
+              bypass_scope_validation=False):
+    # Note, args ckpt_to_load_from, tensor_name_in_ckpt, max_norm and trainable
+    # are not supported on TPU. They are solely for matching the signature of
+    # __new__ of parent class fc._EmbeddingColumn.
+    del bypass_scope_validation
+    # pylint: disable=redundant-keyword-arg
+    return fc._EmbeddingColumn.__new__(
+        cls,
+        categorical_column,
+        dimension,
+        combiner=combiner,
+        layer_creator=layer_creator,
+        ckpt_to_load_from=ckpt_to_load_from,
+        tensor_name_in_ckpt=tensor_name_in_ckpt,
+        max_norm=max_norm,
+        trainable=trainable,
+        use_safe_embedding_lookup=use_safe_embedding_lookup)
+
+  def __init__(self,
+               categorical_column,
+               dimension,
+               combiner='mean',
+               layer_creator=None,
+               ckpt_to_load_from=None,
+               tensor_name_in_ckpt=None,
+               max_norm=None,
+               trainable=True,
+               max_sequence_length=0,
+               learning_rate_fn=None,
+               use_safe_embedding_lookup=True,
+               bypass_scope_validation=False):
+    _TPUBaseEmbeddingColumn.__init__(
+        self,
+        categorical_column,
+        max_sequence_length=max_sequence_length,
+        learning_rate_fn=learning_rate_fn)
+    self._key = None
+    # If true, scope validation is skipped to allow the same column to be used
+    # in multiple variable scopes. By default, this is False, and we expect a
+    # 1:1 mapping between feature columns and scopes.
+    self._bypass_scope_validation = bypass_scope_validation
+
+  def get_combiner(self):
+    return self.combiner
+
+  def get_embedding_table_size(self):
+    """Returns num_ids and width."""
+    return (self.categorical_column._num_buckets, self.dimension)
+
+  def get_feature_key_name(self):
+    """get_feature_key_name."""
+    if self.is_categorical_column_weighted():
+      return self.categorical_column.categorical_column.name
+    return self.categorical_column.name
+
+  def get_weight_key_name(self):
+    """get_weight_key_name."""
+    if self.is_categorical_column_weighted():
+      return self.categorical_column.weight_feature_key
+    return None
+
+  def get_embedding_var_name(self):
+    """get_embedding_var_name."""
+    return self.categorical_column.name
+
+  def get_initializer(self):
+    return self._tpu_initializer
+
+  def is_categorical_column_weighted(self):
+    """Check if the categorical column of the embedding column is weighted."""
+    if isinstance(
+        self.categorical_column,
+        (
+            fc._WeightedCategoricalColumn,  # pylint: disable=protected-access
+            fc_lib.WeightedCategoricalColumn)):
+      return True
+    return False
+
+  def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None):
+    if tpu.under_tpu_inference_context():
+      def host_computation():
+        return fc._EmbeddingColumn._get_dense_tensor(
+            self, inputs, weight_collections, trainable)
+
+      return tpu_replication.outside_compilation(host_computation)
+
+    if _is_running_on_cpu():
+      return fc._EmbeddingColumn._get_dense_tensor(
+          self, inputs, weight_collections, trainable)
+
+    # TPU mode
+    # Get the embeddings from the LazyBuilder.
+    tensor = inputs.get(self.get_feature_key_name())
+
+    # Add to collection for _create_tpu_embedding_variables_and_ops
+    _record_variable_scope_and_name(
+        self.get_embedding_var_name(),
+        'embedding_weights',
+        bypass_scope_validation=self._bypass_scope_validation)
+
+    return tensor
+
+  def _get_sequence_dense_tensor(
+      self, inputs, weight_collections=None, trainable=None):
+    if tpu.under_tpu_inference_context():
+      def host_computation():
+        return fc._EmbeddingColumn._get_sequence_dense_tensor(
+            self, inputs, weight_collections, trainable)
+
+      return tpu_replication.outside_compilation(host_computation)
+
+    if _is_running_on_cpu():
+      return fc._EmbeddingColumn._get_sequence_dense_tensor(
+          self, inputs, weight_collections, trainable)
+
+    tensor = inputs.get(self.get_feature_key_name())
+    tensor_lengths = inputs.get(self.get_sequence_length_feature_key_name())
+
+    # inputs is a _LazyBuilder and for rank 1 tensors, it calls expand_dims(-1).
+    # We need to undo this to match the standard CPU sequence embedding.
+    tensor_lengths = array_ops.squeeze(tensor_lengths, -1)
+
+    # Add to collection for _create_tpu_embedding_variables_and_ops
+    _record_variable_scope_and_name(
+        self.get_embedding_var_name(),
+        'embedding_weights',
+        bypass_scope_validation=self._bypass_scope_validation)
+
+    return fc._SequenceDenseColumn.TensorSequenceLengthPair(
+        dense_tensor=tensor, sequence_length=tensor_lengths)
+
+
+class _TPUSharedEmbeddingColumn(_TPUBaseEmbeddingColumn,
+                                fc._SharedEmbeddingColumn):
+  """Core Shared Embedding Column."""
+
+  def __new__(cls,
+              categorical_column,
+              dimension,
+              combiner='mean',
+              initializer=None,
+              shared_embedding_collection_name=None,
+              ckpt_to_load_from=None,
+              tensor_name_in_ckpt=None,
+              max_norm=None,
+              trainable=True,
+              max_sequence_length=0,
+              learning_rate_fn=None,
+              use_safe_embedding_lookup=True):
+    return fc._SharedEmbeddingColumn.__new__(
+        cls,
+        categorical_column,
+        dimension,
+        combiner=combiner,
+        initializer=initializer,
+        shared_embedding_collection_name=shared_embedding_collection_name,
+        ckpt_to_load_from=ckpt_to_load_from,
+        tensor_name_in_ckpt=tensor_name_in_ckpt,
+        max_norm=max_norm,
+        trainable=trainable,
+        use_safe_embedding_lookup=use_safe_embedding_lookup)
+
+  def __init__(self,
+               categorical_column,
+               dimension,
+               combiner='mean',
+               initializer=None,
+               shared_embedding_collection_name=None,
+               ckpt_to_load_from=None,
+               tensor_name_in_ckpt=None,
+               max_norm=None,
+               trainable=True,
+               max_sequence_length=0,
+               learning_rate_fn=None,
+               use_safe_embedding_lookup=True):
+
+    _TPUBaseEmbeddingColumn.__init__(
+        self,
+        categorical_column,
+        max_sequence_length=max_sequence_length,
+        learning_rate_fn=learning_rate_fn)
+    self._key = None
+
+  def get_combiner(self):
+    return self.combiner
+
+  def get_embedding_table_size(self):
+    """Returns num_ids and width."""
+    return (self.categorical_column._num_buckets, self.dimension)
+
+  def get_feature_key_name(self):
+    """get_feature_key_name."""
+    if self.is_categorical_column_weighted():
+      return self.categorical_column.categorical_column.name
+    return self.categorical_column.name
+
+  def get_weight_key_name(self):
+    """get_weight_key_name."""
+    if self.is_categorical_column_weighted():
+      return self.categorical_column.weight_feature_key
+    return None
+
+  def get_embedding_var_name(self):
+    """get_embedding_var_name."""
+    return self.shared_embedding_collection_name
+
+  def get_initializer(self):
+    return self.initializer
+
+  def is_categorical_column_weighted(self):
+    """Check if the categorical column of the embedding column is weighted."""
+    if isinstance(
+        self.categorical_column,
+        (
+            fc._WeightedCategoricalColumn,  # pylint: disable=protected-access
+            fc_lib.WeightedCategoricalColumn)):
+      return True
+    return False
+
+  def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None):
+    if tpu.under_tpu_inference_context():
+      def host_computation():
+        return fc._SharedEmbeddingColumn._get_dense_tensor(
+            self, inputs, weight_collections, trainable)
+
+      return tpu_replication.outside_compilation(host_computation)
+
+    if _is_running_on_cpu():
+      return fc._SharedEmbeddingColumn._get_dense_tensor(
+          self, inputs, weight_collections, trainable)
+
+    # TPU mode
+    # Get the embeddings from the LazyBuilder.
+    tensor = inputs.get(self.get_feature_key_name())
+
+    # Add to collection for _create_tpu_embedding_variables_and_ops
+    _record_variable_scope_and_name(
+        self.get_embedding_var_name(),
+        'embedding_weights',
+        is_shared_embedding=True)
+    return tensor
+
+  def _get_sequence_dense_tensor(
+      self, inputs, weight_collections=None, trainable=None):
+    if tpu.under_tpu_inference_context():
+      def host_computation():
+        return fc._SharedEmbeddingColumn._get_sequence_dense_tensor(
+            self, inputs, weight_collections, trainable)
+
+      return tpu_replication.outside_compilation(host_computation)
+
+    if _is_running_on_cpu():
+      return fc._SharedEmbeddingColumn._get_sequence_dense_tensor(
+          self, inputs, weight_collections, trainable)
+
+    tensor = inputs.get(self.get_feature_key_name())
+    tensor_lengths = inputs.get(self.get_sequence_length_feature_key_name())
+
+    # Add to collection for _create_tpu_embedding_variables_and_ops
+    _record_variable_scope_and_name(
+        self.get_embedding_var_name(),
+        'embedding_weights',
+        is_shared_embedding=True)
+
+    return fc._SequenceDenseColumn.TensorSequenceLengthPair(
+        dense_tensor=tensor, sequence_length=tensor_lengths)
+
+
+def _record_variable_scope_and_name(embedding_var_name,
+                                    embedding_var_name_in_fc,
+                                    is_shared_embedding=False,
+                                    bypass_scope_validation=False):
+  """Add embedding variable name and scope to collection."""
+  g = ops.get_default_graph()
+  collection = g.get_collection_ref(_TPU_FC_TO_SCOPE)
+  if not collection:
+    collection.append({})
+
+  var_def_dict = collection[0]
+
+  captured_scope = variable_scope.get_variable_scope()
+  captured_scope_name = captured_scope.name
+
+  if embedding_var_name in var_def_dict:
+    if (var_def_dict[embedding_var_name][0] != captured_scope_name and
+        not is_shared_embedding and not bypass_scope_validation):
+      raise ValueError(
+          'For embedding var name {}, the variable scope name is different, '
+          'got {}; expected {}'.format(embedding_var_name,
+                                       captured_scope_name,
+                                       var_def_dict[embedding_var_name][0]))
+    if var_def_dict[embedding_var_name][1] != embedding_var_name_in_fc:
+      raise ValueError(
+          'For embedding var name {}, the embedding name is different, '
+          'got {}; expected {}'.format(embedding_var_name,
+                                       embedding_var_name_in_fc,
+                                       var_def_dict[embedding_var_name][1]))
+  else:
+    var_def_dict[embedding_var_name] = (captured_scope_name,
+                                        embedding_var_name_in_fc)
+
+
+def _is_running_on_cpu():
+  """Returns True if the current context is CPU model."""
+  return tpu_function.get_tpu_context().number_of_shards is None
+
+
+def get_sequence_length_feature_key_name_from_feature_key_name(feature_name):
+  """Gets the name of the sequence length feature from that of the base feature.
+
+  Args:
+    feature_name: The feature key of a sequence column.
+
+  Returns:
+    A string which is the feature key for the associated feature length column.
+  """
+  return feature_name + _SEQUENCE_FEATURE_LENGTH_POSTFIX

tensorflow/python/tpu/feature_column_test.py

@@ -0,0 +1,312 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===================================================================
+"""Tests for python.tpu.feature_column."""
+
+import numpy as np
+
+from tensorflow.python.client import session
+from tensorflow.python.feature_column import feature_column as fc
+from tensorflow.python.feature_column import feature_column_lib as fc_lib
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import variables as variables_lib
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import feature_column as tpu_fc
+
+
+def _initialized_session():
+  sess = session.Session()
+  sess.run(variables_lib.global_variables_initializer())
+  sess.run(lookup_ops.tables_initializer())
+  return sess
+
+
+class EmbeddingColumnTest(test.TestCase):
+
+  def test_defaults(self):
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    embedding_dimension = 2
+    embedding_column = tpu_fc.embedding_column(
+        categorical_column, dimension=embedding_dimension)
+    self.assertIs(categorical_column, embedding_column.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column.dimension)
+    self.assertEqual('mean', embedding_column.combiner)
+    self.assertEqual('aaa_embedding', embedding_column.name)
+    self.assertEqual('aaa_embedding', embedding_column._var_scope_name)
+    self.assertEqual((embedding_dimension,), embedding_column._variable_shape)
+    self.assertEqual({
+        'aaa': parsing_ops.VarLenFeature(dtypes.int64)
+    }, embedding_column._parse_example_spec)
+
+  def test_denylisted_column(self):
+    # HashedCategoricalColumn is denylisted and so will raise an exception.
+    categorical_column = fc_lib.categorical_column_with_hash_bucket(
+        key='aaa', hash_bucket_size=3)
+    embedding_dimension = 2
+    with self.assertRaises(TypeError):
+      tpu_fc.embedding_column(categorical_column, dimension=embedding_dimension)
+
+  def test_custom_column(self):
+    # This column is not in any allowlist but should succeed because
+    # it inherits from V2 CategoricalColumn.
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=10)
+    embedding_dimension = 2
+    embedding_column = tpu_fc.embedding_column(
+        categorical_column, dimension=embedding_dimension)
+    self.assertIs(categorical_column, embedding_column.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column.dimension)
+    self.assertEqual('mean', embedding_column.combiner)
+    self.assertEqual('aaa_embedding', embedding_column.name)
+    self.assertEqual('aaa_embedding', embedding_column._var_scope_name)
+    self.assertEqual((embedding_dimension,), embedding_column._variable_shape)
+    self.assertEqual({'aaa': parsing_ops.VarLenFeature(dtypes.int64)},
+                     embedding_column._parse_example_spec)
+
+  def test_all_constructor_args(self):
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    embedding_dimension = 2
+    embedding_column = tpu_fc.embedding_column(
+        categorical_column,
+        dimension=embedding_dimension,
+        combiner='my_combiner',
+        initializer=lambda: 'my_initializer')
+    self.assertIs(categorical_column, embedding_column.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column.dimension)
+    self.assertEqual('my_combiner', embedding_column.combiner)
+    self.assertEqual('aaa_embedding', embedding_column.name)
+    self.assertEqual('aaa_embedding', embedding_column._var_scope_name)
+    self.assertEqual((embedding_dimension,), embedding_column._variable_shape)
+    self.assertEqual({
+        'aaa': parsing_ops.VarLenFeature(dtypes.int64)
+    }, embedding_column._parse_example_spec)
+
+  @test_util.deprecated_graph_mode_only
+  def test_get_dense_tensor(self):
+    # Inputs.
+    vocabulary_size = 3
+    sparse_input = sparse_tensor.SparseTensorValue(
+        # example 0, ids [2]
+        # example 1, ids [0, 1]
+        # example 2, ids []
+        # example 3, ids [1]
+        indices=((0, 0), (1, 0), (1, 4), (3, 0)),
+        values=(2, 0, 1, 1),
+        dense_shape=(4, 5))
+
+    # Embedding variable.
+    embedding_dimension = 2
+    embedding_values = (
+        (1., 2.),  # id 0
+        (3., 5.),  # id 1
+        (7., 11.)  # id 2
+    )
+
+    def _initializer(shape, dtype, partition_info):
+      self.assertAllEqual((vocabulary_size, embedding_dimension), shape)
+      self.assertEqual(dtypes.float32, dtype)
+      self.assertIsNone(partition_info)
+      return embedding_values
+
+    # Expected lookup result, using combiner='mean'.
+    expected_lookups = (
+        # example 0, ids [2], embedding = [7, 11]
+        (7., 11.),
+        # example 1, ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5]
+        (2., 3.5),
+        # example 2, ids [], embedding = [0, 0]
+        (0., 0.),
+        # example 3, ids [1], embedding = [3, 5]
+        (3., 5.),
+    )
+
+    # Build columns.
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=vocabulary_size)
+    embedding_column = tpu_fc.embedding_column(
+        categorical_column,
+        dimension=embedding_dimension,
+        initializer=_initializer)
+
+    # Provide sparse input and get dense result.
+    embedding_lookup = embedding_column._get_dense_tensor(
+        fc._LazyBuilder({
+            'aaa': sparse_input
+        }))
+
+    # Assert expected embedding variable and lookups.
+    global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+    self.assertItemsEqual(('embedding_weights:0',),
+                          tuple([v.name for v in global_vars]))
+    with _initialized_session():
+      self.assertAllEqual(embedding_values, global_vars[0])
+      self.assertAllEqual(expected_lookups, embedding_lookup)
+
+
+class SharedEmbeddingColumnTest(test.TestCase):
+
+  @test_util.deprecated_graph_mode_only
+  def test_defaults(self):
+    categorical_column_a = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    categorical_column_b = fc_lib.categorical_column_with_identity(
+        key='bbb', num_buckets=3)
+    embedding_dimension = 2
+    embedding_column_b, embedding_column_a = tpu_fc.shared_embedding_columns(
+        [categorical_column_b, categorical_column_a],
+        dimension=embedding_dimension)
+    self.assertIs(categorical_column_a, embedding_column_a.categorical_column)
+    self.assertIs(categorical_column_b, embedding_column_b.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column_a.dimension)
+    self.assertEqual(embedding_dimension, embedding_column_b.dimension)
+    self.assertEqual('mean', embedding_column_a.combiner)
+    self.assertEqual('mean', embedding_column_b.combiner)
+    self.assertIsNotNone(embedding_column_a.initializer)
+    self.assertIsNotNone(embedding_column_b.initializer)
+    self.assertEqual('aaa_bbb_shared_embedding',
+                     embedding_column_a.shared_embedding_collection_name)
+    self.assertEqual('aaa_bbb_shared_embedding',
+                     embedding_column_b.shared_embedding_collection_name)
+    self.assertEqual('aaa_shared_embedding', embedding_column_a.name)
+    self.assertEqual('bbb_shared_embedding', embedding_column_b.name)
+    self.assertEqual('aaa_bbb_shared_embedding',
+                     embedding_column_a._var_scope_name)
+    self.assertEqual('aaa_bbb_shared_embedding',
+                     embedding_column_b._var_scope_name)
+    self.assertEqual((embedding_dimension,), embedding_column_a._variable_shape)
+    self.assertEqual((embedding_dimension,), embedding_column_b._variable_shape)
+    self.assertEqual({
+        'aaa': parsing_ops.VarLenFeature(dtypes.int64)
+    }, embedding_column_a._parse_example_spec)
+    self.assertEqual({
+        'bbb': parsing_ops.VarLenFeature(dtypes.int64)
+    }, embedding_column_b._parse_example_spec)
+
+  @test_util.deprecated_graph_mode_only
+  def test_all_constructor_args(self):
+    categorical_column_a = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    categorical_column_b = fc_lib.categorical_column_with_identity(
+        key='bbb', num_buckets=3)
+    embedding_dimension = 2
+    embedding_column_a, embedding_column_b = tpu_fc.shared_embedding_columns(
+        [categorical_column_a, categorical_column_b],
+        dimension=embedding_dimension,
+        combiner='my_combiner',
+        initializer=lambda: 'my_initializer',
+        shared_embedding_collection_name='var_scope_name')
+    self.assertIs(categorical_column_a, embedding_column_a.categorical_column)
+    self.assertIs(categorical_column_b, embedding_column_b.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column_a.dimension)
+    self.assertEqual(embedding_dimension, embedding_column_b.dimension)
+    self.assertEqual('my_combiner', embedding_column_a.combiner)
+    self.assertEqual('my_combiner', embedding_column_b.combiner)
+    self.assertEqual('my_initializer', embedding_column_a.initializer())
+    self.assertEqual('my_initializer', embedding_column_b.initializer())
+    self.assertEqual('var_scope_name',
+                     embedding_column_a.shared_embedding_collection_name)
+    self.assertEqual('var_scope_name',
+                     embedding_column_b.shared_embedding_collection_name)
+    self.assertEqual('aaa_shared_embedding', embedding_column_a.name)
+    self.assertEqual('bbb_shared_embedding', embedding_column_b.name)
+    self.assertEqual('var_scope_name', embedding_column_a._var_scope_name)
+    self.assertEqual('var_scope_name', embedding_column_b._var_scope_name)
+    self.assertEqual((embedding_dimension,), embedding_column_a._variable_shape)
+    self.assertEqual((embedding_dimension,), embedding_column_b._variable_shape)
+    self.assertEqual({
+        'aaa': parsing_ops.VarLenFeature(dtypes.int64)
+    }, embedding_column_a._parse_example_spec)
+    self.assertEqual({
+        'bbb': parsing_ops.VarLenFeature(dtypes.int64)
+    }, embedding_column_b._parse_example_spec)
+
+  @test_util.deprecated_graph_mode_only
+  def test_get_dense_tensor(self):
+    # Inputs.
+    vocabulary_size = 3
+    # -1 values are ignored.
+    input_a = np.array([
+        [2, -1, -1],  # example 0, ids [2]
+        [0, 1, -1]
+    ])  # example 1, ids [0, 1]
+    input_b = np.array([
+        [0, -1, -1],  # example 0, ids [0]
+        [-1, -1, -1]
+    ])  # example 1, ids []
+    input_features = {'aaa': input_a, 'bbb': input_b}
+
+    # Embedding variable.
+    embedding_dimension = 2
+    embedding_values = (
+        (1., 2.),  # id 0
+        (3., 5.),  # id 1
+        (7., 11.)  # id 2
+    )
+
+    def _initializer(shape, dtype, partition_info):
+      self.assertAllEqual((vocabulary_size, embedding_dimension), shape)
+      self.assertEqual(dtypes.float32, dtype)
+      self.assertIsNone(partition_info)
+      return embedding_values
+
+    # Expected lookup result, using combiner='mean'.
+    expected_lookups_a = (
+        # example 0:
+        (7., 11.),  # ids [2], embedding = [7, 11]
+        # example 1:
+        (2., 3.5),  # ids [0, 1], embedding = mean([1, 2] + [3, 5]) = [2, 3.5]
+    )
+    expected_lookups_b = (
+        # example 0:
+        (1., 2.),  # ids [0], embedding = [1, 2]
+        # example 1:
+        (0., 0.),  # ids [], embedding = [0, 0]
+    )
+
+    # Build columns.
+    categorical_column_a = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=vocabulary_size)
+    categorical_column_b = fc_lib.categorical_column_with_identity(
+        key='bbb', num_buckets=vocabulary_size)
+    embedding_column_a, embedding_column_b = tpu_fc.shared_embedding_columns(
+        [categorical_column_a, categorical_column_b],
+        dimension=embedding_dimension,
+        initializer=_initializer)
+
+    # Provide sparse input and get dense result.
+    embedding_lookup_a = embedding_column_a._get_dense_tensor(
+        fc._LazyBuilder(input_features))
+    embedding_lookup_b = embedding_column_b._get_dense_tensor(
+        fc._LazyBuilder(input_features))
+
+    # Assert expected embedding variable and lookups.
+    global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+    self.assertItemsEqual(('embedding_weights:0',),
+                          tuple([v.name for v in global_vars]))
+    embedding_var = global_vars[0]
+    with _initialized_session():
+      self.assertAllEqual(embedding_values, embedding_var)
+      self.assertAllEqual(expected_lookups_a, embedding_lookup_a)
+      self.assertAllEqual(expected_lookups_b, embedding_lookup_b)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/tpu/feature_column_v2_test.py

@@ -0,0 +1,287 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===================================================================
+"""Tests for python.tpu.feature_column."""
+
+import copy
+
+from absl.testing import parameterized
+
+from tensorflow.python.client import session
+from tensorflow.python.feature_column import feature_column_lib as fc_lib
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import lookup_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables as variables_lib
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import feature_column_v2 as tpu_fc
+from tensorflow.python.tpu import tpu_function
+
+
+def _initialized_session():
+  sess = session.Session()
+  sess.run(variables_lib.global_variables_initializer())
+  sess.run(lookup_ops.tables_initializer())
+  return sess
+
+
+class _TestStateManager(fc_lib.StateManager):
+
+  def __init__(self, trainable=True):
+    self._all_variables = {}
+    self._trainable = trainable
+
+  def create_variable(self,
+                      feature_column,
+                      name,
+                      shape,
+                      dtype=None,
+                      trainable=True,
+                      use_resource=True,
+                      initializer=None):
+    if feature_column not in self._all_variables:
+      self._all_variables[feature_column] = {}
+    var_dict = self._all_variables[feature_column]
+    if name in var_dict:
+      return var_dict[name]
+    else:
+      var = variable_scope.get_variable(
+          name=name,
+          shape=shape,
+          dtype=dtype,
+          trainable=self._trainable and trainable,
+          use_resource=use_resource,
+          initializer=initializer)
+      var_dict[name] = var
+      return var
+
+  def get_variable(self, feature_column, name):
+    return self._all_variables[feature_column][name]
+
+
+class EmbeddingColumnTestV2(test.TestCase, parameterized.TestCase):
+
+  def test_defaults(self):
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    embedding_dimension = 2
+    embedding_column = tpu_fc.embedding_column_v2(
+        categorical_column, dimension=embedding_dimension)
+    # Can't test default initializer as it's a random function.
+    self.assertIs(categorical_column, embedding_column.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column.dimension)
+    self.assertEqual('mean', embedding_column.combiner)
+    self.assertEqual('aaa_embedding', embedding_column.name)
+    self.assertEqual((embedding_dimension,), embedding_column.variable_shape)
+
+  def test_all_constructor_args(self):
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    embedding_dimension = 2
+    embedding_column = tpu_fc.embedding_column_v2(
+        categorical_column,
+        dimension=embedding_dimension,
+        combiner='my_combiner',
+        initializer=lambda: 'my_initializer')
+    self.assertIs(categorical_column, embedding_column.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column.dimension)
+    self.assertEqual('my_combiner', embedding_column.combiner)
+    self.assertEqual('my_initializer', embedding_column.initializer())
+    self.assertEqual('aaa_embedding', embedding_column.name)
+    self.assertEqual((embedding_dimension,), embedding_column.variable_shape)
+    self.assertEqual({
+        'aaa': parsing_ops.VarLenFeature(dtypes.int64)
+    }, embedding_column._parse_example_spec)
+
+  def test_deepcopy(self):
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    embedding_column = tpu_fc.embedding_column_v2(
+        categorical_column, dimension=2)
+    embedding_column_copy = copy.deepcopy(embedding_column)
+    self.assertEqual(embedding_column.dimension,
+                     embedding_column_copy.dimension)
+    self.assertEqual(embedding_column._max_sequence_length,
+                     embedding_column_copy._max_sequence_length)
+
+  def test_with_scope_validation(self):
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    embedding_dimension = 2
+    initializer = init_ops.truncated_normal_initializer(mean=0.0, stddev=.5)
+    embedding_column = tpu_fc._TPUEmbeddingColumnV2(
+        categorical_column=categorical_column,
+        dimension=embedding_dimension,
+        combiner='mean',
+        initializer=initializer,
+        max_sequence_length=0,
+        learning_rate_fn=None,
+        use_safe_embedding_lookup=True,
+        bypass_scope_validation=False)
+    self.assertIs(categorical_column, embedding_column.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column.dimension)
+    state_manager = _TestStateManager()
+    with tpu_function.tpu_shard_context(1):
+      with variable_scope.variable_scope('tower1/scope1'):
+        embedding_column.create_state(state_manager)
+      with variable_scope.variable_scope('tower2/scope2'):
+        # With default scope validation, the same column cannot be used in a new
+        # variable scope.
+        with self.assertRaisesRegex(ValueError,
+                                    'the variable scope name is different'):
+          embedding_column.create_state(state_manager)
+
+  def test_bypass_scope_validation(self):
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    embedding_dimension = 2
+    initializer = init_ops.truncated_normal_initializer(mean=0.0, stddev=.5)
+    embedding_column = tpu_fc._TPUEmbeddingColumnV2(
+        categorical_column=categorical_column,
+        dimension=embedding_dimension,
+        combiner='mean',
+        initializer=initializer,
+        max_sequence_length=0,
+        learning_rate_fn=None,
+        use_safe_embedding_lookup=True,
+        bypass_scope_validation=True)
+    self.assertIs(categorical_column, embedding_column.categorical_column)
+    self.assertEqual(embedding_dimension, embedding_column.dimension)
+    state_manager = _TestStateManager()
+    with tpu_function.tpu_shard_context(1):
+      with variable_scope.variable_scope('tower1/scope1'):
+        embedding_column.create_state(state_manager)
+      with variable_scope.variable_scope('tower2/scope2'):
+        embedding_column.create_state(state_manager)
+
+  def test_deepcopy_with_bypass_scope_validation(self):
+    categorical_column = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=3)
+    embedding_dimension = 2
+    initializer = init_ops.truncated_normal_initializer(mean=0.0, stddev=.5)
+    embedding_column = tpu_fc._TPUEmbeddingColumnV2(
+        categorical_column=categorical_column,
+        dimension=embedding_dimension,
+        combiner='mean',
+        initializer=initializer,
+        max_sequence_length=0,
+        use_safe_embedding_lookup=False,
+        bypass_scope_validation=True)
+    embedding_column_copy = copy.deepcopy(embedding_column)
+    self.assertEqual(embedding_dimension, embedding_column_copy.dimension)
+    self.assertEqual(embedding_column._max_sequence_length,
+                     embedding_column_copy._max_sequence_length)
+    self.assertTrue(embedding_column_copy._bypass_scope_validation)
+    self.assertFalse(embedding_column_copy.use_safe_embedding_lookup)
+
+
+class SharedEmbeddingColumnTestV2(test.TestCase, parameterized.TestCase):
+
+  @test_util.deprecated_graph_mode_only
+  def test_defaults(self):
+    vocabulary_size = 3
+    categorical_column_a = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=vocabulary_size)
+    categorical_column_b = fc_lib.categorical_column_with_identity(
+        key='bbb', num_buckets=vocabulary_size)
+    embedding_dimension = 2
+    embedding_column_b, embedding_column_a = tpu_fc.shared_embedding_columns_v2(
+        [categorical_column_b, categorical_column_a],
+        dimension=embedding_dimension)
+    self.assertIs(categorical_column_a, embedding_column_a.categorical_column)
+    self.assertIs(categorical_column_b, embedding_column_b.categorical_column)
+    self.assertEqual((vocabulary_size, embedding_dimension),
+                     embedding_column_a.get_embedding_table_size())
+    self.assertEqual((vocabulary_size, embedding_dimension),
+                     embedding_column_a.get_embedding_table_size())
+    self.assertEqual('mean', embedding_column_a.combiner)
+    self.assertEqual('mean', embedding_column_b.combiner)
+    self.assertIsNotNone(embedding_column_a.get_initializer())
+    self.assertIsNotNone(embedding_column_b.get_initializer())
+    self.assertEqual('aaa_bbb_shared_embedding',
+                     embedding_column_a.get_embedding_var_name())
+    self.assertEqual('aaa_bbb_shared_embedding',
+                     embedding_column_b.get_embedding_var_name())
+    self.assertEqual('aaa_shared_embedding', embedding_column_a.name)
+    self.assertEqual('bbb_shared_embedding', embedding_column_b.name)
+    self.assertEqual((embedding_dimension,), embedding_column_a.variable_shape)
+    self.assertEqual((embedding_dimension,), embedding_column_b.variable_shape)
+
+  @test_util.deprecated_graph_mode_only
+  def test_all_constructor_args(self):
+    vocabulary_size = 3
+    categorical_column_a = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=vocabulary_size)
+    categorical_column_b = fc_lib.categorical_column_with_identity(
+        key='bbb', num_buckets=vocabulary_size)
+    embedding_dimension = 2
+    embedding_column_a, embedding_column_b = tpu_fc.shared_embedding_columns_v2(
+        [categorical_column_a, categorical_column_b],
+        dimension=embedding_dimension,
+        combiner='my_combiner',
+        initializer=lambda: 'my_initializer',
+        shared_embedding_collection_name='var_scope_name')
+    self.assertIs(categorical_column_a, embedding_column_a.categorical_column)
+    self.assertIs(categorical_column_b, embedding_column_b.categorical_column)
+    self.assertEqual((vocabulary_size, embedding_dimension),
+                     embedding_column_a.get_embedding_table_size())
+    self.assertEqual((vocabulary_size, embedding_dimension),
+                     embedding_column_a.get_embedding_table_size())
+    self.assertEqual('my_combiner', embedding_column_a.combiner)
+    self.assertEqual('my_combiner', embedding_column_b.combiner)
+    self.assertEqual('my_initializer', embedding_column_a.get_initializer()())
+    self.assertEqual('my_initializer', embedding_column_b.get_initializer()())
+    self.assertEqual('var_scope_name',
+                     embedding_column_a.get_embedding_var_name())
+    self.assertEqual('var_scope_name',
+                     embedding_column_b.get_embedding_var_name())
+    self.assertEqual('aaa_shared_embedding', embedding_column_a.name)
+    self.assertEqual('bbb_shared_embedding', embedding_column_b.name)
+    self.assertEqual((embedding_dimension,), embedding_column_a.variable_shape)
+    self.assertEqual((embedding_dimension,), embedding_column_b.variable_shape)
+
+  def test_deepcopy(self):
+    vocabulary_size = 3
+    categorical_column_a = fc_lib.categorical_column_with_identity(
+        key='aaa', num_buckets=vocabulary_size)
+    categorical_column_b = fc_lib.categorical_column_with_identity(
+        key='bbb', num_buckets=vocabulary_size)
+    embedding_dimension = 2
+    columns = tpu_fc.shared_embedding_columns_v2(
+        [categorical_column_b, categorical_column_a],
+        dimension=embedding_dimension)
+    columns_copy = copy.deepcopy(columns)
+    self.assertEqual(
+        [column._shared_embedding_collection_name for column in columns],
+        [column._shared_embedding_collection_name for column in columns_copy])
+
+
+class DeviceSpecificEmbeddingColumnTestV2(test.TestCase,
+                                          parameterized.TestCase):
+
+  @test_util.deprecated_graph_mode_only
+  def test_error_dense_shape_invalid(self):
+    categorical_column_input = fc_lib.categorical_column_with_identity(
+        key='inp', num_buckets=5)
+    with self.assertRaisesRegex(ValueError, 'tensor_core_shape must be size 2'):
+      tpu_fc.shared_embedding_columns_v2([categorical_column_input],
+                                         dimension=20,
+                                         tensor_core_shape=[None, 20, 15])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/tpu/functional.py

@@ -0,0 +1,19 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Functional operations."""
+
+from tensorflow.python.tpu.ops import tpu_ops
+
+TPUPartitionedCall = tpu_ops.tpu_partitioned_call  # pylint: disable=invalid-name

tensorflow/python/tpu/ops/BUILD

@@ -0,0 +1,98 @@
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "ops",
+    srcs = ["tpu_ops.py"],
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:tpu_ops_gen",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/tpu:tpu_function",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+tf_gen_op_wrapper_py(
+    name = "gen_tpu_embedding_ops",
+    out = "gen_tpu_embedding_ops.py",
+    extra_py_deps = [
+        "//tensorflow/python:pywrap_tfe",
+        "//tensorflow/python/util:dispatch",
+        "//tensorflow/python/util:deprecation",
+        "//tensorflow/python/util:tf_export",
+    ],
+    py_lib_rule = pytype_strict_library,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core/ops:tpu_configuration_ops_op_lib",
+        "//tensorflow/core/ops:tpu_embedding_ops_op_lib",
+        "//tensorflow/core/tpu/ops:tpu_embedding_ops",
+    ],
+)
+
+tf_gen_op_wrapper_py(
+    name = "gen_xla_ops",
+    out = "gen_xla_ops.py",
+    api_def_srcs = [
+        "//tensorflow/core/api_def:base_api_def",
+        "//tensorflow/core/api_def:python_api_def",
+    ],
+    op_allowlist = [
+        "ConvertToCooTensor",
+        "GetMinibatchesInCsrWithPhysicalReplica",
+        "GetMinibatchSplitsWithPhysicalReplica",
+        "StoreMinibatchStatisticsInFdo",
+        "ConvertToListOfSparseCoreCooTensors",
+        "SortListOfSparseCoreCooTensors",
+        "ConvertToSparseCoreCsrWrappedCooTensor",
+        "GetStatsFromListOfSparseCoreCooTensors",
+        "GlobalIterId",
+        "TPUCopyWithDynamicShape",
+        "TPUAnnotateTensorsWithDynamicShape",
+        "XlaSparseDenseMatmul",
+        "XlaSparseDenseMatmulWithCsrInput",
+        "XlaSparseDenseMatmulIntWithCsrInput",
+        "XlaSparseDenseMatmulCustomCombinerOnTcWithCsrInput",
+        "XlaSparseDenseMatmulGrad",
+        "XlaSparseDenseMatmulIntGradWithCsrInput",
+        "XlaSparseDenseMatmulGradWithCsrInput",
+        "XlaSparseDenseMatmulCustomCombinerOnTcGradWithSgdAndCsrInput",
+        "XlaSparseDenseMatmulCustomCombinerOnTcGradWithAdagradAndCsrInput",
+        "XlaSparseDenseMatmulCustomCombinerOnTcGradWithAdagradMomentumAndCsrInput",
+        "XlaSparseDenseMatmulCustomCombinerOnTcGradWithAdamAndCsrInput",
+        "XlaSparseDenseMatmulCustomCombinerOnTcGradWithFtrlAndCsrInput",
+        "XlaSparseDenseMatmulCustomCombinerOnTcGradWithCsrInput",
+        "XlaSparseDenseMatmulGradWithSgdAndCsrInput",
+        "XlaSparseDenseMatmulGradWithAdagradAndCsrInput",
+        "XlaSparseDenseMatmulGradWithAdagradMomentumAndCsrInput",
+        "XlaSparseDenseMatmulGradWithAdamAndCsrInput",
+        "XlaSparseDenseMatmulGradWithFtrlAndCsrInput",
+        "XlaSparseDenseMatmulWithStaticBufferSize",
+        "XlaSparseDenseMatmulGradWithSgdAndStaticBufferSize",
+        "XlaSparseDenseMatmulGradWithAdagradAndStaticBufferSize",
+        "XlaSparseDenseMatmulGradWithAdagradMomentumAndStaticBufferSize",
+        "XlaSparseDenseMatmulGradWithAdamAndStaticBufferSize",
+        "XlaSparseDenseMatmulGradWithFtrlAndStaticBufferSize",
+        "XlaSparseCoreSgd",
+        "XlaSparseCoreAdagrad",
+        "XlaSparseCoreAdagradMomentum",
+        "XlaSparseCoreAdam",
+        "XlaSparseCoreFtrl",
+    ],
+    deps = [
+        "//tensorflow/core/tpu/ops:sparse_core_ops",
+        "//tensorflow/core/tpu/ops:sparse_core_preprocess_ops",
+        "//tensorflow/core/tpu/ops:tpu_copy_with_dynamic_shape_op",
+    ],
+)

tensorflow/python/tpu/ops/tpu_ops.py

@@ -0,0 +1,604 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Operations for TPUs."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+# pylint: disable=wildcard-import,unused-import
+from tensorflow.python.ops import gen_tpu_ops
+from tensorflow.python.ops.gen_tpu_ops import *
+# pylint: enable=wildcard-import,unused-import
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.tpu import tpu_function
+from tensorflow.python.util.tf_export import tf_export
+
+
+ops.NotDifferentiable("TPUReplicatedInput")
+
+
+def _create_default_group_assignment():
+  num_shards = tpu_function.get_tpu_context().number_of_shards
+  if num_shards is None:
+    logging.warning(
+        "cross_replica_sum should be used within a tpu_shard_context, but "
+        "got unset number_of_shards. Assuming 1.")
+    num_shards = 1
+  group_assignment = [list(range(num_shards))]
+  return group_assignment
+
+
+def all_to_all(x,
+               concat_dimension,
+               split_dimension,
+               split_count,
+               group_assignment=None,
+               name=None):
+  """Exchange data across TPU replicas.
+
+  Args:
+    x: The local tensor.
+    concat_dimension: The dimension number to concatenate.
+    split_dimension: The dimension number to split.
+    split_count: The number of splits, this number must equal to the sub-group
+      size(group_assignment.get_shape()[1])
+    group_assignment: Optional 2d int32 lists with shape [num_groups,
+      num_replicas_per_group]. `group_assignment[i]` represents the replica ids
+      in the ith subgroup.
+    name: Optional op name.
+
+  Returns:
+    A `Tensor` which is concatenated by data from different replicas.
+  """
+  if group_assignment is None:
+    group_assignment = _create_default_group_assignment()
+  return gen_tpu_ops.all_to_all(
+      x,
+      group_assignment,
+      concat_dimension=concat_dimension,
+      split_dimension=split_dimension,
+      split_count=split_count,
+      name=name)
+
+
+@ops.RegisterGradient("AllToAll")
+def _all_to_all_grad(op, grad):
+  # The gradient of a all-to-all is also a all-to-all but the
+  # split_dimension and concat_dimension is swapped.
+  # The gradient with respect to group_assignment is None.
+  return [
+      gen_tpu_ops.all_to_all(
+          grad,
+          op.inputs[1],
+          concat_dimension=op.get_attr("split_dimension"),
+          split_dimension=op.get_attr("concat_dimension"),
+          split_count=op.get_attr("split_count")), None
+  ]
+
+
+@tf_export(v1=["tpu.cross_replica_sum"])
+def cross_replica_sum(x, group_assignment=None, name=None):
+  """Sum the input tensor across replicas according to group_assignment.
+
+  Args:
+    x: The local tensor to the sum.
+    group_assignment: Optional 2d int32 lists with shape [num_groups,
+      num_replicas_per_group]. `group_assignment[i]` represents the replica ids
+      in the ith subgroup.
+    name: Optional op name.
+
+  Returns:
+    A `Tensor` which is summed across replicas.
+  """
+  if group_assignment is None:
+    group_assignment = _create_default_group_assignment()
+
+  return gen_tpu_ops.cross_replica_sum(x, group_assignment, name=name)
+
+
+def collective_permute(x, source_target_pairs, name=None):
+  """Permute the input tensor across replicas given source_target_pairs.
+
+  For each source_target_pair <a, b>, we send replica a's input to replica b.
+  Each replica id must only appear once in the source column. Also it must
+  only appear once in the target column.
+  For the replica id not in the target column, this op returns a zero tensor
+  with the same shape and dtype of the input x.
+
+  For example, suppose there are 4 TPU instances: `[A, B, C, D]`. Passing
+  source_target_pairs=`[[0,1],[1,2],[2,3]]` gets the outputs:
+  `[0, A, B, C]`.
+
+  Args:
+    x: The local tensor to be permuted.
+    source_target_pairs: 2d int lists with shape [num_pairs, 2].
+      source_target_pairs[i][0] represents the source replica id and
+      source_target_pairs[i][1] represents the target replica id.
+    name: Optional op name.
+
+  Returns:
+    A `Tensor` which is permuted.
+  """
+  return gen_tpu_ops.collective_permute(x, source_target_pairs, name=name)
+
+
+@ops.RegisterGradient("CollectivePermute")
+def _collective_permute_grad(op, grad):
+  # The gradient of a collective permute operation is also a collective
+  # permute, but with source/target pairs reversed. The gradient with respect
+  # to input argument `source_target_pairs` is `None`.
+  source_target_pairs = op.inputs[1][:, ::-1]
+  return [gen_tpu_ops.collective_permute(grad, source_target_pairs), None]
+
+
+@ops.RegisterGradient("CrossReplicaSum")
+def _cross_replica_sum_grad(op, grad):
+  # The gradient of a cross replica sum is also a cross-replica sum.
+  # The gradient with respect to group_assignment is None.
+  return [gen_tpu_ops.cross_replica_sum(grad, op.inputs[1]), None]
+
+
+# This extra type checking exists to give a more helpful error message.
+_SUPPORTED_INFEED_DTYPES = frozenset([
+    dtypes.bool, dtypes.int32, dtypes.int64, dtypes.bfloat16, dtypes.float32,
+    dtypes.complex64, dtypes.uint32, dtypes.uint8, dtypes.int8
+])
+
+
+@ops.RegisterGradient("TPUEmbeddingActivations")
+def _embedding_activations_grad(activations_op, grad_wrt_activations):
+  """Saves the gradient of embedding activations ops in a graph collection."""
+  g = ops.get_default_graph()
+  table_id = activations_op.get_attr("table_id")
+  lookup_id = activations_op.get_attr("lookup_id")
+  table_gradients = g.get_collection_ref("tpu_embedding_gradients_table_%d" %
+                                         table_id)
+
+  if not table_gradients:
+    raise RuntimeError(
+        "Gradients for TPUEmbedding have been generated in non-training mode."
+        "This is not expected. Consider putting your Optimizer.minimize code "
+        "behind the training mode condition check\n")
+
+  if lookup_id < 0 or lookup_id >= len(table_gradients):
+    raise RuntimeError(
+        "Gradients (w.r.t. TPUEmbedding activations) generated for table_id {} "
+        "and lookup_id {}. The lookup_id attribute is outside the expected "
+        "range [0, {}).".format(table_id, lookup_id, len(table_gradients)))
+
+  if table_gradients[lookup_id] is not None:
+    raise RuntimeError(
+        "Duplicate gradients (w.r.t. TPUEmbedding activations) generated for "
+        "table_id {} and lookup_id {}. This happens when there are multiple "
+        "calls to tf.gradients in a graph containing TPU embeddings. "
+        "TF cannot identify which gradient to use for updating the embedding "
+        "variables. Consider placing tf.StopGradient around tensors where "
+        "variable update is not required. Previous gradients were generated by "
+        "the following callstack: {}.".format(
+            table_id, lookup_id, table_gradients[lookup_id].op.traceback))
+
+  table_gradients[lookup_id] = array_ops.identity(grad_wrt_activations)
+  return [
+      # RegisterGradient requires that value be returned for all inputs. Since
+      # the first argument (tpu_gradient_variable_{table_name}) has shape [1],
+      # we will return zeros(shape=[1]). The actual gradient w.r.t. the
+      # embedding activations (grad_wrt_activations) has the same shape as the
+      # activations returned by  embedding_activations.
+      array_ops.zeros(arg.shape, dtype=dtypes.float32)
+      for arg in activations_op.inputs
+  ]
+
+
+def infeed_dequeue(dtype, shape, name=None):
+  """A placeholder op for a value that will be fed into the computation.
+
+  Args:
+    dtype: A `tf.DType`. The type of elements in the tensor.
+    shape: A `tf.TensorShape` or list of `ints`. The shape of the tensor.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `Tensor` of type `dtype`.
+    A tensor that will be provided using the infeed mechanism.
+
+  Raises:
+    TypeError: If 'dtype` is not a supported infeed type.
+  """
+  if dtype not in _SUPPORTED_INFEED_DTYPES:
+    raise TypeError(
+        "Operation '{}' has type {} which is not a supported TPU infeed type. "
+        "Supported types are: {}".format(name, dtype,
+                                         list(_SUPPORTED_INFEED_DTYPES)))
+
+  return gen_tpu_ops.infeed_dequeue(dtype, shape, name=name)
+
+
+# pylint: disable=redefined-outer-name
+def infeed_dequeue_tuple(dtypes, shapes, name=None):
+  """A placeholder op for values fed into the TPU simultaneously as a tuple.
+
+  Args:
+    dtypes: A list of `tf.DType`s that has length `>= 1`. The element types of
+      each element in `outputs`.
+    shapes: A list of shapes (each a `tf.TensorShape` or list of `ints`). The
+      shapes of each tensor in `outputs`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A list of `Tensor` objects of type `dtypes`.
+    A list of tensors that will be provided using the infeed mechanism.
+
+  Raises:
+    TypeError: If a type in 'dtypes` is not a supported infeed type.
+  """
+  for dtype in dtypes:
+    if dtype not in _SUPPORTED_INFEED_DTYPES:
+      raise TypeError(
+          "{} is not a supported TPU infeed type. Supported types are: "
+          "{}".format(dtype, list(_SUPPORTED_INFEED_DTYPES)))
+  return gen_tpu_ops.infeed_dequeue_tuple(dtypes, shapes, name=name)
+
+
+# pylint: enable=redefined-outer-name
+
+
+# pylint: disable=protected-access
+def send_tpu_embedding_gradients(inputs,
+                                 config,
+                                 learning_rates=None,
+                                 name=None):
+  """A placeholder op for feeding per-sample gradients to the embedding layer.
+
+  Args:
+    inputs: A TensorList of gradients with which to update embedding tables.
+      This argument has the same length and shapes as the return value of
+      RecvTPUEmbeddingActivations, but contains gradients of the model's loss
+      with respect to the embedding activations. The embedding tables are
+      updated from these gradients via the optimizers specified in the TPU
+      embedding configuration given to tpu.initialize_system.
+    config: Serialized TPUEmbeddingConfiguration proto.
+    learning_rates: A TensorList of float32 scalars, one for each dynamic
+        learning rate tag: see the comments in
+          //third_party/tensorflow/core/protobuf/tpu/
+          optimization_parameters.proto. Multiple tables can share the same
+          dynamic learning rate tag as specified in the configuration. If the
+          learning rates for all tables are constant, this list should be empty.
+    name: A name for the operation (optional).
+
+  Returns:
+    A SendTPUEmbeddingGradients operation.
+  """
+  if learning_rates is None:
+    learning_rates = []
+  return gen_tpu_ops.send_tpu_embedding_gradients(
+      inputs=inputs, learning_rates=learning_rates, config=config, name=name)
+
+
+send_tpu_embedding_gradients.__doc__ = (
+    gen_tpu_ops.send_tpu_embedding_gradients.__doc__)
+
+
+# pylint: disable=protected-access
+def enqueue_tpu_embedding_integer_batch(batch,
+                                        device_ordinal,
+                                        mode_override=None,
+                                        name=None):
+  """A placeholder op for enqueueing embedding IDs to the TPU.
+
+  Args:
+    batch: A list of 1D tensors, one for each embedding table, containing the
+      indices into the tables.
+    device_ordinal: The TPU device to use. Should be >= 0 and less than the
+      number of TPU cores in the task on which the node is placed.
+    mode_override: A string input that overrides the mode specified in the
+      TPUEmbeddingConfiguration. Supported values are {'unspecified',
+      'inference', 'train', 'backward_pass_only'}. When set to 'unspecified',
+      the mode set in TPUEmbeddingConfiguration is used, otherwise mode_override
+      is used (optional).
+    name: A name for the operation (optional).
+
+  Returns:
+    An EnqueueTPUEmbeddingIntegerBatch operation.
+  """
+  if mode_override is None:
+    mode_override = "unspecified"
+  return gen_tpu_ops.enqueue_tpu_embedding_integer_batch(
+      batch=batch,
+      device_ordinal=device_ordinal,
+      mode_override=mode_override,
+      name=name)
+
+
+enqueue_tpu_embedding_integer_batch.__doc__ = (
+    gen_tpu_ops.enqueue_tpu_embedding_integer_batch.__doc__)
+
+
+# pylint: disable=protected-access
+def enqueue_tpu_embedding_sparse_batch(sample_indices,
+                                       embedding_indices,
+                                       aggregation_weights,
+                                       device_ordinal,
+                                       combiners=None,
+                                       mode_override=None,
+                                       name=None):
+  """A placeholder op for enqueueing embedding IDs to the TPU.
+
+  Args:
+    sample_indices: A list of rank 1 Tensors specifying the training example and
+      feature to which the corresponding embedding_indices and
+      aggregation_weights values belong. sample_indices[i] must equal b * nf +
+      f, where nf is the number of features from the corresponding table, f is
+      in [0, nf), and b is in [0, batch size). Both int32 and int64 are allowed,
+      and will be converted to int32 internally.
+    embedding_indices: A list of rank 1 Tensors, indices into the embedding
+      tables. Both int32 and int64 are allowed and will be converted to int32
+      internally.
+    aggregation_weights: A list of rank 1 Tensors containing per sample -- i.e.,
+      per (training example, feature) -- aggregation weights. Both float32 and
+      float64 are allowed and will be converted to float32 internally.
+    device_ordinal: The TPU device to use. Should be >= 0 and less than the
+      number of TPU cores in the task on which the node is placed.
+    combiners: A list of string scalars, one for each embedding table that
+      specify how to normalize the embedding activations after weighted
+      summation. Supported combiners are 'mean', 'sum', or 'sqrtn'. It is
+      invalid to have the sum of the weights be 0 for 'mean' or the sum of the
+      squared weights be 0 for 'sqrtn'. If combiners isn't passed, the default
+      is to use 'sum' for all tables (optional).
+    mode_override: A string input that overrides the mode specified in the
+      TPUEmbeddingConfiguration. Supported values are {'unspecified',
+      'inference', 'train', 'backward_pass_only'}. When set to 'unspecified',
+      the mode set in TPUEmbeddingConfiguration is used, otherwise mode_override
+      is used (optional).
+    name: A name for the operation (optional).
+
+  Returns:
+    An EnqueueTPUEmbeddingSparseBatch operation.
+  """
+  if mode_override is None:
+    mode_override = "unspecified"
+  return gen_tpu_ops.enqueue_tpu_embedding_sparse_batch(
+      sample_indices=sample_indices,
+      embedding_indices=embedding_indices,
+      aggregation_weights=aggregation_weights,
+      device_ordinal=device_ordinal,
+      combiners=combiners,
+      mode_override=mode_override,
+      name=name)
+
+
+enqueue_tpu_embedding_sparse_batch.__doc__ = (
+    gen_tpu_ops.enqueue_tpu_embedding_sparse_batch.__doc__)
+
+
+# pylint: disable=protected-access
+def enqueue_tpu_embedding_sparse_tensor_batch(sample_indices,
+                                              embedding_indices,
+                                              aggregation_weights,
+                                              table_ids,
+                                              device_ordinal,
+                                              max_sequence_lengths=None,
+                                              num_features=None,
+                                              combiners=None,
+                                              mode_override=None,
+                                              name=None):
+  """A placeholder op for enqueueing embedding IDs to the TPU.
+
+  Args:
+    sample_indices: A list of rank 2 Tensors specifying the training example to
+      which the corresponding embedding_indices and aggregation_weights values
+      belong. It corresponds to sp_ids.indices in embedding_lookup_sparse(). If
+      the size of its first dimension is 0, we assume each embedding_indices
+      belongs to a different sample. Both int32 and int64 are allowed and will
+      be converted to int32 internally.
+    embedding_indices: A list of rank 1 Tensors, indices into the embedding
+      tables. It corresponds to sp_ids.values in embedding_lookup_sparse(). Both
+      int32 and int64 are allowed and will be converted to int32 internally.
+    aggregation_weights: A list of rank 1 Tensors containing per training
+      example aggregation weights. It corresponds to sp_weights.values in
+      embedding_lookup_sparse(). If the size of its first dimension is 0, we
+      assume all weights are 1. Both float32 and float64 are allowed and will be
+      converted to float32 internally.
+    table_ids: A list of integers specifying the identifier of the embedding
+      table (offset of TableDescriptor in the TPUEmbeddingConfiguration) to
+      lookup the corresponding input. The ith input is looked up using
+      table_ids[i]. The size of the table_ids list must be equal to that of
+      sample_indices, embedding_indices and aggregation_weights.
+    device_ordinal: The TPU device to use. Should be >= 0 and less than the
+      number of TPU cores in the task on which the node is placed.
+    max_sequence_lengths: A list of integers, the size of which is equal to
+      sample_indices. If equal to 0, the corresponding feature is considered to
+      be a non-sequence feature, If greater than 0, the corresponding feature is
+      a sequence feature with the given maximal length. If None, then we assume
+      a list of all zeroes.
+    num_features: A list of integers, the size of which is equal to
+      sample_indices. If non-empty, entries in this list must be at least 1. For
+      each batch element, we will take num_features rows of the input tensor for
+      embedding lookup. E.g., when sample_indices is empty, the embedding
+      indices must be of shape (batch_size*num_features).
+    combiners: A list of string scalars, one for each embedding table that
+      specify how to normalize the embedding activations after weighted
+      summation. Supported combiners are 'mean', 'sum', or 'sqrtn'. It is
+      invalid to have the sum of the weights be 0 for 'mean' or the sum of the
+      squared weights be 0 for 'sqrtn'. If combiners isn't passed, the default
+      is to use 'sum' for all tables (optional).
+    mode_override: A string input that overrides the mode specified in the
+      TPUEmbeddingConfiguration. Supported values are {'unspecified',
+      'inference', 'train', 'backward_pass_only'}. When set to 'unspecified',
+      the mode set in TPUEmbeddingConfiguration is used, otherwise mode_override
+      is used (optional).
+    name: A name for the operation (optional).
+
+  Returns:
+    An EnqueueTPUEmbeddingSparseTensorBatch operation.
+  """
+  if mode_override is None:
+    mode_override = "unspecified"
+  return gen_tpu_ops.enqueue_tpu_embedding_sparse_tensor_batch(
+      sample_indices=sample_indices,
+      embedding_indices=embedding_indices,
+      aggregation_weights=aggregation_weights,
+      table_ids=table_ids,
+      device_ordinal=device_ordinal,
+      max_sequence_lengths=max_sequence_lengths,
+      combiners=combiners,
+      mode_override=mode_override,
+      num_features=num_features,
+      name=name)
+
+
+enqueue_tpu_embedding_sparse_tensor_batch.__doc__ = (
+    gen_tpu_ops.enqueue_tpu_embedding_sparse_tensor_batch.__doc__)
+
+
+# pylint: disable=protected-access
+def enqueue_tpu_embedding_ragged_tensor_batch(sample_splits,
+                                              embedding_indices,
+                                              aggregation_weights,
+                                              table_ids,
+                                              device_ordinal,
+                                              max_sequence_lengths=None,
+                                              num_features=None,
+                                              combiners=None,
+                                              mode_override=None,
+                                              name=None):
+  """A placeholder op for enqueueing embedding IDs to the TPU.
+
+  Args:
+    sample_splits: A list of rank 1 Tensors specifying the break points for
+      splitting embedding_indices and aggregation_weights into rows. It
+      corresponds to ids.row_splits in embedding_lookup(), when ids is a
+      RaggedTensor. Both int32 and int64 are allowed and will be converted to
+      int32 internally.
+    embedding_indices: A list of rank 1 Tensors, indices into the embedding
+      tables. It corresponds to ids.values in embedding_lookup(), when ids is a
+      RaggedTensor. Both int32 and int64 are allowed and will be converted to
+      int32 internally.
+    aggregation_weights: A list of rank 1 Tensors containing per training
+      example aggregation weights. It corresponds to the values field of a
+      RaggedTensor with the same row_splits as ids in embedding_lookup(), when
+      ids is a RaggedTensor. Both float32 and float64 are allowed and will be
+      converted to float32 internally.
+    table_ids: A list of integers specifying the identifier of the embedding
+      table (offset of TableDescriptor in the TPUEmbeddingConfiguration) to
+      lookup the corresponding input. The ith input is looked up using
+      table_ids[i]. The size of the table_ids list must be equal to that of
+      sample_indices, embedding_indices and aggregation_weights.
+    device_ordinal: The TPU device to use. Should be >= 0 and less than the
+      number of TPU cores in the task on which the node is placed.
+    max_sequence_lengths: A list of integers, the size of which is equal to
+      sample_indices. If equal to 0, the corresponding feature is considered to
+      be a non-sequence feature, If greater than 0, the corresponding feature is
+      a sequence feature with the given maximal length. If None, then we assume
+      a list of all zeroes.
+    num_features: A list of integers, the size of which must be equal to
+      sample_indices. If non-empty, entries in this list must be at least 1. For
+      each batch element, we will take num_features rows of the input tensor for
+      embedding lookup. E.g., when sample_indices is empty, the embedding
+      indices must be of shape (batch_size*num_features).
+    combiners: A list of string scalars, one for each embedding table that
+      specify how to normalize the embedding activations after weighted
+      summation. Supported combiners are 'mean', 'sum', or 'sqrtn'. It is
+      invalid to have the sum of the weights be 0 for 'mean' or the sum of the
+      squared weights be 0 for 'sqrtn'. If combiners isn't passed, the default
+      is to use 'sum' for all tables (optional).
+    mode_override: A string input that overrides the mode specified in the
+      TPUEmbeddingConfiguration. Supported values are {'unspecified',
+      'inference', 'training', 'backward_pass_only'}. When set to 'unspecified',
+      the mode set in TPUEmbeddingConfiguration is used, otherwise mode_override
+      is used (optional).
+    name: A name for the operation (optional).
+
+  Returns:
+    An EnqueueTPUEmbeddingRaggedTensorBatch operation.
+  """
+  if mode_override is None:
+    mode_override = "unspecified"
+  return gen_tpu_ops.enqueue_tpu_embedding_ragged_tensor_batch(
+      sample_splits=sample_splits,
+      embedding_indices=embedding_indices,
+      aggregation_weights=aggregation_weights,
+      table_ids=table_ids,
+      device_ordinal=device_ordinal,
+      max_sequence_lengths=max_sequence_lengths,
+      combiners=combiners,
+      mode_override=mode_override,
+      num_features=num_features,
+      name=name)
+
+
+enqueue_tpu_embedding_ragged_tensor_batch.__doc__ = (
+    gen_tpu_ops.enqueue_tpu_embedding_ragged_tensor_batch.__doc__)
+
+
+def enqueue_tpu_embedding_arbitrary_tensor_batch(sample_indices_or_row_splits,
+                                                 embedding_indices,
+                                                 aggregation_weights,
+                                                 device_ordinal,
+                                                 combiners=None,
+                                                 mode_override=None,
+                                                 name=None):
+  """A placeholder op for enqueueing embedding IDs to the TPU.
+
+  Args:
+    sample_indices_or_row_splits: A list of rank 1 or 2 Tensors. When rank 2,
+      the tensors specify the training example to which the corresponding
+      embedding_indices and aggregation_weights values belong. If the size of
+      its first dimension is 0, we assume each embedding_indices belongs to a
+      different sample. Both int32 and int64 are allowed and will be converted
+      to int32 internally. When rank 1, the tensors specify the row splits for
+      splitting embedding_indices and aggregation_weights into rows. It
+      corresponds to ids.row_splits in embedding_lookup(), when ids is a
+      RaggedTensor. When enqueuing N-D ragged tensor, only the last dimension is
+      allowed to be ragged. the row splits is 1-D dense tensor. When empty, we
+      assume a dense tensor is passed to the op. Both int32 and int64 are
+      allowed and will be converted to int32 internally.
+    embedding_indices: A list of rank 1 Tensors, indices into the embedding
+      tables. Both int32 and int64 are allowed and will be converted to int32
+      internally.
+    aggregation_weights: A list of rank 1 Tensors containing per training
+      example aggregation weights. Both float32 and float64 are allowed and will
+      be converted to float32 internally.
+    device_ordinal: The TPU device to use. Should be >= 0 and less than the
+      number of TPU cores in the task on which the node is placed.
+    combiners: A list of string scalars, one for each embedding table that
+      specify how to normalize the embedding activations after weighted
+      summation. Supported combiners are 'mean', 'sum', or 'sqrtn'. It is
+      invalid to have the sum of the weights be 0 for 'mean' or the sum of the
+      squared weights be 0 for 'sqrtn'. If combiners isn't passed, the default
+      is to use 'sum' for all tables (optional).
+    mode_override: A string input that overrides the mode specified in the
+      TPUEmbeddingConfiguration. Supported values are {'unspecified',
+      'inference', 'training', 'backward_pass_only'}. When set to 'unspecified',
+      the mode set in TPUEmbeddingConfiguration is used, otherwise mode_override
+      is used (optional).
+    name: A name for the operation (optional).
+
+  Returns:
+    An EnqueueTPUEmbeddingArbitraryTensorBatch operation.
+  """
+  if mode_override is None:
+    mode_override = "unspecified"
+  return gen_tpu_ops.enqueue_tpu_embedding_arbitrary_tensor_batch(
+      sample_indices_or_row_splits=sample_indices_or_row_splits,
+      embedding_indices=embedding_indices,
+      aggregation_weights=aggregation_weights,
+      device_ordinal=device_ordinal,
+      combiners=combiners,
+      mode_override=mode_override,
+      name=name)
+
+
+enqueue_tpu_embedding_arbitrary_tensor_batch.__doc__ = (
+    gen_tpu_ops.enqueue_tpu_embedding_arbitrary_tensor_batch.__doc__)

tensorflow/python/tpu/ops/tpu_ordinal_selector_op.py

@@ -0,0 +1,17 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Operations to select TPU core to run."""
+

tensorflow/python/tpu/preempted_hook.py

@@ -0,0 +1,89 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Implementation of the SessionRunHook for preemptible Cloud TPUs."""
+
+import logging as _logging
+import os
+import threading
+import time
+
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.training import session_run_hook
+
+
+class CloudTPUPreemptedHook(session_run_hook.SessionRunHook):
+  """The SessionRunHook for preemptible Cloud TPUs.
+
+  This is an implementation of SessionRunHook for the pre-emptible Google Cloud
+  TPU service. It attempts to close the session if the TPU is preempted, and
+  exits the coordinator process if the session cannot be closed.
+  """
+
+  def __init__(self, cluster):
+    self._cluster = cluster
+
+  def after_create_session(self, session, coord):
+    if tpu_cluster_resolver.is_running_in_gce():
+      self._tpu_poller = _TPUPollingThread(self._cluster, session)
+      self._tpu_poller.start()
+
+  def end(self, session):
+    self._tpu_poller.stop()
+
+
+class _TPUPollingThread(threading.Thread):
+  """A thread that polls the state of a TPU node.
+
+  When the node transitions into a TERMINAL state (PREEMPTED, TERMINATED)
+  that's considered as not recoverable by the underlying infrastructure,
+  it attempts to close the session, and exits the entire process if the
+  session.close() stucks.
+  """
+
+  def __init__(self, cluster, session):
+    super(_TPUPollingThread, self).__init__()
+
+    self.daemon = True
+    self._running = True
+    self._session_closed = False
+    self._cluster = cluster
+    self._session = session
+    self._interval = 30
+
+    # Some of the Google API libraries are quite chatty, so disable them.
+    for name in ['googleapiclient.discovery', 'oauth2client.client']:
+      _logging.getLogger(name).setLevel(_logging.WARNING)
+
+  def stop(self):
+    self._running = False
+    self._session_closed = True
+    self.join()
+
+  def run(self):
+    if not tpu_cluster_resolver.is_running_in_gce():
+      logging.warning(
+          'TPUPollingThread is running in a non-GCE environment, exiting...')
+      self._running = False
+      return
+
+    while self._running:
+      recoverable = self._cluster._cloud_tpu_client.recoverable()  # pylint: disable=protected-access
+      if not recoverable:
+        logging.warning(
+            'TPUPollingThread found TPU %s in state %s',
+            self._cluster._tpu, self._cluster._cloud_tpu_client.state())  # pylint: disable=protected-access
+        os._exit(1)  # pylint: disable=protected-access
+      time.sleep(self._interval)

tensorflow/python/tpu/profiler/BUILD

@@ -0,0 +1,62 @@
+load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow:__subpackages__",
+    ],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "profiler",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/core/profiler:profiler_analysis_proto_py",
+        "//tensorflow/python/util:all_util",
+        "@tsl//tsl/profiler/protobuf:trace_events_proto_py",
+    ],
+)
+
+py_library(
+    name = "profiler_analysis_pb2_grpc",
+    srcs = ["profiler_analysis_pb2_grpc.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = ["//tensorflow/core/profiler:profiler_analysis_proto_py"],
+)
+
+py_library(
+    name = "capture_tpu_profile_lib",
+    srcs = [
+        "capture_tpu_profile.py",
+        "version.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:versions",
+        "//tensorflow/python/platform:gfile",
+        "//tensorflow/python/platform:tf_logging",
+        "//tensorflow/python/profiler:profiler_client",
+        "//tensorflow/python/profiler:profiler_v2",
+        "@absl_py//absl:app",
+        "@absl_py//absl/flags",
+        "@pypi//packaging",
+    ],
+)
+
+py_binary(
+    name = "capture_tpu_profile_bin",
+    srcs = ["capture_tpu_profile.py"],
+    main = "capture_tpu_profile.py",
+    strict_deps = True,
+    deps = [
+        ":capture_tpu_profile_lib",
+        "@absl_py//absl/flags",
+        "@pypi//packaging",
+    ],
+)

tensorflow/python/tpu/profiler/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Classes for TPU trace events."""
+
+# pylint: disable=wildcard-import,unused-import
+from tensorflow.core.profiler.profiler_analysis_pb2 import *
+from tensorflow.python.util.all_util import remove_undocumented
+from tsl.profiler.protobuf.trace_events_pb2 import *
+# pylint: enable=wildcard-import,unused-import
+
+
+_allowed_symbols = ['Trace', 'Resource', 'Device', 'TraceEvent']
+
+remove_undocumented(__name__, _allowed_symbols)

tensorflow/python/tpu/profiler/capture_tpu_profile.py

@@ -0,0 +1,206 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Cloud TPU profiler client."""
+
+import os
+import sys
+
+from absl import app
+from absl import flags
+from packaging.version import Version
+
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver as resolver
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import versions
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.profiler import profiler_client
+from tensorflow.python.profiler import profiler_v2 as profiler
+from tensorflow.python.tpu.profiler import version as profiler_version
+
+FLAGS = flags.FLAGS
+
+# Cloud TPU Cluster Resolvers
+flags.DEFINE_string(
+    'gcp_project', None,
+    'Project name for the Cloud TPU-enabled project. If not specified, we '
+    'will attempt to automatically detect the GCE project from metadata.')
+flags.DEFINE_string(
+    'tpu_zone',
+    None,
+    help='GCE zone where the Cloud TPU is located in. If not specified, we '
+    'will attempt to automatically detect the GCE project from metadata.')
+flags.DEFINE_string(
+    'tpu', None, 'Name of the Cloud TPU for Cluster Resolvers. You must '
+    'specify either this flag or --service_addr.')
+
+# Tool specific parameters
+flags.DEFINE_string(
+    'service_addr', None, 'Address of TPU profiler service e.g. '
+    'localhost:8466, you must specify either this flag or --tpu.')
+flags.DEFINE_string(
+    'workers_list', None, 'The list of worker TPUs that we are about to profile'
+    ' e.g. 10.0.1.2:8466, 10.0.1.3:8466. You can specify this flag with --tpu '
+    'or --service_addr to profile a subset of tpu nodes. You can also use only'
+    '--tpu and leave this flag unspecified to profile all the tpus.')
+flags.DEFINE_string(
+    'logdir', None, 'Path of TensorBoard log directory e.g. /tmp/tb_log, '
+    'gs://tb_bucket')
+flags.DEFINE_integer('duration_ms', 0,
+                     'Duration of tracing or monitoring in ms.')
+flags.DEFINE_integer(
+    'num_tracing_attempts', 3, 'Automatically retry N times when no trace '
+    'event is collected.')
+flags.DEFINE_boolean('include_dataset_ops', True, 'Deprecated.')
+flags.DEFINE_integer(
+    'host_tracer_level', 2, 'Adjust host tracer level to control the verbosity '
+    ' of the TraceMe event being collected.')
+
+# Monitoring parameters
+flags.DEFINE_integer(
+    'monitoring_level', 0, 'Choose a monitoring level between '
+    '1 and 2 to monitor your TPU job continuously. Level 2 is more verbose than'
+    ' level 1 and shows more metrics.')
+flags.DEFINE_integer(
+    'num_queries', 100,
+    'This script will run monitoring for num_queries before it stops.')
+flags.DEFINE_boolean('display_timestamp', True, 'Deprecated.')
+
+
+def get_workers_list(cluster_resolver):
+  """Returns a comma separated list of TPU worker host:port pairs.
+
+  Gets cluster_spec from cluster_resolver. Use the worker's task indices to
+  obtain and return a list of host:port pairs.
+
+  Args:
+    cluster_resolver: TensorFlow TPUClusterResolver instance.
+
+  Returns:
+    A string of comma separated list of host:port pairs. For example:
+    '10.2.0.1:8466,10.2.0.2:8466,10.2.0.3:8466,10.2.0.4:8466'
+
+  Raises:
+    UnavailableError: cluster_resolver doesn't contain a valid cluster_spec.
+  """
+  worker_job_name = 'worker'
+  cluster_spec = cluster_resolver.cluster_spec()
+  if not cluster_spec:
+    raise errors.UnavailableError(
+        'None', 'None',
+        'Cluster spec not found, your client must run in GCE environment.')
+  task_indices = cluster_spec.task_indices(worker_job_name)
+  workers_list = [
+      cluster_spec.task_address(worker_job_name, i).replace(':8470', ':8466')
+      for i in task_indices
+  ]
+  return ','.join(workers_list)
+
+
+def monitoring_helper(service_addr, duration_ms, monitoring_level, num_queries):
+  """Helper function to print monitoring results.
+
+  Helper function to print monitoring results for num_queries times.
+
+  Args:
+    service_addr: Address of the TPU profiler service.
+    duration_ms: Duration of one monitoring sample in milliseconds.
+    monitoring_level: An integer between 1 and 2. Level 2 is more verbose than
+      level 1 and shows more metrics.
+    num_queries: Number of monitoring samples to collect.
+  """
+  if monitoring_level <= 0 or monitoring_level > 2:
+    sys.exit('Please choose a monitoring level between 1 and 2.')
+
+  for query in range(0, num_queries):
+    res = profiler_client.monitor(service_addr, duration_ms, monitoring_level)
+    print('Cloud TPU Monitoring Results (Sample ', query, '):\n\n', res)
+
+
+def run_main():
+  app.run(main)
+
+
+def main(unused_argv=None):
+  logging.set_verbosity(logging.INFO)
+  tf_version = versions.__version__
+  print('TensorFlow version %s detected' % tf_version)
+  print('Welcome to the Cloud TPU Profiler v%s' % profiler_version.__version__)
+
+  if Version(tf_version) < Version('2.2.0'):
+    sys.exit('You must install tensorflow >= 2.2.0 to use this plugin.')
+
+  if not FLAGS.service_addr and not FLAGS.tpu:
+    sys.exit('You must specify either --service_addr or --tpu.')
+
+  tpu_cluster_resolver = None
+  if FLAGS.service_addr:
+    if FLAGS.tpu:
+      logging.warn('Both --service_addr and --tpu are set. Ignoring '
+                   '--tpu and using --service_addr.')
+    service_addr = FLAGS.service_addr
+  else:
+    try:
+      tpu_cluster_resolver = (
+          resolver.TPUClusterResolver([FLAGS.tpu],
+                                      zone=FLAGS.tpu_zone,
+                                      project=FLAGS.gcp_project))
+      service_addr = tpu_cluster_resolver.get_master()
+    except (ValueError, TypeError):
+      sys.exit('Failed to find TPU %s in zone %s project %s. You may use '
+               '--tpu_zone and --gcp_project to specify the zone and project of'
+               ' your TPU.' % (FLAGS.tpu, FLAGS.tpu_zone, FLAGS.gcp_project))
+  service_addr = service_addr.replace('grpc://', '').replace(':8470', ':8466')
+
+  workers_list = ''
+  if FLAGS.workers_list is not None:
+    workers_list = FLAGS.workers_list
+  elif tpu_cluster_resolver is not None:
+    workers_list = get_workers_list(tpu_cluster_resolver)
+
+  # If profiling duration was not set by user or set to a non-positive value,
+  # we set it to a default value of 1000ms.
+  duration_ms = FLAGS.duration_ms if FLAGS.duration_ms > 0 else 1000
+
+  if FLAGS.monitoring_level > 0:
+    print('Since monitoring level is provided, profile', service_addr, ' for ',
+          FLAGS.duration_ms, ' ms and show metrics for ', FLAGS.num_queries,
+          ' time(s).')
+    monitoring_helper(service_addr, duration_ms, FLAGS.monitoring_level,
+                      FLAGS.num_queries)
+  else:
+    if not FLAGS.logdir:
+      sys.exit('You must specify either --logdir or --monitoring_level.')
+
+    if not gfile.Exists(FLAGS.logdir):
+      gfile.MakeDirs(FLAGS.logdir)
+
+    try:
+      if Version(tf_version) < Version('2.3.0'):
+        profiler_client.trace(service_addr, os.path.expanduser(FLAGS.logdir),
+                              duration_ms, workers_list,
+                              FLAGS.num_tracing_attempts)
+      else:
+        options = profiler.ProfilerOptions(
+            host_tracer_level=FLAGS.host_tracer_level)
+        profiler_client.trace(service_addr, os.path.expanduser(FLAGS.logdir),
+                              duration_ms, workers_list,
+                              FLAGS.num_tracing_attempts, options)
+    except errors.UnavailableError:
+      sys.exit(0)
+
+
+if __name__ == '__main__':
+  run_main()

tensorflow/python/tpu/profiler/pip_package/BUILD

@@ -0,0 +1,15 @@
+# Description:
+#  Tools for building the Cloud TPU profiler pip package.
+
+package(default_visibility = ["//visibility:private"])
+
+licenses(["notice"])  # Apache 2.0
+
+sh_binary(
+    name = "build_pip_package",
+    srcs = ["build_pip_package.sh"],
+    data = [
+        "setup.py",
+        "//tensorflow/python/tpu/profiler:capture_tpu_profile_bin",
+    ],
+)

tensorflow/python/tpu/profiler/pip_package/README

@@ -0,0 +1,2 @@
+This package contains a Python wrapper around Tensorflow Python APIs that are
+used to profile Cloud TPU.

tensorflow/python/tpu/profiler/pip_package/build_pip_package.sh

@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+
+set -e
+
+if [ "$(uname)" = "Darwin" ]; then
+  sedi="sed -i ''"
+else
+  sedi="sed -i"
+fi
+
+PACKAGE_NAME="cloud_tpu_profiler"
+PIP_PACKAGE="tensorflow/python/tpu/profiler/pip_package"
+RUNFILES="bazel-bin/tensorflow/python/tpu/profiler/pip_package/build_pip_package.runfiles/org_tensorflow/tensorflow/python/tpu/profiler"
+
+function main() {
+  if [ $# -lt 1 ] ; then
+    echo "No destination dir provided"
+    exit 1
+  fi
+
+  DEST=$1
+  TMPDIR=$(mktemp -d -t tmp.XXXXXXXXXX)
+
+  echo $(date) : "=== Using tmpdir: ${TMPDIR}"
+
+  cp ${PIP_PACKAGE}/README ${TMPDIR}
+  cp ${PIP_PACKAGE}/setup.py ${TMPDIR}
+  mkdir ${TMPDIR}/${PACKAGE_NAME}
+  cp -a ${RUNFILES}/. ${TMPDIR}/${PACKAGE_NAME}/
+
+  # Fix the import statements to reflect the copied over path.
+  find  ${TMPDIR}/${PACKAGE_NAME} -name \*.py |
+    xargs $sedi -e '
+      s/^from tensorflow.python.tpu.profiler/from '${PACKAGE_NAME}'/
+  '
+  echo $(ls $TMPDIR)
+
+  pushd ${TMPDIR}
+  echo $(date) : "=== Building wheel"
+  echo $(pwd)
+  python setup.py bdist_wheel >/dev/null
+  python3 setup.py bdist_wheel >/dev/null
+  mkdir -p ${DEST}
+  cp dist/* ${DEST}
+  popd
+  rm -rf ${TMPDIR}
+  echo $(date) : "=== Output wheel file is in: ${DEST}"
+}
+
+main "$@"

tensorflow/python/tpu/profiler/pip_package/setup.py

@@ -0,0 +1,58 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Cloud TPU profiler package."""
+
+from setuptools import setup
+
+from cloud_tpu_profiler.version import __version__
+
+CONSOLE_SCRIPTS = [
+    'capture_tpu_profile=cloud_tpu_profiler.capture_tpu_profile:run_main',
+]
+
+setup(
+    name='cloud_tpu_profiler',
+    version=__version__.replace('-', ''),
+    description='Trace and profile Cloud TPU performance',
+    long_description='Tools for capture TPU profile',
+    url='https://www.tensorflow.org/tfrc/',
+    author='Google Inc.',
+    author_email='packages@tensorflow.org',
+    packages=['cloud_tpu_profiler'],
+    entry_points={
+        'console_scripts': CONSOLE_SCRIPTS,
+    },
+    classifiers=[
+        'Development Status :: 5 - Production/Stable',
+        'Intended Audience :: Developers',
+        'Intended Audience :: Education',
+        'Intended Audience :: Science/Research',
+        'License :: OSI Approved :: Apache Software License',
+        'Programming Language :: Python :: 2',
+        'Programming Language :: Python :: 2.7',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.4',
+        'Programming Language :: Python :: 3.5',
+        'Programming Language :: Python :: 3.6',
+        'Topic :: Scientific/Engineering',
+        'Topic :: Scientific/Engineering :: Mathematics',
+        'Topic :: Scientific/Engineering :: Artificial Intelligence',
+        'Topic :: Software Development',
+        'Topic :: Software Development :: Libraries',
+        'Topic :: Software Development :: Libraries :: Python Modules',
+    ],
+    license='Apache 2.0',
+    keywords='tensorflow performance tpu',
+)

tensorflow/python/tpu/profiler/profiler_analysis_pb2_grpc.py

@@ -0,0 +1,121 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+#
+# Do not use pylint on generated code.
+# pylint: disable=missing-docstring,g-short-docstring-punctuation,g-no-space-after-docstring-summary,invalid-name,line-too-long,unused-argument,g-doc-args
+import grpc
+
+from tensorflow.core.profiler import profiler_analysis_pb2 as third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+
+
+class ProfileAnalysisStub(object):
+  """//////////////////////////////////////////////////////////////////////////////
+
+  ProfileAnalysis service provide entry point for profiling TPU and for
+  serving profiled data to Tensorboard through GRPC
+  //////////////////////////////////////////////////////////////////////////////
+  """
+
+  def __init__(self, channel):
+    """Constructor.
+
+    Args:
+      channel: A grpc.Channel.
+    """
+    self.NewSession = channel.unary_unary(
+        '/tensorflow.ProfileAnalysis/NewSession',
+        request_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+        .NewProfileSessionRequest.SerializeToString,
+        response_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+        .NewProfileSessionResponse.FromString,
+    )
+    self.EnumSessions = channel.unary_unary(
+        '/tensorflow.ProfileAnalysis/EnumSessions',
+        request_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+        .EnumProfileSessionsAndToolsRequest.SerializeToString,
+        response_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+        .EnumProfileSessionsAndToolsResponse.FromString,
+    )
+    self.GetSessionToolData = channel.unary_unary(
+        '/tensorflow.ProfileAnalysis/GetSessionToolData',
+        request_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+        .ProfileSessionDataRequest.SerializeToString,
+        response_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+        .ProfileSessionDataResponse.FromString,
+    )
+
+
+class ProfileAnalysisServicer(object):
+  """//////////////////////////////////////////////////////////////////////////////
+
+  ProfileAnalysis service provide entry point for profiling TPU and for
+  serving profiled data to Tensorboard through GRPC
+  //////////////////////////////////////////////////////////////////////////////
+  """
+
+  def NewSession(self, request, context):
+    """Starts a profiling session, blocks until it completes.
+
+    TPUProfileAnalysis service delegate this to TPUProfiler service.
+    Populate the profiled data in repository, then return status to caller.
+    """
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+  def EnumSessions(self, request, context):
+    """Enumerate existing sessions and return available profile tools."""
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+  def GetSessionToolData(self, request, context):
+    """Retrieve specific tool's data for specific session."""
+    context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+    context.set_details('Method not implemented!')
+    raise NotImplementedError('Method not implemented!')
+
+
+def add_ProfileAnalysisServicer_to_server(servicer, server):
+  rpc_method_handlers = {
+      'NewSession':
+          grpc.unary_unary_rpc_method_handler(
+              servicer.NewSession,
+              request_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+              .NewProfileSessionRequest.FromString,
+              response_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+              .NewProfileSessionResponse.SerializeToString,
+          ),
+      'EnumSessions':
+          grpc.unary_unary_rpc_method_handler(
+              servicer.EnumSessions,
+              request_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+              .EnumProfileSessionsAndToolsRequest.FromString,
+              response_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+              .EnumProfileSessionsAndToolsResponse.SerializeToString,
+          ),
+      'GetSessionToolData':
+          grpc.unary_unary_rpc_method_handler(
+              servicer.GetSessionToolData,
+              request_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+              .ProfileSessionDataRequest.FromString,
+              response_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
+              .ProfileSessionDataResponse.SerializeToString,
+          ),
+  }
+  generic_handler = grpc.method_handlers_generic_handler(
+      'tensorflow.ProfileAnalysis', rpc_method_handlers)
+  server.add_generic_rpc_handlers((generic_handler,))

tensorflow/python/tpu/profiler/version.py

@@ -0,0 +1,20 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Cloud TPU profiler version information."""
+
+# Cloud TPU profiler uses semantic versioning, see http://semver.org/.
+# A version string consists of
+# major_version.minor_version.patch_version-build_metadata.
+__version__ = "1.14.1-a0"

tensorflow/python/tpu/pywrap_sparse_core_layout.cc

@@ -0,0 +1,45 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <pybind11/pybind11.h>
+
+#include "pybind11/cast.h"  // from @pybind11
+#include "pybind11/detail/common.h"  // from @pybind11
+#include "pybind11_abseil/status_casters.h"  // from @pybind11_abseil
+#include "pybind11_protobuf/native_proto_caster.h"  // from @pybind11_protobuf
+#include "tensorflow/core/tpu/kernels/sparse_core_layout.h"
+#include "tensorflow/core/tpu/kernels/sparse_core_layout.pb.h"
+
+namespace tensorflow::tpu {
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_pywrap_sparse_core_layout, m) {
+  py::class_<SparseCoreLayoutStacker>(m, "SparseCoreLayoutStacker")
+      .def(py::init<int, bool, int>(), py::arg("num_partitions"),
+           py::arg("disable_table_stacking"),
+           py::arg("sparse_cores_per_partition"))
+      .def("SetActivationMemoryBytesLimit",
+           &SparseCoreLayoutStacker::SetActivationMemoryBytesLimit)
+      .def("SetVariableShardBytesLimit",
+           &SparseCoreLayoutStacker::SetVariableShardBytesLimit)
+      .def("SetStackingEnabled", &SparseCoreLayoutStacker::SetStackingEnabled)
+      .def("AddTable", &SparseCoreLayoutStacker::AddTable,
+           py::arg("table_name"), py::arg("table_height"),
+           py::arg("table_width"), py::arg("group"), py::arg("output_samples"),
+           py::arg("num_features"))
+      .def("GetLayouts", &SparseCoreLayoutStacker::GetLayouts);
+}
+
+}  // namespace tensorflow::tpu

tensorflow/python/tpu/pywrap_tpu_embedding.cc

@@ -0,0 +1,35 @@
+/* Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+#include "tensorflow/core/tpu/kernels/sparse_core_ops_utils.h"
+#include "tensorflow/python/lib/core/pybind11_lib.h"
+
+PYBIND11_MODULE(_pywrap_tpu_embedding, m) {
+  m.def("stack_tables",
+        [](const std::vector<int64_t>& table_heights,
+           const std::vector<int64_t>& table_widths,
+           const std::vector<int64_t>& table_num_samples,
+           const std::vector<int64_t>& table_groups,
+           const std::vector<std::string>& table_names, int64_t num_tpu_chips) {
+          return tensorflow::GetTableStacks(table_heights, table_widths,
+                                            table_num_samples, table_groups,
+                                            table_names, num_tpu_chips);
+        });
+}

tensorflow/python/tpu/tensor_tracer.proto

@@ -0,0 +1,131 @@
+// Copyright 2026 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow;
+
+import "tensorflow/core/framework/graph.proto";
+
+// Tensor Tracer Report proto gives information about the trace including:
+// - TensorTracerConfig: version, device, num replicas, trace mode.
+// - Graphdef, e.g., list of operations, tensors
+// - TracedTensorDef:
+//    * Name of the tensor
+//    * Tracepoint name if provided.
+//    * Index of the tensor in the compact cache if traced.
+//    * Explanation for why the tensor is traced or not.
+message TensorTracerReport {
+  TensorTracerConfig config = 1;
+
+  // Tensorflow graph.
+  tensorflow.GraphDef graphdef = 2;
+
+  // A map from tensor name to its TracedTensorDef.
+  map<string, TracedTensorDef> tensordef = 3;
+
+  // The fingerprint of the TensorTracerReport (fingerprint calculation excludes
+  // this field and graphdef).
+  string fingerprint = 4;
+
+  // The function_name passed to the function_callback
+  // that produced this TensorTracerReport
+  string concrete_function_name = 5;
+
+  // The index of the last stack frame where the stack traces for all output
+  // operations in the graph have the same value.
+  int32 last_common_frame_no = 6;
+
+  // List of names of output tensors of the function being traced.
+  repeated string outputs = 7;
+
+  // Information about the number of tensors traced and skipped.
+  TracingStats tracing_stats = 8;
+
+  message TensorTracerConfig {
+    // Tensor tracer version, e.g. hostcall, outside compilation.
+    string version = 1;
+    // Traced device, CPU, TPU...
+    string device = 2;
+
+    // Trace mode, norm, summary, full-trace.
+    string trace_mode = 3;
+
+    // Number of cores, e.g. TPU cores, in the system.
+    int32 num_cores = 4;
+
+    // Number of hosts, e.g. compute nodes in the system.
+    int32 num_hosts = 5;
+
+    // Keep submode as string for backward compatibility.
+    string submode = 6;
+
+    // Keep num cores per host for backward compatibility.
+    int32 num_cores_per_host = 7;
+
+    // Id of the included cores, if a subset of cores are traced.
+    repeated int32 included_cores = 8;
+
+    // The names of the signatures corresponding to the cache indices.
+    repeated string signatures = 9;
+  }
+
+  message TracingStats {
+    // The total number of tensors in the function.
+    int32 total_tensors = 1;
+    // The number of traced tensors in the function.
+    int32 traced_tensors = 2;
+    // Counts of traced tensors by op type.
+    map<string, int32> traced_tensor_types = 3;
+    // The number of tensors added by Tensor Tracer.
+    int32 added_tensors = 4;
+  }
+
+  message TracedTensorDef {
+    // Name of the tensor as appears in tf graph.
+    string name = 1;
+    // Cache index of the tensor. This may be different than topological index.
+    int32 cache_index = 2;
+    // If trace points are provided, corresponding tracepoint name of the
+    // tensor. Trace points are placed on the edges (tensors) in the tensorflow
+    // graph, and they force tensor tracer to trace the corresponding tensor.
+    // Tracepoints can be added using the programatic interface
+    // tensor_tracer.tensor_tracepoint(tensor, trace_point_name) function.
+    // This will add a trace point with the given trace_point_name for the given
+    // tensor. If a trace_point is provided for the tensor,
+    // trace_point name will be used for the rest of the analysis instead of
+    // tensor names. One can use trace_point_name's to compare two models with
+    // arbitrary tensor names by providing the same trace point name for the
+    // tensors that are comparable.
+    string trace_point_name = 3;
+    // Whether the tensor is traced or not.
+    bool is_traced = 4;
+    // Detailed explanation why the tensor is traced or not.
+    string explanation = 5;
+    // Detailed stack of operation
+    Stack op_stack_info = 6;
+
+    message Stack {
+      // Function names from stack
+      repeated string stack_fn_names = 1;
+      // Line in stack
+      repeated string stack_lines = 2;
+      // Filenames from stack
+      repeated string stack_filenames = 3;
+      // Line number in file from stack
+      repeated int32 stack_linenos = 4;
+    }
+  }
+}

tensorflow/python/tpu/tensor_tracer_flags.py

@@ -0,0 +1,504 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========================================================================
+"""Utilities to handle tensor tracer parameters."""
+
+
+import os
+import os.path
+import re
+from absl import flags
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import tf_logging as logging
+
+TRACE_MODE_PART_TENSOR = 'part-tensor'
+TRACE_MODE_FULL_TENSOR = 'full-tensor'
+TRACE_MODE_FULL_TENSOR_SUMMARY = 'full_tensor_summary'
+
+TRACE_MODE_NAN_INF = 'nan-inf'
+TRACE_MODE_NORM = 'norm'
+TRACE_MODE_MAX_ABS = 'max-abs'
+TRACE_MODE_SUMMARY = 'summary'
+TRACE_MODE_HISTORY = 'history'
+# summary mode to collects a finite set of signatures for each traced tensor,
+# (such as norm, max, min, mean) and dumps it using tb summaries.
+
+# Full tensor mode dumps the whole tensor values for the traced tensors without
+# any processing on them; using tb summaries.
+
+_SUBMODE_BRIEF = 'brief'
+_SUBMODE_DETAILED = 'detailed'
+
+_FLAG_SINGLE_QUOTE_PAT = re.compile(r"\s*--([^=]+)='([^']*)'")
+_FLAG_DOUBLE_QUOTE_PAT = re.compile(r'\s*--([^=]+)="([^"]*)"')
+_FLAG_NO_QUOTE_PAT = re.compile(r'\s*--([^=]+)=(\S*)')
+_FLAG_NO_EQUAL_PAT = re.compile(r'\s*--([^=]+)\s*')
+
+FLAGS_ENV_VAR = 'TENSOR_TRACER_FLAGS'
+FLAG_NAME_ENABLE = 'enable'
+FLAG_NAME_TRACE_MODE = 'trace_mode'
+FLAG_NAME_TRACE_SCALAR_OPS = 'trace_scalar'
+FLAG_NAME_SUBMODE = 'submode'
+FLAG_NAME_EXCLUDED_OPNAMES = 'excluded_opnames'
+FLAG_NAME_EXCLUDED_OPTYPES = 'excluded_optypes'
+FLAG_NAME_INCLUDED_OPNAMES = 'included_opnames'
+FLAG_NAME_INCLUDED_OPTYPES = 'included_optypes'
+FLAG_NAME_TRACE_LEVEL = 'trace_level'
+FLAG_NAME_TRACE_DIR = 'trace_dir'
+FLAG_NAME_REPORT_FILE = 'report_file'
+FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR = 'use_test_undeclared_outputs_dir'
+FLAG_NAME_OP_RANGE = 'op_range'
+# Folder to dump the pre (before tensor tracer updates) and post graphs (after
+# tensor tracer updates).
+FLAG_NAME_DUMP_BEFORE_AFTER_GRAPHS = 'dump_graphs'
+FLAG_NAME_SUMMARY_SIGNATURES = 'signatures'
+FLAG_NAME_SUMMARY_PER_CORE = 'collect_summary_per_core'
+FLAG_NAME_TEMP_CACHE_VAR = 'use_temp_cache'
+FLAG_NAME_INSPECT_TRACE = 'inspect_trace'
+FLAG_NAME_FINGERPRINT_DIR = 'use_fingerprint_subdirectory'
+FLAG_FLUSH_SUMMARY = 'flush_summaries'
+
+
+VALID_FLAG_NAMES = [
+    FLAG_NAME_ENABLE, FLAG_NAME_TRACE_MODE,
+    FLAG_NAME_TRACE_SCALAR_OPS,
+    FLAG_NAME_SUBMODE, FLAG_NAME_EXCLUDED_OPNAMES,
+    FLAG_NAME_EXCLUDED_OPTYPES, FLAG_NAME_INCLUDED_OPNAMES,
+    FLAG_NAME_INCLUDED_OPTYPES, FLAG_NAME_TRACE_DIR,
+    FLAG_NAME_REPORT_FILE,
+    FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR,
+    FLAG_NAME_OP_RANGE,
+    FLAG_NAME_DUMP_BEFORE_AFTER_GRAPHS, FLAG_NAME_TRACE_LEVEL,
+    FLAG_NAME_SUMMARY_SIGNATURES, FLAG_NAME_SUMMARY_PER_CORE,
+    FLAG_NAME_TEMP_CACHE_VAR, FLAG_NAME_FINGERPRINT_DIR,
+    FLAG_NAME_INSPECT_TRACE, FLAG_FLUSH_SUMMARY,
+]
+
+_OP_RANGE_PAT = re.compile(r'(\d+):(\d+)')
+_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR = 'TEST_UNDECLARED_OUTPUTS_DIR'
+
+_TT_DEFAULT_TRACE_LEVEL = 3
+_TT_PREFIX = 'tensor_tracer'
+
+_TT_NORM = 'norm'
+_TT_MAX = 'max'
+_TT_MAX_ABS = 'max-abs'
+_TT_MIN = 'min'
+_TT_SPARSITY = 'sparsity'
+_TT_MEAN = 'mean'
+_TT_VAR = 'var'
+_TT_SIZE = 'size'
+
+TT_SUMMARY_NORM = '%s_%s' % (_TT_PREFIX, _TT_NORM)
+TT_SUMMARY_MAX = '%s_%s' % (_TT_PREFIX, _TT_MAX)
+TT_SUMMARY_MAX_ABS = '%s_%s' % (_TT_PREFIX, _TT_MAX_ABS)
+TT_SUMMARY_MIN = '%s_%s' % (_TT_PREFIX, _TT_MIN)
+TT_SUMMARY_SPARSITY = '%s_%s' % (_TT_PREFIX, _TT_SPARSITY)
+TT_SUMMARY_MEAN = '%s_%s' % (_TT_PREFIX, _TT_MEAN)
+TT_SUMMARY_VAR = '%s_%s' % (_TT_PREFIX, _TT_VAR)
+TT_SUMMARY_SIZE = '%s_%s' % (_TT_PREFIX, _TT_SIZE)
+
+TT_SUMMARY_SIGNATURES = (TT_SUMMARY_NORM, TT_SUMMARY_MAX, TT_SUMMARY_MIN,
+                         TT_SUMMARY_SPARSITY, TT_SUMMARY_MEAN, TT_SUMMARY_VAR,
+                         TT_SUMMARY_SIZE, TT_SUMMARY_MAX_ABS)
+
+FLAGS = flags.FLAGS
+
+DELTA_THRESHOLD = flags.DEFINE_float(
+    'delta_threshold',
+    default=0.5,
+    help=('Log if history based diff crosses this threshold.'))
+TT_CHECK_FILTER = flags.DEFINE_bool(
+    'tt_check_filter',
+    default=False,
+    help='Terminate early to check op name filtering.')
+TT_SINGLE_CORE_SUMMARIES = flags.DEFINE_bool(
+    'tt_single_core_summaries',
+    default=False,
+    help='Report single core metric and avoid aggregation.')
+
+
+class TTParameters(object):
+  """A class that handles the parameters of Tensor Tracer."""
+
+  def __init__(self, env=None):
+    if env:
+      self._env = env
+    else:
+      self._env = os.environ
+    self._validate_flag_names()
+    self.trace_mode = self._get_trace_mode()
+    self.submode = self._get_submode()
+    self.trace_dir = self._get_trace_dir()
+    self.report_file_path = self._get_report_filepath()
+    self.op_range = self._get_op_range()
+    self.excluded_opname_re_list = self._flag_value_to_re_list(
+        FLAG_NAME_EXCLUDED_OPNAMES)
+    self.excluded_optype_re_list = self._flag_value_to_re_list(
+        FLAG_NAME_EXCLUDED_OPTYPES)
+
+    self.included_opname_re_list = self._flag_value_to_re_list(
+        FLAG_NAME_INCLUDED_OPNAMES)
+    self.included_optype_re_list = self._flag_value_to_re_list(
+        FLAG_NAME_INCLUDED_OPTYPES)
+
+    self.trace_scalar_ops = self.is_flag_on(FLAG_NAME_TRACE_SCALAR_OPS)
+    self.use_compact_trace = self.trace_mode in (TRACE_MODE_NAN_INF,
+                                                 TRACE_MODE_NORM,
+                                                 TRACE_MODE_HISTORY,
+                                                 TRACE_MODE_MAX_ABS,
+                                                 TRACE_MODE_SUMMARY)
+    self.use_temp_cache_var = self.is_flag_on(FLAG_NAME_TEMP_CACHE_VAR)
+    self.inspect_trace = self.is_flag_on(FLAG_NAME_INSPECT_TRACE)
+    self.use_fingerprint_subdir = self.is_flag_on(FLAG_NAME_FINGERPRINT_DIR)
+
+    _, self.graph_dump_path = self.get_flag_value(
+        FLAG_NAME_DUMP_BEFORE_AFTER_GRAPHS)
+    self.trace_level = self._get_flag_int_value(FLAG_NAME_TRACE_LEVEL,
+                                                _TT_DEFAULT_TRACE_LEVEL)
+    self.summary_signatures = self._get_summary_signatures()
+    self.collect_summary_per_core = self.is_flag_on(FLAG_NAME_SUMMARY_PER_CORE)
+    # TODO(b/199284834): Will be resolved with referenced bug.
+    if self.collect_summary_per_core:
+      logging.warning('Aggregate signatures are approximate for mean, variance'
+                      ' and sparsity.')
+    self.flush_summaries_with_outside_compile = self.is_flag_on(
+        FLAG_FLUSH_SUMMARY)
+    # Do not produce errors or warnings if Tensor Tracer is not enabled.
+    if self.is_enabled():
+      self._check_flag_errors()
+
+  def _check_flag_errors(self):
+    if self.trace_mode in (TRACE_MODE_SUMMARY, TRACE_MODE_FULL_TENSOR_SUMMARY):
+      if not self.trace_dir:
+        raise ValueError('trace_dir must be explicitly provided in '
+                         'TENSOR_TRACER_FLAGS when summary mode is used.')
+
+  def _get_report_filepath(self):
+    """Sets the path of the output report file."""
+
+    found, report_file_path = self.get_flag_value(FLAG_NAME_REPORT_FILE)
+    if found and report_file_path and self.use_test_undeclared_outputs_dir():
+      if os.path.isabs(report_file_path):
+        raise ValueError('If use_test_undeclared_outputs_dir is set,'
+                         'report_file_path cannot be an absolute path (%s)'
+                         %report_file_path)
+      outputs_dir = self._env.get(_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR)
+      report_file_path = os.path.join(outputs_dir, report_file_path)
+    return report_file_path
+
+  def _get_op_range(self):
+    """Sets the index range of the Ops that we will consider tracing."""
+    found, op_range = self.get_flag_value(FLAG_NAME_OP_RANGE)
+    if not found or not op_range:
+      op_range = (-1, -1)  # this means including all ops.
+      return op_range
+    match = _OP_RANGE_PAT.match(op_range)
+    if not match:
+      op_range = (-1, -1)  # this means including all ops.
+      return op_range
+    op_range = (int(match.group(1)), int(match.group(2)))
+    return op_range
+
+  def _get_trace_dir(self):
+    found, trace_dir = self.get_flag_value(FLAG_NAME_TRACE_DIR)
+    if found and trace_dir and self.use_test_undeclared_outputs_dir():
+      raise ValueError(
+          'Cannot not use --%s and --%s at the same time' %
+          (FLAG_NAME_TRACE_DIR, FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR))
+    if self.use_test_undeclared_outputs_dir():
+      trace_dir = self._env.get(_TEST_UNDECLARED_OUTPUTS_DIR_ENV_VAR)
+    return trace_dir
+
+  def _get_trace_mode(self):
+    """Checks if the given trace mode is valid."""
+
+    found, trace_mode = self.get_flag_value(FLAG_NAME_TRACE_MODE)
+    if not found or not trace_mode:
+      trace_mode = TRACE_MODE_NORM
+    valid_trace_modes = [
+        TRACE_MODE_NAN_INF, TRACE_MODE_PART_TENSOR, TRACE_MODE_FULL_TENSOR,
+        TRACE_MODE_NORM, TRACE_MODE_MAX_ABS,
+        TRACE_MODE_SUMMARY, TRACE_MODE_FULL_TENSOR_SUMMARY,
+        TRACE_MODE_HISTORY
+    ]
+    if trace_mode not in valid_trace_modes:
+      raise ValueError('Invalid trace mode "%s" given to the Tensor_Tracer.'
+                       'Valid trace modes are: %s'%(trace_mode,
+                                                    valid_trace_modes))
+    return trace_mode
+
+  def is_brief_mode(self):
+    return self.submode == _SUBMODE_BRIEF
+
+  def _get_submode(self):
+    """Checks if the given submode is valid."""
+
+    found, submode = self.get_flag_value(FLAG_NAME_SUBMODE)
+    if not found or not submode:
+      submode = _SUBMODE_DETAILED
+    if not submode:
+      return
+    valid_submodes = [_SUBMODE_DETAILED, _SUBMODE_BRIEF]
+    if submode not in valid_submodes:
+      raise ValueError('Invalid submode "%s" given to the Tensor_Tracer.'
+                       'Valid submodes are: %s'%(submode,
+                                                 valid_submodes))
+    return submode
+
+  @staticmethod
+  def match_next_flag(tt_flags, pos):
+    """Returns the match for the next TensorTracer flag.
+
+    Args:
+       tt_flags: a string that contains the flags.
+       pos: where in flags to start the search.
+
+    Returns:
+       A pair where the first element is the regular-expression
+       match found and the second element indicates if the match
+       has a value.
+    """
+
+    match = _FLAG_DOUBLE_QUOTE_PAT.match(tt_flags, pos)
+    if match:
+      return match, True
+    match = _FLAG_SINGLE_QUOTE_PAT.match(tt_flags, pos)
+    if match:
+      return match, True
+    match = _FLAG_NO_QUOTE_PAT.match(tt_flags, pos)
+    if match:
+      return match, True
+    match = _FLAG_NO_EQUAL_PAT.match(tt_flags, pos)
+    if match:
+      # The flag is found but is not given a value.
+      return match, False
+    # The flag is not found.
+    return None, False
+
+  def _validate_flag_names(self):
+    """Validates if the TensorTrace flags passed are valid."""
+    tensor_tracer_flags = self._env.get(FLAGS_ENV_VAR)
+    if not tensor_tracer_flags:
+      return
+    pos = 0
+    while True:
+      match, _ = TTParameters.match_next_flag(tensor_tracer_flags, pos)
+      if not match:
+        break
+      flag_name = match.group(1)
+      if flag_name not in VALID_FLAG_NAMES:
+        raise ValueError(
+            'The flag name "%s" passed via the environment variable "%s" '
+            'is invalid. Valid flag names are:'
+            '\n%s' % (flag_name, FLAGS_ENV_VAR, VALID_FLAG_NAMES))
+      pos = match.end()
+
+  def _supported_signatures(self):
+    """Returns a tuple of supported signatures."""
+    return TT_SUMMARY_SIGNATURES
+
+  def _get_summary_signatures(self):
+    """Verifies and returns the summary signatures.
+
+    Returns:
+      A dictionary of the signature identifiers {signature: index} that will be
+      computed when trace_mode is summary.
+    """
+    signatures = self._flag_value_as_list(FLAG_NAME_SUMMARY_SIGNATURES)
+    supported_signatures = self._supported_signatures()
+
+    tt_signatures = []
+    for signature in signatures:
+      signature_with_prefix = '%s_%s' % (_TT_PREFIX, signature)
+      if signature in supported_signatures:
+        tt_signatures.append(signature)
+      elif signature_with_prefix in supported_signatures:
+        tt_signatures.append(signature_with_prefix)
+      else:
+        logging.warning('Unknown signature:%s. Supported signatures: %s' %
+                        (signature, supported_signatures))
+    if not tt_signatures:
+      # Default case collects norm and max only.
+      return {TT_SUMMARY_MAX_ABS: 0, TT_SUMMARY_NORM: 1}
+    else:
+      return {signature: idx for idx, signature in enumerate(tt_signatures)}
+
+  def get_signature_to_agg_fn_map(self):
+    """Returns a map that contains the aggregate function for each signature."""
+    # TODO(b/199284834): Aggregations are not accurate for mean and sparsity if
+    # cores have a different number of elements. Variance uses the maximal core
+    # variance.
+    return {TRACE_MODE_NORM: linalg_ops.norm,
+            TRACE_MODE_HISTORY: math_ops.reduce_max,
+            TRACE_MODE_MAX_ABS: math_ops.reduce_max,
+            TRACE_MODE_NAN_INF: math_ops.reduce_max,
+            TT_SUMMARY_NORM: linalg_ops.norm,
+            TT_SUMMARY_MAX: math_ops.reduce_max,
+            TT_SUMMARY_MAX_ABS:
+                lambda t, axis=0: math_ops.reduce_max(math_ops.abs(t),  # pylint: disable=g-long-lambda
+                                                      axis=axis),
+            TT_SUMMARY_MIN: math_ops.reduce_min,
+            # Exact if each part has the same number of values.
+            TT_SUMMARY_SPARSITY: math_ops.reduce_mean,
+            TT_SUMMARY_MEAN: math_ops.reduce_mean,
+            TT_SUMMARY_VAR: math_ops.reduce_max,  # Simply reduce max variance.
+            TT_SUMMARY_SIZE: math_ops.reduce_sum}
+
+  def _flag_value_as_list(self, wanted_flag_name):
+    """Returns the string list of a TensorTracer flag.
+
+    Args:
+      wanted_flag_name: the name of the flag we are looking for.
+
+    Returns:
+      The list value of the flag.
+    """
+    string_value_list = []
+    found, flag_value = self.get_flag_value(wanted_flag_name)
+
+    if found:
+      assert flag_value is not None
+      string_value_list = flag_value.split(',')
+    return string_value_list
+
+  def _flag_value_as_int_list(self, wanted_flag_name):
+    """Returns the integer list of a TensorTracer flag.
+
+    Args:
+      wanted_flag_name: the name of the flag we are looking for.
+
+    Returns:
+      the value of the flag.
+    Raises:
+      RuntimeError: If supposedly deadcode is reached.
+    """
+    int_list = []
+    found, flag_value = self.get_flag_value(wanted_flag_name)
+
+    if found and flag_value:
+      try:
+        integer_values = flag_value.split(',')
+        int_list = [int(int_val) for int_val in integer_values]
+      except ValueError:
+        logging.warning('Cannot convert %s to int for flag %s', int_list,
+                        wanted_flag_name)
+    return int_list
+
+  def _get_flag_int_value(self, wanted_flag_name, default_value):
+    """Returns the int value of a TensorTracer flag.
+
+    Args:
+      wanted_flag_name: the name of the flag we are looking for.
+      default_value: the default value for the flag, if not provided.
+    Returns:
+      the value of the flag.
+    Raises:
+      RuntimeError: If supposedly deadcode is reached.
+    """
+    flag_int_value = default_value
+    found, flag_value = self.get_flag_value(wanted_flag_name)
+
+    if found:
+      try:
+        flag_int_value = int(flag_value)
+      except ValueError:
+        logging.warning('Cannot convert %s to int for flag %s' % (
+            flag_int_value, wanted_flag_name))
+    return flag_int_value
+
+  def get_flag_value(self, wanted_flag_name):
+    """Returns the value of a TensorTracer flags.
+
+    Args:
+      wanted_flag_name: the name of the flag we are looking for.
+
+    Returns:
+      A pair where the first element indicates if the flag is
+      found and the second element is the value of the flag.
+
+    Raises:
+      RuntimeError: If supposedly deadcode is reached.
+    """
+
+    tensor_tracer_flags = self._env.get(FLAGS_ENV_VAR)
+    if not tensor_tracer_flags:
+      return False, None
+    pos = 0
+    while True:
+      match, has_value = TTParameters.match_next_flag(
+          tensor_tracer_flags, pos)
+      if not match:
+        return False, None
+      flag_name = match.group(1)
+      if has_value:
+        flag_value = match.group(2)
+      else:
+        flag_value = None
+      if flag_name == wanted_flag_name:
+        return True, flag_value
+      pos = match.end()
+    raise RuntimeError('Invalid tensor tracer flag. Could not recognize %s.' %
+                       flag_name)
+
+  def _flag_value_to_re_list(self, flag_name):
+    """Converts list of strings to compiled RE."""
+
+    re_list = []
+    found, flag_value = self.get_flag_value(flag_name)
+    if not found or not flag_value:
+      return re_list
+    list_of_values = flag_value.split(',')
+    for v in list_of_values:
+      r = re.compile(v)
+      re_list.append(r)
+    return re_list
+
+  def is_flag_on(self, flag_name):
+    """Returns True if the given flag is on."""
+
+    found, flag_value = self.get_flag_value(flag_name)
+    if not found:
+      return False
+    if flag_value is None:
+      return True
+    # Depends on the flag value.
+    flag_value = flag_value.lower()
+    enabled = flag_value in ['1', 't', 'true', 'y', 'yes']
+    return enabled
+
+  def is_enabled(self):
+    """Returns True if TensorTracer is enabled."""
+
+    if self.is_flag_on(FLAG_NAME_ENABLE):
+      logging.debug('Tensor Tracer is enabled with flags %s.',
+                    self._env.get(FLAGS_ENV_VAR))
+      return True
+    else:
+      return False
+
+  def use_test_undeclared_outputs_dir(self):
+    """Decides the output directory of the report and trace files.
+
+    Args:
+       None.
+
+    Returns:
+       True if the output files should be written to the
+       test-undeclared-outputs-directory defined via an
+       env variable.
+    """
+
+    return self.is_flag_on(FLAG_NAME_USE_TEST_UNDECLARED_OUTPUTS_DIR)

tensorflow/python/tpu/tensor_tracer_report.py

@@ -0,0 +1,436 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ========================================================================
+"""Tensor Tracer report generation utilities."""
+
+import collections
+import hashlib
+import os
+
+from tensorflow.python.platform import gfile
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.tpu import tensor_tracer_pb2
+
+_TRACER_LOG_PREFIX = ' [>>>TT>>>]'
+_MARKER_SECTION_BEGIN = '!!!!!!! section-begin:'
+_MARKER_SECTION_END = '!!!!!!! section-end:'
+
+_SECTION_NAME_CONFIG = 'configuration'
+_SECTION_NAME_REASON = 'reason'
+_SECTION_NAME_OP_LIST = 'op-list'
+_SECTION_NAME_TENSOR_LIST = 'tensor-list'
+_SECTION_NAME_CACHE_INDEX_MAP = 'cache-index-map'
+_SECTION_NAME_GRAPH = 'graph'
+_SECTION_NAME_TENSOR_TRACER_CHECKPOINT = 'tensor_tracer_checkpoint'
+
+_FIELD_NAME_VERSION = 'version:'
+_FIELD_NAME_DEVICE = 'device:'
+_FIELD_NAME_TRACE_MODE = 'trace-mode:'
+_FIELD_NAME_SUBMODE = 'submode:'
+_FIELD_NAME_NUM_REPLICAS = 'num-replicas:'
+_FIELD_NAME_NUM_REPLICAS_PER_HOST = 'num-replicas-per-host:'
+_FIELD_NAME_NUM_HOSTS = 'num-hosts:'
+_FIELD_NAME_NUM_OPS = 'number-of-ops:'
+_FIELD_NAME_NUM_TENSORS = 'number-of-tensors:'
+_FIELD_NAME_NUM_CACHE_INDICES = 'number-of-indices:'
+_FIELD_NAME_TOPOLOGICAL_SORT_SUCCEED = 'topological-sort-succeed:'
+
+_CURRENT_VERSION = 'use-outside-compilation'
+
+_TT_REPORT_PROTO = 'tensor_tracer_report.report_pb'
+
+
+def topological_sort(g):
+  """Performs topological sort on the given graph.
+
+  Args:
+     g: the graph.
+
+  Returns:
+     A pair where the first element indicates if the topological
+     sort succeeded (True if there is no cycle found; False if a
+     cycle is found) and the second element is either the sorted
+     list of nodes or the cycle of nodes found.
+  """
+  def _is_loop_edge(op):
+    """Returns true if the op is the end of a while-loop creating a cycle."""
+    return op.type in ['NextIteration']
+
+  def _in_op_degree(op):
+    """Returns the number of incoming edges to the given op.
+
+    The edge calculation skips the edges that come from 'NextIteration' ops.
+    NextIteration creates a cycle in the graph. We break cycles by treating
+    this op as 'sink' and ignoring all outgoing edges from it.
+    Args:
+      op: Tf.Operation
+    Returns:
+      the number of incoming edges.
+    """
+    count = 0
+    for op in op.control_inputs + [in_tensor.op for in_tensor in op.inputs]:
+      if not _is_loop_edge(op):
+        count += 1
+    return count
+
+  sorted_ops = []
+  op_in_degree = {op: _in_op_degree(op) for op in g.get_operations()}
+
+  frontier = [op for (op, degree) in op_in_degree.items() if degree == 0]
+  frontier.sort(key=lambda op: op.name)
+  while frontier:
+    op = frontier.pop()
+    # Remove the op from graph, and remove its outgoing edges.
+    sorted_ops.append(op)
+    if _is_loop_edge(op):
+      continue
+    # pylint: disable=protected-access
+    consumers = list(op._control_outputs)
+    # pylint: enable=protected-access
+    for out_tensor in op.outputs:
+      consumers += [consumer_op for consumer_op in out_tensor.consumers()]
+    consumers.sort(key=lambda op: op.name)
+    for consumer in consumers:
+      # For each deleted edge shift the bucket of the vertex.
+      op_in_degree[consumer] -= 1
+      if op_in_degree[consumer] == 0:
+        frontier.append(consumer)
+      if op_in_degree[consumer] < 0:
+        raise ValueError('consumer:%s degree mismatch'%consumer.name)
+
+  left_ops = set(op for (op, degree) in op_in_degree.items() if degree > 0)
+  if left_ops:
+    return (True, left_ops)
+  else:
+    assert len(g.get_operations()) == len(sorted_ops)
+    return (False, sorted_ops)
+
+
+class TensorTracerConfig(object):
+  """Tensor Tracer config object."""
+
+  def __init__(self):
+    self.version = _CURRENT_VERSION
+    self.device_type = None
+    self.num_replicas = None
+    self.num_replicas_per_host = None
+    self.num_hosts = None
+
+
+class TensorTraceOrder(object):
+  """Class that is responsible from storing the trace-id of the tensors."""
+
+  def __init__(self, graph_order, traced_tensors):
+    self.graph_order = graph_order
+    self.traced_tensors = traced_tensors
+    self._create_tensor_maps()
+
+  def _create_tensor_maps(self):
+    """Creates tensor to cache id maps."""
+    self.tensorname_to_cache_idx = {}
+    self.cache_idx_to_tensor_idx = []
+    for out_tensor in self.traced_tensors:
+      tensor_name = out_tensor.name
+      if tensor_name in self.tensorname_to_cache_idx:
+        raise ValueError('Tensor name {} should not be already in '
+                         'tensorname_to_cache_idx'.format(tensor_name))
+      if tensor_name not in self.graph_order.tensor_to_idx:
+        raise ValueError(
+            'Tensor name {} is not in the tensor_to_idx, tensor_to_idx={} '
+            .format(tensor_name, self.graph_order.tensor_to_idx))
+      tensor_idx = self.graph_order.tensor_to_idx[tensor_name]
+      cache_idx = len(self.tensorname_to_cache_idx)
+      self.tensorname_to_cache_idx[tensor_name] = cache_idx
+      self.cache_idx_to_tensor_idx.append(tensor_idx)
+      if len(self.tensorname_to_cache_idx) != len(
+          self.cache_idx_to_tensor_idx):
+        raise RuntimeError(
+            'len(self.tensorname_to_cache_idx) must equal'
+            'len(self.cache_idx_to_tensor_idx), got '
+            'len(self.tensorname_to_cache_idx)={}, '
+            'len(self.cache_idx_to_tensor_idx)={}'
+            .format(
+                len(self.tensorname_to_cache_idx),
+                len(self.cache_idx_to_tensor_idx)))
+
+
+def sort_tensors_and_ops(graph):
+  """Returns a wrapper that has consistent tensor and op orders."""
+  graph_wrapper = collections.namedtuple('GraphWrapper',
+                                         ['graph', 'operations', 'op_to_idx',
+                                          'tensors', 'tensor_to_idx',
+                                          'contains_cycle',
+                                          'topological_order_or_cycle'])
+  contains_cycle, topological_order_or_cycle = topological_sort(graph)
+  if not contains_cycle:
+    operations = topological_order_or_cycle
+  else:
+    operations = graph.get_operations()
+  op_to_idx = {op.name: index for index, op
+               in enumerate(operations)}
+  tensors = []
+  for op in operations:
+    tensors.extend(op.outputs)
+  tensor_to_idx = {tensor.name: index for index, tensor in
+                   enumerate(tensors)}
+  return graph_wrapper(graph=graph, operations=operations, op_to_idx=op_to_idx,
+                       tensors=tensors, tensor_to_idx=tensor_to_idx,
+                       contains_cycle=contains_cycle,
+                       topological_order_or_cycle=topological_order_or_cycle)
+
+
+class OpenReportFile(object):
+  """Context manager for writing report file."""
+
+  def __init__(self, tt_parameters):
+    if not tt_parameters.report_file_path:
+      self._report_file = None
+      return
+    try:
+      self._report_file = gfile.Open(tt_parameters.report_file_path, 'w')
+    except IOError as e:
+      raise e
+
+  def __enter__(self):
+    return self._report_file
+
+  def __exit__(self, unused_type, unused_value, unused_traceback):
+    if self._report_file:
+      self._report_file.close()
+
+
+def proto_fingerprint(message_proto):
+  serialized_message = message_proto.SerializeToString()
+  hasher = hashlib.sha256(serialized_message)
+  return hasher.hexdigest()
+
+
+class TTReportHandle(object):
+  """Utility class responsible from creating a tensor tracer report."""
+
+  def __init__(self):
+    self.instrument_records = {}
+    self._report_file = None
+
+  def instrument(self, name, explanation):
+    self.instrument_records[name] = explanation
+
+  def instrument_op(self, op, explanation):
+    self.instrument(op.name, explanation)
+
+  def instrument_tensor(self, tensor, explanation):
+    self.instrument(tensor.name, explanation)
+
+  def create_report_proto(self, tt_config, tt_parameters, tensor_trace_order,
+                          tensor_trace_points, collected_signature_types):
+    """Creates and returns a proto that stores tensor tracer configuration.
+
+    Args:
+      tt_config: TensorTracerConfig object holding information about the run
+        environment (device, # cores, # hosts), and tensor tracer version
+        information.
+      tt_parameters: TTParameters objects storing the user provided parameters
+        for tensor tracer.
+      tensor_trace_order: TensorTraceOrder object storing a topological order of
+        the graph.
+      tensor_trace_points: Progromatically added trace_points/checkpoints.
+      collected_signature_types: The signature types collected, e,g, norm,
+        max, min, mean...
+    Returns:
+      TensorTracerReport proto.
+    """
+    report = tensor_tracer_pb2.TensorTracerReport()
+    report.config.version = tt_config.version
+    report.config.device = tt_config.device_type
+    report.config.num_cores = tt_config.num_replicas
+    report.config.num_hosts = tt_config.num_hosts
+    report.config.num_cores_per_host = tt_config.num_replicas_per_host
+    report.config.submode = tt_parameters.submode
+    report.config.trace_mode = tt_parameters.trace_mode
+
+    for signature_name, _ in sorted(collected_signature_types.items(),
+                                    key=lambda x: x[1]):
+      report.config.signatures.append(signature_name)
+
+    for tensor in tensor_trace_order.graph_order.tensors:
+      tensor_def = tensor_tracer_pb2.TensorTracerReport.TracedTensorDef()
+      tensor_def.name = tensor.name
+      if tensor.name in tensor_trace_order.tensorname_to_cache_idx:
+        tensor_def.is_traced = True
+        tensor_def.cache_index = (
+            tensor_trace_order.tensorname_to_cache_idx[tensor.name])
+      else:
+        # To prevent small changes affecting the fingerprint calculation, avoid
+        # writing the untraced tensors to metadata. Fingerprints will be
+        # different only when the list of the traced tensors are different.
+        if tt_parameters.use_fingerprint_subdir:
+          continue
+        tensor_def.is_traced = False
+
+      if tensor.name in tensor_trace_points:
+        tensor_def.trace_point_name = tensor_trace_points[tensor.name]
+      if tensor.name in self.instrument_records:
+        tensor_def.explanation = self.instrument_records[tensor.name]
+      elif tensor.op.name in self.instrument_records:
+        tensor_def.explanation = self.instrument_records[tensor.op.name]
+      report.tensordef[tensor.name].CopyFrom(tensor_def)
+    report.fingerprint = proto_fingerprint(report)
+    logging.info('TensorTracerProto fingerprint is %s.',
+                 report.fingerprint)
+    tf_graph = tensor_trace_order.graph_order.graph
+    report.graphdef.CopyFrom(tf_graph.as_graph_def())
+    return report
+
+  def report_proto_path(self, trace_dir, summary_tag_name):
+    """Returns the path where report proto should be written.
+
+    Args:
+      trace_dir: String denoting the trace directory.
+      summary_tag_name: Name of the unique tag that relates to
+                        the report.
+    Returns:
+      A string denoting the path to the report proto.
+    """
+    filename = _TT_REPORT_PROTO + '.' + summary_tag_name.replace('/', '_')
+    return os.path.join(trace_dir, filename)
+
+  def write_report_proto(self, report_path, report_proto, tt_parameters):
+    """Writes the given report proto under trace_dir."""
+    gfile.MakeDirs(tt_parameters.trace_dir)
+    with gfile.GFile(report_path, 'wb') as f:
+      f.write(report_proto.SerializeToString())
+
+  def create_report(self, tt_config, tt_parameters,
+                    tensor_trace_order, tensor_trace_points):
+    """Creates a report file and writes the trace information."""
+    with OpenReportFile(tt_parameters) as self._report_file:
+      self._write_config_section(tt_config, tt_parameters)
+      self._write_op_list_section(tensor_trace_order.graph_order)
+      self._write_tensor_list_section(tensor_trace_order.graph_order)
+      self._write_trace_points(tensor_trace_points)
+      self._write_cache_index_map_section(tensor_trace_order)
+      self._write_reason_section()
+      self._write_graph_section(tensor_trace_order.graph_order)
+
+  def _write_trace_points(self, tensor_trace_points):
+    """Writes the list of checkpoints."""
+    self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN,
+                                  _SECTION_NAME_TENSOR_TRACER_CHECKPOINT))
+    for (tensor, checkpoint_name) in tensor_trace_points:
+      self._write_report('%s %s\n'%(tensor.name, checkpoint_name))
+    self._write_report('%s %s\n'%(_MARKER_SECTION_END,
+                                  _SECTION_NAME_TENSOR_TRACER_CHECKPOINT))
+
+  def _write_report(self, content):
+    """Writes the given content to the report."""
+
+    line = '%s %s'%(_TRACER_LOG_PREFIX, content)
+    if self._report_file:
+      self._report_file.write(line)
+    else:
+      logging.info(line)
+
+  def _write_config_section(self, tt_config, tt_parameters):
+    """Writes the config section of the report."""
+
+    self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_CONFIG))
+    self._write_report('%s %s\n'%(_FIELD_NAME_VERSION, tt_config.version))
+    self._write_report('%s %s\n'%(_FIELD_NAME_DEVICE, tt_config.device_type))
+    self._write_report('%s %s\n'%(_FIELD_NAME_TRACE_MODE,
+                                  tt_parameters.trace_mode))
+    self._write_report('%s %s\n'%(_FIELD_NAME_SUBMODE,
+                                  tt_parameters.submode))
+    self._write_report('%s %s\n'%(_FIELD_NAME_NUM_REPLICAS,
+                                  tt_config.num_replicas))
+    self._write_report('%s %s\n'%(_FIELD_NAME_NUM_REPLICAS_PER_HOST,
+                                  tt_config.num_replicas_per_host))
+    self._write_report('%s %s\n'%(_FIELD_NAME_NUM_HOSTS, tt_config.num_hosts))
+    self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_CONFIG))
+
+  def _write_reason_section(self):
+    """Writes the reason section of the report."""
+
+    self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_REASON))
+    for key in sorted(self.instrument_records):
+      self._write_report('"%s" %s\n'%(key, self.instrument_records[key]))
+    self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_REASON))
+
+  def _write_op_list_section(self, graph_order):
+    """Writes the Op-list section of the report."""
+
+    self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_OP_LIST))
+    self._write_report('%s %d\n'%(_FIELD_NAME_NUM_OPS,
+                                  len(graph_order.operations)))
+    for i in range(0, len(graph_order.operations)):
+      op = graph_order.operations[i]
+      line = '%d "%s" %s'%(i, op.name, op.type)
+      for out_tensor in op.outputs:
+        if out_tensor.name not in graph_order.tensor_to_idx:
+          raise ValueError(
+              'out_tensor is not in tensor_to_idx. out_tensor={}, '
+              'tensor_to_idx={}'
+              .format(out_tensor.name, graph_order.tensor_to_idx))
+        line += ' %d'%graph_order.tensor_to_idx[out_tensor.name]
+      line += '\n'
+      self._write_report(line)
+    self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_OP_LIST))
+
+  def _write_tensor_list_section(self, graph_order):
+    """Writes the tensor-list section of the report."""
+
+    self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN,
+                                  _SECTION_NAME_TENSOR_LIST))
+    self._write_report('%s %d\n'%(_FIELD_NAME_NUM_TENSORS,
+                                  len(graph_order.tensors)))
+    for i in range(0, len(graph_order.tensors)):
+      tensor = graph_order.tensors[i]
+      line = '%d "%s"'%(i, tensor.name)
+      consumers = tensor.consumers()
+      consumers.sort(key=lambda op: op.name)
+      for consumer_op in consumers:
+        if consumer_op.name not in graph_order.op_to_idx:
+          raise ValueError(
+              'consumer_op is not in op_to_idx.  '
+              'got consumer_op={}, op_to_idx={}'
+              .format(consumer_op.name, graph_order.op_to_idx))
+        line += ' %d'%graph_order.op_to_idx[consumer_op.name]
+      line += '\n'
+      self._write_report(line)
+    self._write_report('%s %s\n'%(_MARKER_SECTION_END,
+                                  _SECTION_NAME_TENSOR_LIST))
+
+  def _write_cache_index_map_section(self, tensor_trace_order):
+    """Writes the mapping from cache index to tensor index to the report."""
+    self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN,
+                                  _SECTION_NAME_CACHE_INDEX_MAP))
+    self._write_report('%s %d\n'%(
+        _FIELD_NAME_NUM_CACHE_INDICES,
+        len(tensor_trace_order.cache_idx_to_tensor_idx)))
+    for cache_idx in range(0, len(tensor_trace_order.cache_idx_to_tensor_idx)):
+      tensor_idx = tensor_trace_order.cache_idx_to_tensor_idx[cache_idx]
+      line = '%d %d\n'%(cache_idx, tensor_idx)
+      self._write_report(line)
+    self._write_report('%s %s\n'%(_MARKER_SECTION_END,
+                                  _SECTION_NAME_CACHE_INDEX_MAP))
+
+  def _write_graph_section(self, graph_order):
+    """Writes the graph section of the report."""
+
+    self._write_report('%s %s\n'%(_MARKER_SECTION_BEGIN, _SECTION_NAME_GRAPH))
+    self._write_report('%s %s\n'%(_FIELD_NAME_TOPOLOGICAL_SORT_SUCCEED,
+                                  not graph_order.contains_cycle))
+    l = list(graph_order.topological_order_or_cycle)
+    for i in range(0, len(l)):
+      self._write_report('%d "%s"\n'%(i, l[i].name))
+    self._write_report('%s %s\n'%(_MARKER_SECTION_END, _SECTION_NAME_GRAPH))

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()

tensorflow/python/tpu/topology.py

@@ -0,0 +1,241 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ======================================
+"""Defines the `Topology` class, that describes a TPU fabric topology."""
+
+import numpy as np
+
+from tensorflow.core.protobuf.tpu import topology_pb2
+from tensorflow.python.util import numpy_compat
+from tensorflow.python.util.tf_export import tf_export
+
+
+def _tpu_device_name(job, task, device):
+  """Returns the device name for the TPU `device` on `task` of `job`."""
+  if job is None:
+    return "/task:%d/device:TPU:%d" % (task, device)
+  else:
+    return "/job:%s/task:%d/device:TPU:%d" % (job, task, device)
+
+
+def _tpu_host_device_name(job, task):
+  """Returns the device name for the CPU device on `task` of `job`."""
+  if job is None:
+    return "/task:%d/device:CPU:0" % task
+  else:
+    return "/job:%s/task:%d/device:CPU:0" % (job, task)
+
+
+@tf_export("tpu.experimental.Topology")
+class Topology(object):
+  """Describes a set of TPU devices.
+
+  Represents both the shape of the physical mesh, and the mapping between
+  TensorFlow TPU devices to physical mesh coordinates.
+  """
+
+  def __init__(self, serialized=None, mesh_shape=None, device_coordinates=None):
+    """Builds a Topology object.
+
+    If `serialized` is not `None`, the topology is parsed from `serialized` and
+    the other arguments are ignored. Otherwise, the topology is computed from
+    `mesh_shape` and `device_coordinates`.
+
+    Args:
+      serialized: A serialized `TopologyProto`, or `None`. If not `None`, the
+        serialized proto is parsed to discover the topology.
+      mesh_shape: A sequence of 4 positive integers, or `None`. If not `None`,
+        the shape of the TPU topology, in number of cores. Ignored if
+        `serialized` is not `None`.
+      device_coordinates: A rank 3 numpy array that describes the mapping from
+        TensorFlow TPU devices to TPU fabric coordinates, or `None`. If
+        specified, array is a rank 3 int32 array with shape
+        `[tasks, devices, axis]`.  `tasks` is the number of tasks in the TPU
+        cluster, `devices` is the number of TPU devices per task, and `axis` is
+        the number of axes in the TPU cluster topology. Each entry gives the
+        `axis`-th coordinate in the topology of a task/device pair. TPU
+        topologies are 4-dimensional, with dimensions `(x, y, z, core number)`.
+        This arg is ignored if `serialized is not `None`.
+
+    Raises:
+      ValueError: If `serialized` does not describe a well-formed topology.
+      ValueError: If `serialized` is `None` and `mesh_shape` is not a sequence
+        of 4 positive integers.
+      ValueError: If `serialized` is `None` and `device_coordinates` is not a
+        rank 3 numpy int32 array that describes a valid coordinate mapping.
+    """
+
+    self._serialized = serialized
+
+    if serialized:
+      self._parse_topology(serialized)
+    else:
+      self._mesh_shape = numpy_compat.np_asarray(mesh_shape, dtype=np.int32)
+      self._device_coordinates = numpy_compat.np_asarray(device_coordinates,
+                                                         dtype=np.int32)
+      if len(self._mesh_shape) != 4 or any(self._mesh_shape < 1):
+        raise ValueError("`mesh_shape` must be a sequence of 4 positive "
+                         f"entries; got `mesh_shape={self._mesh_shape}`")
+
+      if (len(self._device_coordinates.shape) != 3 or
+          self._device_coordinates.shape[2] != len(self._mesh_shape)):
+        raise ValueError(
+            "`device_coordinates` must be a rank 3 int32 array "
+            "with minor dimension equal to the `mesh_shape` rank"
+            "got device_coordinates.shape={} len(device_coordinates.shape)={} device_coordinates.shape[2]={} mesh_shape={}, len(mesh_shape)={}"
+            .format(self._device_coordinates.shape,
+                    len(self._device_coordinates.shape),
+                    self._device_coordinates.shape[2], self._mesh_shape,
+                    len(self._mesh_shape)))
+
+    self._topology_tasks, self._topology_devices = self._invert_topology()
+
+    # Coordinates of devices that are missing
+    self._missing_devices = np.argwhere(self._topology_tasks < 0)
+
+  def _parse_topology(self, serialized):
+    """Parses a serialized `TopologyProto` into `self`."""
+    proto = topology_pb2.TopologyProto()
+    proto.ParseFromString(serialized)
+
+    self._mesh_shape = np.array(proto.mesh_shape, dtype=np.int32)
+    if len(self._mesh_shape) != 4 or any(self._mesh_shape < 1):
+      raise ValueError("`mesh_shape` must be a vector of size 4 with positive "
+                       "entries; got {}".format(self._mesh_shape))
+
+    if proto.num_tasks < 0:
+      raise ValueError("`num_tasks` must be >= 0; got {}".format(
+          proto.num_tasks))
+    if proto.num_tpu_devices_per_task < 0:
+      raise ValueError("`num_tpu_devices_per_task` must be >= 0; got {}".format(
+          proto.num_tpu_devices_per_task))
+
+    expected_coordinates_size = (
+        proto.num_tasks * proto.num_tpu_devices_per_task * len(
+            proto.mesh_shape))
+    if len(proto.device_coordinates) != expected_coordinates_size:
+      raise ValueError("`device_coordinates` must have shape num_tasks ({}) * "
+                       "num_tpu_devices_per_task ({}) * len(mesh_shape) ({}); "
+                       "got shape {}".format(proto.num_tasks,
+                                             proto.num_tpu_devices_per_task,
+                                             proto.mesh_shape,
+                                             len(proto.device_coordinates)))
+
+    coords = np.array(proto.device_coordinates, dtype=np.int32)
+    if any(coords < 0):
+      raise ValueError(
+          "All values in `device_coordinates` must be >= 0, got {}"
+          .format(coords))
+    coords = coords.reshape((proto.num_tasks, proto.num_tpu_devices_per_task,
+                             len(proto.mesh_shape)))
+    self._device_coordinates = coords
+
+  def _invert_topology(self):
+    """Inverts a [task,device,axis] topology to [x,y,z] -> task/device maps."""
+    tasks = np.full(list(self.mesh_shape), -1, dtype=np.int32)
+    devices = np.full(list(self.mesh_shape), -1, dtype=np.int32)
+    for task in range(self.device_coordinates.shape[0]):
+      for device in range(self.device_coordinates.shape[1]):
+        x, y, z, core = self.device_coordinates[task, device, :]
+        tasks[x, y, z, core] = task
+        devices[x, y, z, core] = device
+    return tasks, devices
+
+  @property
+  def mesh_shape(self):
+    """A rank 1 int32 array describing the shape of the TPU topology."""
+    return self._mesh_shape
+
+  @property
+  def mesh_rank(self):
+    """Returns the number of dimensions in the mesh."""
+    return len(self._mesh_shape)
+
+  @property
+  def device_coordinates(self):
+    """Describes the mapping from TPU devices to topology coordinates.
+
+    Returns:
+      A rank 3 int32 array with shape `[tasks, devices, axis]`.
+      `tasks` is the number of tasks in the TPU cluster, `devices` is the number
+      of TPU devices per task, and `axis` is the number of axes in the TPU
+      cluster topology. Each entry gives the `axis`-th coordinate in the
+      topology of a task/device pair. TPU topologies are 4-dimensional, with
+      dimensions `(x, y, z, core number)`.
+    """
+    return self._device_coordinates
+
+  @property
+  def missing_devices(self):
+    """Array of indices of missing devices."""
+    return self._missing_devices
+
+  def task_ordinal_at_coordinates(self, device_coordinates):
+    """Returns the TensorFlow task number attached to `device_coordinates`.
+
+    Args:
+      device_coordinates: An integer sequence describing a device's physical
+        coordinates in the TPU fabric.
+
+    Returns:
+      Returns the TensorFlow task number that contains the TPU device with those
+      physical coordinates.
+    """
+    return self._topology_tasks[tuple(device_coordinates)]
+
+  def tpu_device_ordinal_at_coordinates(self, device_coordinates):
+    """Returns the TensorFlow device number at `device_coordinates`.
+
+    Args:
+      device_coordinates: An integer sequence describing a device's physical
+        coordinates in the TPU fabric.
+
+    Returns:
+      Returns the TensorFlow device number within the task corresponding to
+      attached to the device with those physical coordinates.
+    """
+    return self._topology_devices[tuple(device_coordinates)]
+
+  def cpu_device_name_at_coordinates(self, device_coordinates, job=None):
+    """Returns the CPU device attached to a logical core."""
+    return _tpu_host_device_name(
+        job, self._topology_tasks[tuple(device_coordinates)])
+
+  def tpu_device_name_at_coordinates(self, device_coordinates, job=None):
+    """Returns the name of the TPU device assigned to a logical core."""
+    return _tpu_device_name(job,
+                            self._topology_tasks[tuple(device_coordinates)],
+                            self._topology_devices[tuple(device_coordinates)])
+
+  @property
+  def num_tasks(self):
+    """Returns the number of TensorFlow tasks in the TPU slice."""
+    return self._device_coordinates.shape[0]
+
+  @property
+  def num_tpus_per_task(self):
+    """Returns the number of TPU devices per task in the TPU slice."""
+    return self._device_coordinates.shape[1]
+
+  def serialized(self):
+    """Returns the serialized form of the topology."""
+    if self._serialized is None:
+      proto = topology_pb2.TopologyProto()
+      proto.mesh_shape[:] = list(self._mesh_shape)
+      proto.num_tasks = self._device_coordinates.shape[0]
+      proto.num_tpu_devices_per_task = self._device_coordinates.shape[1]
+      proto.device_coordinates.extend(list(self._device_coordinates.flatten()))
+      self._serialized = proto.SerializeToString()
+
+    return self._serialized

tensorflow/python/tpu/topology_test.py

@@ -0,0 +1,41 @@
+# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Tests for topology.py."""
+
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import topology
+
+
+class TopologyTest(test.TestCase):
+
+  def testSerialization(self):
+    """Tests if the class is able to generate serialized strings."""
+    original_topology = topology.Topology(
+        mesh_shape=[1, 1, 1, 2],
+        device_coordinates=[[[0, 0, 0, 0], [0, 0, 0, 1]]],
+    )
+    serialized_str = original_topology.serialized()
+    new_topology = topology.Topology(serialized=serialized_str)
+
+    # Make sure the topology recovered from serialized str is same as the
+    # original topology.
+    self.assertAllEqual(
+        original_topology.mesh_shape, new_topology.mesh_shape)
+    self.assertAllEqual(
+        original_topology.device_coordinates, new_topology.device_coordinates)
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/tpu/tpu.bzl

@@ -0,0 +1,66 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Provides python test rules for Cloud TPU."""
+
+load("@rules_python//python:py_test.bzl", "py_test")
+load(
+    "//tensorflow/python/tpu:tpu_test_wrapper.bzl",
+    _get_kwargs_for_wrapping = "get_kwargs_for_wrapping",
+)
+
+def tpu_py_test(
+        name,
+        tags = None,
+        disable_v2 = False,
+        disable_v3 = False,
+        disable_v3_4chips = True,
+        disable_experimental = False,
+        disable_mlir_bridge = True,
+        disable_tfrt = None,
+        args = [],
+        test_rule = py_test,
+        **kwargs):
+    """Generates identical unit test variants for various Cloud TPU versions.
+
+    TODO(rsopher): actually generate v2 vs v3 tests.
+
+    Args:
+        name: Name of test. Will be prefixed by accelerator versions.
+        tags: BUILD tags to apply to tests.
+        disable_v2: If true, don't generate TPU v2 tests.
+        disable_v3: If true, don't generate TPU v3 tests.
+        disable_v3_4chips: If true, don't generate TPU v3 2x2 tests.
+        disable_experimental: Unused.
+        disable_mlir_bridge: Unused.
+        disable_tfrt: Unused.
+        args: Arguments to apply to tests.
+        **kwargs: Additional named arguments to apply to tests.
+    """
+
+    test_rule(
+        **_get_kwargs_for_wrapping(
+            name,
+            tags,
+            args,
+            **kwargs
+        )
+    )
+
+def tpu_py_strict_test(**kwargs):
+    tpu_py_test(**kwargs)
+
+def internal_create_sanitizer_settings():
+    """Stub definition for an external rule."""
+    pass

tensorflow/python/tpu/tpu_embedding_base.py

@@ -0,0 +1,145 @@
+# 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 Embeddings Mid level APIs."""
+
+import functools
+from typing import Any, Dict, Iterable, Optional, Union, Text
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import variables as tf_variables
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.trackable import autotrackable
+from tensorflow.python.util import nest
+
+
+class TPUEmbeddingBase(autotrackable.AutoTrackable):
+  """The TPUEmbedding Base class.
+
+  This class only contains the basic logic to check the feature config and table
+  config for the tpu embedding mid level APIs.
+  """
+
+  def __init__(
+      self,
+      feature_config: Union[tpu_embedding_v2_utils.FeatureConfig, Iterable],  # pylint:disable=g-bare-generic
+      optimizer: Optional[tpu_embedding_v2_utils._Optimizer] = None):  # pylint:disable=protected-access
+    """Creates the TPUEmbeddingBase object."""
+    self._feature_config = feature_config
+    self._output_shapes = []
+    for feature in nest.flatten(feature_config):
+      self._output_shapes.append(feature.output_shape)
+    # Set table order here to the order of the first occurrence of the table in
+    # a feature provided by the user. The order of this struct must be fixed
+    # to provide the user with deterministic behavior over multiple
+    # instantiations.
+    self._table_config = []
+    for feature in nest.flatten(feature_config):
+      if feature.table not in self._table_config:
+        self._table_config.append(feature.table)
+
+    # Ensure tables have unique names. Also error check the optimizer as we
+    # specifically don't do that in the TableConfig class to allow high level
+    # APIs that are built on this to use strings/other classes to represent
+    # optimizers (before they are passed to this class).
+    table_names = []
+    for i, table in enumerate(self._table_config):
+      if table.optimizer is None:
+        # TODO(bfontain) Should we allow some sort of optimizer merging here?
+        table.optimizer = optimizer
+      if (table.optimizer is not None and
+          not isinstance(table.optimizer, tpu_embedding_v2_utils._Optimizer)):  # pylint: disable=protected-access
+        raise ValueError("{} is an unsupported optimizer class. Please pass an "
+                         "instance of one of the optimizer classes under "
+                         "tf.tpu.experimental.embedding.".format(
+                             type(table.optimizer)))
+      if table.name is None:
+        table.name = "table_{}".format(i)
+      if table.name in table_names:
+        raise ValueError("Tables must have a unique name. "
+                         f"Multiple tables with name {table.name} found.")
+      table_names.append(table.name)
+
+    self._built = False
+
+  @property
+  def embedding_tables(self):
+    """Returns a dict of embedding tables, keyed by `TableConfig`."""
+    raise NotImplementedError
+
+  def _create_variables(self, table: tpu_embedding_v2_utils.TableConfig,
+                        trainable: bool) -> Dict[Text, tf_variables.Variable]:
+    """Create all variables including table variables and slot variables."""
+    variable_shape = (table.vocabulary_size, table.dim)
+
+    def getter(name, shape, dtype, initializer, trainable):
+      del shape
+      # _add_variable_with_custom_getter clears the shape sometimes, so we
+      # take the global shape from outside the getter.
+      initial_value = functools.partial(
+          initializer, variable_shape, dtype=dtype)
+      return tf_variables.Variable(
+          name=name,
+          initial_value=initial_value,
+          shape=variable_shape,
+          dtype=dtype,
+          trainable=trainable)
+
+    def variable_creator(name, initializer, trainable=True):
+      # Use add_variable_with_custom_getter here so that we take advantage of
+      # the checkpoint loading to allow restore before the variables get
+      # created which avoids double initialization.
+      return self._add_variable_with_custom_getter(
+          name=name,
+          initializer=initializer,
+          shape=variable_shape,
+          dtype=dtypes.float32,
+          getter=getter,
+          trainable=trainable)
+
+    parameters = variable_creator(
+        table.name, table.initializer, trainable=trainable)
+
+    def slot_creator(name, initializer):
+      return variable_creator(table.name + "/" + name, initializer, False)
+
+    if table.optimizer is not None:
+      slot_vars = table.optimizer._create_slots(parameters, slot_creator)  # pylint: disable=protected-access
+    else:
+      slot_vars = {}
+    slot_vars["parameters"] = parameters
+    return slot_vars
+
+  def _create_variables_and_slots(self):
+    """Create variables and slots variables for TPU embeddings."""
+    raise NotImplementedError
+
+  def build(self):
+    """Create variables and slots variables for TPU embeddings."""
+    if self._built:
+      return
+    self._variables = self._create_variables_and_slots()
+    self._built = True
+
+  def __call__(self, features: Any, weights: Optional[Any] = None) -> Any:
+    """Call the mid level api to do embedding lookup."""
+    if not self._built:
+      self.build()
+    return self.embedding_lookup(features, weights)
+
+  def embedding_lookup(self,
+                       features: Any,
+                       weights: Optional[Any] = None) -> Any:
+    """Lookup the embedding table using the input features."""
+    raise NotImplementedError

tensorflow/python/tpu/tpu_embedding_for_serving.py

@@ -0,0 +1,625 @@
+# 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.
+# ==============================================================================
+"""Mid level API for Serving TPU Embeddings."""
+import functools
+from typing import Any, Dict, Iterable, Optional, Union
+
+from absl import logging
+
+from tensorflow.core.tpu.kernels import sparse_core_layout_pb2
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import tpu_strategy
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework.constant_op import constant as tf_constant
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops import variables as tf_variables
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.tpu import tpu_embedding_base
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu import tpu_embedding_v3_utils
+from tensorflow.python.trackable import base as trackable_base
+from tensorflow.python.types import core
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("tpu.experimental.embedding.TPUEmbeddingForServing")
+class TPUEmbeddingForServing(tpu_embedding_base.TPUEmbeddingBase):
+  """The TPUEmbedding mid level API running on CPU for serving.
+
+  Note: This class is intended to be used for embedding tables that are trained
+  on TPU and to be served on CPU. Therefore the class should be only initialized
+  under non-TPU strategy. Otherwise an error will be raised.
+
+  You can first train your model using the TPUEmbedding class and save the
+  checkpoint. Then use this class to restore the checkpoint to do serving.
+
+  First train a model and save the checkpoint.
+  ```python
+  model = model_fn(...)
+  strategy = tf.distribute.TPUStrategy(...)
+  with strategy.scope():
+    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
+        feature_config=feature_config,
+        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
+
+  # Your custom training code.
+
+  checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
+  checkpoint.save(...)
+
+  ```
+
+  Then restore the checkpoint and do serving.
+  ```python
+
+  # Restore the model on CPU.
+  model = model_fn(...)
+  embedding = tf.tpu.experimental.embedding.TPUEmbeddingForServing(
+        feature_config=feature_config,
+        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
+
+  checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
+  checkpoint.restore(...)
+
+  result = embedding(...)
+  table = embedding.embedding_table
+  ```
+
+  NOTE: This class can also be used to do embedding training on CPU. But it
+  requires the conversion between keras optimizer and embedding optimizers so
+  that the slot variables can stay consistent between them.
+  """
+
+  def __init__(
+      self,
+      feature_config: Union[tpu_embedding_v2_utils.FeatureConfig, Iterable],  # pylint:disable=g-bare-generic
+      optimizer: Optional[tpu_embedding_v2_utils._Optimizer],
+      experimental_sparsecore_restore_info: Optional[Dict[str, Any]] = None,
+  ):  # pylint:disable=protected-access
+    """Creates the TPUEmbeddingForServing mid level API object.
+
+    ```python
+    embedding = tf.tpu.experimental.embedding.TPUEmbeddingForServing(
+        feature_config=tf.tpu.experimental.embedding.FeatureConfig(
+            table=tf.tpu.experimental.embedding.TableConfig(
+                dim=...,
+                vocabulary_size=...)))
+    ```
+
+    Args:
+      feature_config: A nested structure of
+        `tf.tpu.experimental.embedding.FeatureConfig` configs.
+      optimizer: An instance of one of `tf.tpu.experimental.embedding.SGD`,
+        `tf.tpu.experimental.embedding.Adagrad` or
+        `tf.tpu.experimental.embedding.Adam`. When not created under TPUStrategy
+        may be set to None to avoid the creation of the optimizer slot
+        variables, useful for optimizing memory consumption when exporting the
+        model for serving where slot variables aren't needed.
+      experimental_sparsecore_restore_info: Information from the sparse core
+        training, required to restore from checkpoint for serving (like number
+        of TPU devices used `num_tpu_devices`.)
+
+    Raises:
+      RuntimeError: If created under TPUStrategy.
+    """
+    super(TPUEmbeddingForServing, self).__init__(feature_config, optimizer)
+    self._strategy = distribute_lib.get_strategy()
+    if isinstance(
+        self._strategy, (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV2)
+    ):
+      raise RuntimeError("Serving on TPU is not yet supported.")
+
+  @property
+  def variables(
+      self,
+  ) -> Dict[
+      tpu_embedding_v2_utils.TableConfig, Dict[str, tf_variables.Variable]
+  ]:
+    """Returns a dict of variables, keyed by `TableConfig`, then by slot name."""
+    self._maybe_build()
+    return self._variables
+
+  @property
+  def embedding_tables(
+      self,
+  ) -> Dict[tpu_embedding_v2_utils.TableConfig, tf_variables.Variable]:
+    """Returns a dict of embedding tables, keyed by `TableConfig`."""
+    self._maybe_build()
+    # Only return the tables and not the slot variables.
+    return {
+        table: self._variables[table.name]["parameters"]
+        for table in self._table_config
+    }
+
+  def _maybe_build(self):
+    if not self._built:
+      # This can be called while tracing a function, so we wrap the
+      # initialization code with init_scope so it runs eagerly, this means that
+      # it will not be included the function graph generated by tracing so that
+      # we can be sure that we only initialize the TPU for embeddings exactly
+      # once.
+      with ops.init_scope():
+        self.build()
+
+  # TODO(silkyarora) Update the tests for all TPU embedding to expect this
+  # possibly empty information in checkpoints.
+  def _maybe_delete_sc_layouts_from_checkpoint(self):
+    # Remove the sparse_core_table_layouts from the checkpoint, it is only
+    # required for sparsecore.
+    if (
+        hasattr(
+            self,
+            tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY,
+        )
+        and not self._get_sparse_core_table_layouts_str()
+    ):
+      delattr(
+          self,
+          tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY,
+      )
+
+  def build(self):
+    """Create variables and slots variables for TPU embeddings."""
+    super().build()
+    self._maybe_delete_sc_layouts_from_checkpoint()
+
+  def _track_restore_info_for_cpu(self) -> None:
+    def getter(name, shape, dtype, initializer, trainable):
+      del shape
+      # _add_variable_with_custom_getter clears the shape sometimes, so we
+      # take the global shape from outside the getter.
+      initial_value = functools.partial(initializer, dtype=dtype)
+      return tf_variables.Variable(
+          name=name,
+          initial_value=initial_value,
+          shape=None,
+          dtype=dtype,
+          trainable=trainable,
+      )
+
+    def empty_string(dtype: dtypes.DType):
+      return tf_constant("", dtype=dtype)
+
+    # _add_variable_with_custom_getter is used here to restore from checkpoint
+    # at creation time. The layouts from sparse core must be restored from
+    # checkpoint and before any other tables are restored
+    setattr(
+        self,
+        tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY,
+        self._add_variable_with_custom_getter(
+            name=tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY,
+            initializer=empty_string,
+            dtype=dtypes.string,
+            getter=getter,
+            trainable=False,
+        ),
+    )
+
+  def _get_sparse_core_table_layouts_str(self) -> bytes:
+    layouts_str = getattr(
+        self,
+        tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY,
+    )
+    return layouts_str.read_value().numpy()
+
+  def _trackable_children(
+      self, save_type=trackable_base.SaveType.CHECKPOINT, **kwargs: Any
+  ):
+    # Remove the trackables added to make sparsecore checkpoint restore work.
+    # These are not required for serializing the model.
+    tc = super()._trackable_children(save_type, **kwargs)
+    if save_type == trackable_base.SaveType.SAVEDMODEL:
+      if tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY in tc:
+        tc.pop(tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY, None)
+      sclt = [
+          k
+          for k, v in tc.items()
+          if isinstance(
+              v, tpu_embedding_v3_utils.SparseCoreStackedTableTrackable
+          )
+      ]
+      for k in sclt:
+        tc.pop(k, None)
+    return tc
+
+  def _create_variables_from_stacked_tables(self):
+    sc_layouts = sparse_core_layout_pb2.SparseCoreTableLayouts()
+    sc_layouts.ParseFromString(self._get_sparse_core_table_layouts_str())
+    stacked_table_name_to_layouts = {}
+    for layout in sc_layouts.tables:
+      stacked_tables_list = stacked_table_name_to_layouts.setdefault(
+          layout.stacked_table_name, []
+      )
+      stacked_tables_list.append(layout)
+    table_to_config = {table.name: table for table in self._table_config}
+    variables = {}
+    for stacked_table_name, layouts in stacked_table_name_to_layouts.items():
+      logging.info(
+          "Loading stacked table state variables(%s) for %s tables",
+          stacked_table_name,
+          len(layouts),
+      )
+      stacked_var_trackable = (
+          tpu_embedding_v3_utils.SparseCoreStackedTableTrackable(
+              layouts, table_to_config
+          )
+      )
+      # The stacked table is added as trackable to the embedding so that the
+      # checkpoint key corresponsing to stacked table is read.
+      self._track_trackable(stacked_var_trackable, stacked_table_name)
+      variables.update(stacked_var_trackable.get_vars())
+    return variables
+
+  def _create_variables_and_slots(
+      self,
+  ) -> Dict[str, Dict[str, tf_variables.Variable]]:
+    """Create variables for TPU embeddings.
+
+    Returns:
+      A dict of dicts. The outer dict is keyed by the table names and the inner
+      dicts are keyed by 'parameters' and the slot variable names.
+    """
+    self._track_restore_info_for_cpu()
+    variables = {}
+    # If there are stacked variables from SC checkpoint process those
+    # first
+    stacked_variables = self._create_variables_from_stacked_tables()
+    for table in self._table_config:
+      if table.name in stacked_variables:
+        variables[table.name] = {"parameters": stacked_variables[table.name]}
+      else:
+        variables[table.name] = self._create_variables(table, trainable=True)
+    return variables
+
+  def embedding_lookup(
+      self, features: Any, weights: Optional[Any] = None
+  ) -> Any:
+    """Apply standard lookup ops on CPU.
+
+    Args:
+      features: A nested structure of `tf.Tensor`s, `tf.SparseTensor`s or
+        `tf.RaggedTensor`s, with the same structure as `feature_config`. Inputs
+        will be downcast to `tf.int32`. Only one type out of `tf.SparseTensor`
+        or `tf.RaggedTensor` is supported per call.
+      weights: If not `None`, a nested structure of `tf.Tensor`s,
+        `tf.SparseTensor`s or `tf.RaggedTensor`s, matching the above, except
+        that the tensors should be of float type (and they will be downcast to
+        `tf.float32`). For `tf.SparseTensor`s we assume the `indices` are the
+        same for the parallel entries from `features` and similarly for
+        `tf.RaggedTensor`s we assume the row_splits are the same.
+
+    Returns:
+      A nested structure of Tensors with the same structure as input features.
+    """
+    return cpu_embedding_lookup(
+        features, weights, self.embedding_tables, self._feature_config
+    )
+
+
+def _ragged_embedding_lookup_with_reduce(
+    table: tf_variables.Variable,
+    ragged: ragged_tensor.RaggedTensor,
+    weights: ragged_tensor.RaggedTensor,
+    combiner: str,
+) -> core.Tensor:
+  """Compute a ragged lookup followed by a reduce on axis 1.
+
+  Args:
+    table: The embedding table.
+    ragged: A RaggedTensor of ids to look up.
+    weights: A RaggedTensor of weights (or None).
+    combiner: One of "mean", "sum", "sqrtn".
+
+  Returns:
+    A Tensor.
+  """
+  if weights is None:
+    weights = array_ops.ones_like(ragged, dtype=table.dtype)
+  weights = array_ops.expand_dims(weights, axis=2)
+  ragged_result = embedding_ops.embedding_lookup(table, ragged)
+  ragged_result = math_ops.reduce_sum(ragged_result * weights, axis=1)
+  if combiner == "mean":
+    ragged_result = math_ops.div_no_nan(
+        ragged_result, math_ops.reduce_sum(weights, axis=1)
+    )
+  elif combiner == "sqrtn":
+    ragged_result = math_ops.div_no_nan(
+        ragged_result,
+        math_ops.sqrt(math_ops.reduce_sum(weights * weights, axis=1)),
+    )
+  return ragged_result
+
+
+@tf_export("tpu.experimental.embedding.serving_embedding_lookup")
+def cpu_embedding_lookup(
+    inputs: Any,
+    weights: Optional[Any],
+    tables: Dict[tpu_embedding_v2_utils.TableConfig, tf_variables.Variable],
+    feature_config: Union[
+        tpu_embedding_v2_utils.FeatureConfig, Iterable  # pylint:disable=g-bare-generic
+    ],
+) -> Any:
+  """Apply standard lookup ops with `tf.tpu.experimental.embedding` configs.
+
+  This function is a utility which allows using the
+  `tf.tpu.experimental.embedding` config objects with standard lookup functions.
+  This can be used when exporting a model which uses
+  `tf.tpu.experimental.embedding.TPUEmbedding` for serving on CPU. In particular
+  `tf.tpu.experimental.embedding.TPUEmbedding` only supports lookups on TPUs and
+  should not be part of your serving graph.
+
+  Note that TPU specific options (such as `max_sequence_length`) in the
+  configuration objects will be ignored.
+
+  In the following example we take a trained model (see the documentation for
+  `tf.tpu.experimental.embedding.TPUEmbedding` for the context) and create a
+  saved model with a serving function that will perform the embedding lookup and
+  pass the results to your model:
+
+  ```python
+  model = model_fn(...)
+  embedding = tf.tpu.experimental.embedding.TPUEmbedding(
+      feature_config=feature_config,
+      batch_size=1024,
+      optimizer=tf.tpu.experimental.embedding.SGD(0.1))
+  checkpoint = tf.train.Checkpoint(model=model, embedding=embedding)
+  checkpoint.restore(...)
+
+  @tf.function(input_signature=[{'feature_one': tf.TensorSpec(...),
+                                 'feature_two': tf.TensorSpec(...),
+                                 'feature_three': tf.TensorSpec(...)}])
+  def serve_tensors(embedding_features):
+    embedded_features = tf.tpu.experimental.embedding.serving_embedding_lookup(
+        embedding_features, None, embedding.embedding_tables,
+        feature_config)
+    return model(embedded_features)
+
+  model.embedding_api = embedding
+  tf.saved_model.save(model,
+                      export_dir=...,
+                      signatures={'serving_default': serve_tensors})
+
+  ```
+
+  NOTE: It's important to assign the embedding API object to a member of your
+  model as `tf.saved_model.save` only supports saving variables as one
+  `Trackable` object. Since the model's weights are in `model` and the
+  embedding table are managed by `embedding`, we assign `embedding` to an
+  attribute of `model` so that tf.saved_model.save can find the embedding
+  variables.
+
+  NOTE: The same `serve_tensors` function and `tf.saved_model.save` call will
+  work directly from training.
+
+  Args:
+    inputs: a nested structure of Tensors, SparseTensors or RaggedTensors.
+    weights: a nested structure of Tensors, SparseTensors or RaggedTensors or
+      None for no weights. If not None, structure must match that of inputs, but
+      entries are allowed to be None.
+    tables: a dict of mapping TableConfig objects to Variables.
+    feature_config: a nested structure of FeatureConfig objects. The keys of
+      feature_config is a superset of inputs.
+
+  Returns:
+    A nested structure of Tensors with the same structure as inputs.
+  """
+
+  flat_inputs = nest.flatten_with_joined_string_paths(inputs)
+  flat_weights = [None] * len(flat_inputs)
+  if weights is not None:
+    nest.assert_same_structure(inputs, weights)
+    flat_weights = nest.flatten(weights)
+  flat_features = dict(nest.flatten_with_joined_string_paths(feature_config))
+
+  input_keys = {key for key, _ in flat_inputs}
+  if not input_keys.issubset(flat_features.keys()):
+    raise ValueError(
+        "Inputs are not a subset of feature_config. Inputs keys are {}, but"
+        " feature_config keys are {}".format(input_keys, flat_features.keys())
+    )
+  outputs = []
+  for (path, inp), weight in zip(flat_inputs, flat_weights):
+    feature = flat_features[path]
+    table = tables[feature.table]
+
+    if weight is not None:
+      if isinstance(inp, tensor.Tensor):
+        raise ValueError(
+            "Weight specified for {}, but input is dense.".format(path)
+        )
+      elif type(weight) is not type(inp):
+        raise ValueError(
+            "Weight for {} is of type {} but it does not match type of the "
+            "input which is {}.".format(path, type(weight), type(inp))
+        )
+      elif feature.max_sequence_length > 0:
+        raise ValueError(
+            "Weight specified for {}, but this is a sequence feature.".format(
+                path
+            )
+        )
+
+    if isinstance(inp, tensor.Tensor):
+      if feature.max_sequence_length > 0:
+        raise ValueError(
+            "Feature {} is a sequence feature but a dense tensor "
+            "was passed.".format(path)
+        )
+      outputs.append(embedding_ops.embedding_lookup_v2(table, inp))
+
+    elif isinstance(inp, sparse_tensor.SparseTensor):
+      outputs.append(
+          _embedding_lookup_for_sparse_tensor(inp, weight, table, feature)
+      )
+    elif isinstance(inp, ragged_tensor.RaggedTensor):
+      outputs.append(
+          _embedding_lookup_for_ragged_tensor(inp, weight, table, feature)
+      )
+    else:
+      raise ValueError(
+          "Input {} is type {}. Tensor, SparseTensor or "
+          "RaggedTensor expected.".format(path, type(inp))
+      )
+  return nest.pack_sequence_as(inputs, outputs)
+
+
+def _embedding_lookup_for_sparse_tensor(
+    inp: sparse_tensor.SparseTensor,
+    weight: Optional[sparse_tensor.SparseTensor],
+    table: tf_variables.Variable,
+    feature: tpu_embedding_v2_utils.FeatureConfig,
+) -> tensor.Tensor:
+  """Embedding lookup for sparse tensor based on its feature config.
+
+  Args:
+    inp: a single SparseTensor input.
+    weight: None or SparseTensor which has the same shape of the input.
+    table: a table variable.
+    feature: a feature config.
+
+  Returns:
+    Embedding lookup result.
+  """
+  inp_rank = inp.shape.rank
+  # The input rank can be None for sequence input tensor.
+  if (
+      not feature.output_shape
+      and feature.max_sequence_length > 0
+      and (inp_rank is None or inp_rank == 2)
+  ):
+    batch_size = math_ops.cast(array_ops.shape(inp)[0], dtype=dtypes.int64)
+    sparse_shape = array_ops_stack.stack(
+        [batch_size, feature.max_sequence_length], axis=0
+    )
+    # TPU Embedding truncates sequences to max_sequence_length, and if we
+    # don't truncate, scatter_nd will error out if the index was out of
+    # bounds.
+    truncated_inp = sparse_ops.sparse_slice(
+        inp, start=[0, 0], size=sparse_shape
+    )
+
+    dense_output_shape = array_ops_stack.stack(
+        [batch_size, feature.max_sequence_length, feature.table.dim], axis=0
+    )
+    return array_ops.scatter_nd(
+        truncated_inp.indices,
+        array_ops.gather(table.read_value(), truncated_inp.values),
+        dense_output_shape,
+    )
+  else:
+    if feature.max_sequence_length > 0:
+      logging.warning(
+          (
+              "max_sequence_length setting will be ignored because the rank of"
+              " the input tensor is %d which is not 2."
+          ),
+          inp_rank,
+      )
+    if (
+        not feature.validate_weights_and_indices
+        and inp_rank is not None
+        and inp_rank <= 2
+    ):
+      return embedding_ops.embedding_lookup_sparse_v2(
+          table, inp, sp_weights=weight, combiner=feature.table.combiner
+      )
+    else:
+      return embedding_ops.safe_embedding_lookup_sparse_v2(
+          table, inp, sparse_weights=weight, combiner=feature.table.combiner
+      )
+
+
+def _embedding_lookup_for_ragged_tensor(
+    inp: ragged_tensor.RaggedTensor,
+    weight: Optional[ragged_tensor.RaggedTensor],
+    table: tf_variables.Variable,
+    feature: tpu_embedding_v2_utils.FeatureConfig,
+) -> tensor.Tensor:
+  """Embedding lookup for ragged tensor based on its feature config.
+
+  Args:
+    inp: a single rank 2 RaggedTensor input.
+    weight: None or RaggedTensor which has the same shape of the input.
+    table: a table variable.
+    feature: a feature config.
+
+  Returns:
+    Embedding lookup result.
+
+  Raises:
+    ValueError: if input ragged tensor is not rank 2 or output shape set in the
+      feature config doesn't match with the first dim size of the input.
+  """
+  if inp.shape.rank != 2:
+    raise ValueError(
+        "Only rank 2 ragged tensor is supported, but got rank {}".format(
+            inp.shape.rank
+        )
+    )
+  batch_size = inp.shape[0]
+  if feature.output_shape:
+    output_batch_size = math_ops.reduce_prod(feature.output_shape)
+    # If the output batch size matches the data batch size, treat it as
+    # normal ragged input.
+    if output_batch_size == batch_size:
+      ragged_output = _ragged_embedding_lookup_with_reduce(
+          table, inp, weight, feature.table.combiner
+      )
+      ragged_output = array_ops.reshape(
+          ragged_output, shape=feature.output_shape + [feature.table.dim]
+      )
+    # If the data batch size is a factor of the output batch size, the
+    # divide result will be the sequence length. Ignore the weights and
+    # combiner.
+    elif output_batch_size > batch_size and output_batch_size % batch_size == 0:
+      ragged_output = embedding_ops.embedding_lookup_v2(table, inp)
+      # Pad or truncate in the sequence dimension
+      ragged_output = ragged_output.to_tensor(
+          shape=[batch_size, output_batch_size // batch_size, feature.table.dim]
+      )
+      # Reshape to desire output shape.
+      ragged_output = array_ops.reshape(
+          ragged_output, feature.output_shape + [feature.table.dim]
+      )
+    else:
+      raise ValueError(
+          "Output shape set in the FeatureConfig should be the factor of "
+          "the input data batch size. But instead got output shape {}, "
+          "input data batch size {}".format(feature.output_shape, batch_size)
+      )
+  else:
+    if feature.max_sequence_length > 0:
+      output_shape = [
+          batch_size,
+          feature.max_sequence_length,
+          feature.table.dim,
+      ]
+      ragged_lookup = embedding_ops.embedding_lookup_v2(table, inp)
+      # Unlike scatter_nd, RaggedTensor.to_tensor truncates to the given
+      # shape.
+      ragged_output = ragged_lookup.to_tensor(shape=output_shape)
+    else:
+      ragged_output = _ragged_embedding_lookup_with_reduce(
+          table, inp, weight, feature.table.combiner
+      )
+  return ragged_output

tensorflow/python/tpu/tpu_embedding_for_serving_test.py

@@ -0,0 +1,385 @@
+# 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 CPU."""
+
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+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 init_ops_v2
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_for_serving
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.util import nest
+
+
+class TPUEmbeddingForServingTest(test.TestCase):
+
+  def setUp(self):
+    super(TPUEmbeddingForServingTest, 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]
+
+  def _create_mid_level(self):
+    optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+    return tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=self.feature_config, optimizer=optimizer)
+
+  def _get_dense_tensors(self, dtype=dtypes.int32):
+    feature0 = constant_op.constant(self.feature_watched_values, dtype=dtype)
+    feature1 = constant_op.constant(self.feature_favorited_values, dtype=dtype)
+    feature2 = constant_op.constant(self.feature_friends_values, dtype=dtype)
+    return (feature0, feature1, feature2)
+
+  def test_cpu_dense_lookup(self):
+    mid_level = self._create_mid_level()
+    features = self._get_dense_tensors()
+    results = mid_level(features, weights=None)
+    all_lookups = []
+    for feature, config in zip(nest.flatten(features), self.feature_config):
+      table = mid_level.embedding_tables[config.table].numpy()
+      all_lookups.append(table[feature.numpy()])
+    self.assertAllClose(results, nest.pack_sequence_as(results, all_lookups))
+
+  def test_cpu_dense_lookup_with_weights(self):
+    mid_level = self._create_mid_level()
+    features = self._get_dense_tensors()
+    weights = self._get_dense_tensors(dtype=dtypes.float32)
+
+    with self.assertRaisesRegex(
+        ValueError, 'Weight specified for .*, but input is dense.'):
+      mid_level(features, weights=weights)
+
+  def _get_sparse_tensors(self, dtype=dtypes.int32):
+    feature0 = sparse_tensor.SparseTensor(
+        indices=self.feature_watched_indices,
+        values=constant_op.constant(self.feature_watched_values, dtype=dtype),
+        dense_shape=[self.data_batch_size, 2])
+    feature1 = sparse_tensor.SparseTensor(
+        indices=self.feature_favorited_indices,
+        values=constant_op.constant(self.feature_favorited_values, dtype=dtype),
+        dense_shape=[self.data_batch_size, 2])
+    feature2 = sparse_tensor.SparseTensor(
+        indices=self.feature_friends_indices,
+        values=constant_op.constant(self.feature_friends_values, dtype=dtype),
+        dense_shape=[self.data_batch_size, 3])
+    return (feature0, feature1, feature2)
+
+  def test_cpu_sparse_lookup(self):
+    mid_level = self._create_mid_level()
+    features = self._get_sparse_tensors()
+    results = mid_level(features, weights=None)
+    reduced = []
+    for feature, config in zip(nest.flatten(features), self.feature_config):
+      table = mid_level.embedding_tables[config.table].numpy()
+      all_lookups = table[feature.values.numpy()]
+      # With row starts we can use reduceat in numpy. Get row starts from the
+      # ragged tensor API.
+      ragged = ragged_tensor.RaggedTensor.from_sparse(feature)
+      row_starts = ragged.row_starts().numpy()
+      reduced.append(np.add.reduceat(all_lookups, row_starts))
+      if config.table.combiner == 'mean':
+        # for mean, divide by the row lengths.
+        reduced[-1] /= np.expand_dims(ragged.row_lengths().numpy(), axis=1)
+    self.assertAllClose(results, nest.pack_sequence_as(results, reduced))
+
+  def test_cpu_sparse_lookup_with_weights(self):
+    mid_level = self._create_mid_level()
+    features = self._get_sparse_tensors()
+    weights = self._get_sparse_tensors(dtype=dtypes.float32)
+    results = mid_level(features, weights=weights)
+    weighted_sum = []
+    for feature, weight, config in zip(nest.flatten(features),
+                                       nest.flatten(weights),
+                                       self.feature_config):
+      table = mid_level.embedding_tables[config.table].numpy()
+      # Expand dims here needed to broadcast this multiplication properly.
+      weight = np.expand_dims(weight.values.numpy(), axis=1)
+      all_lookups = table[feature.values.numpy()] * weight
+      # With row starts we can use reduceat in numpy. Get row starts from the
+      # ragged tensor API.
+      row_starts = ragged_tensor.RaggedTensor.from_sparse(feature).row_starts()
+      row_starts = row_starts.numpy()
+      weighted_sum.append(np.add.reduceat(all_lookups, row_starts))
+      if config.table.combiner == 'mean':
+        weighted_sum[-1] /= np.add.reduceat(weight, row_starts)
+    self.assertAllClose(results, nest.pack_sequence_as(results,
+                                                       weighted_sum))
+
+  def test_cpu_sparse_lookup_with_non_sparse_weights(self):
+    mid_level = self._create_mid_level()
+    features = self._get_sparse_tensors()
+    weights = self._get_dense_tensors(dtype=dtypes.float32)
+    with self.assertRaisesRegex(
+        ValueError, 'but it does not match type of the input which is'):
+      mid_level(features, weights=weights)
+
+  def _get_ragged_tensors(self, dtype=dtypes.int32):
+    feature0 = ragged_tensor.RaggedTensor.from_row_lengths(
+        values=constant_op.constant(self.feature_watched_values, dtype=dtype),
+        row_lengths=self.feature_watched_row_lengths)
+    feature1 = ragged_tensor.RaggedTensor.from_row_lengths(
+        values=constant_op.constant(self.feature_favorited_values, dtype=dtype),
+        row_lengths=self.feature_favorited_row_lengths)
+    feature2 = ragged_tensor.RaggedTensor.from_row_lengths(
+        values=constant_op.constant(self.feature_friends_values, dtype=dtype),
+        row_lengths=self.feature_friends_row_lengths)
+    return (feature0, feature1, feature2)
+
+  def test_cpu_ragged_lookup_with_weights(self):
+    mid_level = self._create_mid_level()
+    features = self._get_ragged_tensors()
+    weights = self._get_ragged_tensors(dtype=dtypes.float32)
+    results = mid_level(features, weights=weights)
+    weighted_sum = []
+    for feature, weight, config in zip(nest.flatten(features),
+                                       nest.flatten(weights),
+                                       self.feature_config):
+      table = mid_level.embedding_tables[config.table].numpy()
+      # Expand dims here needed to broadcast this multiplication properly.
+      weight = np.expand_dims(weight.values.numpy(), axis=1)
+      all_lookups = table[feature.values.numpy()] * weight
+      row_starts = feature.row_starts().numpy()
+      weighted_sum.append(np.add.reduceat(all_lookups, row_starts))
+      if config.table.combiner == 'mean':
+        weighted_sum[-1] /= np.add.reduceat(weight, row_starts)
+    self.assertAllClose(results, nest.pack_sequence_as(results,
+                                                       weighted_sum))
+
+  def test_cpu_partial_structure_for_features(self):
+    mid_level = self._create_mid_level()
+    # Remove one element of the tuple, the inputs are a subset of
+    # feature_config and will be able to excute.
+    features = tuple(self._get_sparse_tensors()[:2])
+    results = mid_level(features, weights=None)
+    reduced = []
+    for i, feature in enumerate(nest.flatten(features)):
+      config = self.feature_config[i]
+      table = mid_level.embedding_tables[config.table].numpy()
+      all_lookups = table[feature.values.numpy()]
+      # With row starts we can use reduceat in numpy. Get row starts from the
+      # ragged tensor API.
+      ragged = ragged_tensor.RaggedTensor.from_sparse(feature)
+      row_starts = ragged.row_starts().numpy()
+      reduced.append(np.add.reduceat(all_lookups, row_starts))
+      if config.table.combiner == 'mean':
+        # for mean, divide by the row lengths.
+        reduced[-1] /= np.expand_dims(ragged.row_lengths().numpy(), axis=1)
+    self.assertAllClose(results, nest.pack_sequence_as(results, reduced))
+
+  def test_cpu_invalid_structure_for_features(self):
+    mid_level = self._create_mid_level()
+    # Add one element of the tuple, self.feature_config has 3 so we need to
+    # pass no more than 3 elements or None.
+    features = self._get_sparse_tensors() + (self._get_sparse_tensors()[0],)
+    with self.assertRaises(ValueError):
+      mid_level(features, weights=None)
+
+  def test_cpu_invalid_structure_for_weights(self):
+    mid_level = self._create_mid_level()
+    features = self._get_sparse_tensors()
+    # Remove one element of the tuple, inputs has 3 so we need to pass 3 (or
+    # None).
+    weights = tuple(self._get_dense_tensors(dtype=dtypes.float32)[:2])
+    with self.assertRaises(ValueError):
+      mid_level(features, weights=weights)
+
+  def _numpy_sequence_lookup(self, table, indices, values, batch_size,
+                             max_sequence_length, dim):
+    # First we truncate to max_sequence_length.
+    valid_entries = np.nonzero(indices[:, 1] < max_sequence_length)[0]
+    indices = indices[valid_entries]
+    values = values[valid_entries]
+    # Then we gather the values
+    lookup = table[values]
+    # Then we scatter them into the result array.
+    scatter_result = np.zeros([batch_size, max_sequence_length, dim])
+    for i, index in enumerate(indices):
+      scatter_result[index[0], index[1], :] = lookup[i]
+    return scatter_result
+
+  def test_cpu_sequence_lookup_sparse(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)
+    mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=feature_config, optimizer=optimizer)
+    features = self._get_sparse_tensors()[2:3]
+    result = mid_level(features, weights=None)
+
+    golden = self._numpy_sequence_lookup(
+        mid_level.embedding_tables[self.table_user].numpy(),
+        features[0].indices.numpy(),
+        features[0].values.numpy(),
+        self.data_batch_size,
+        feature_config[0].max_sequence_length,
+        self.table_user.dim)
+
+    self.assertAllClose(result[0], golden)
+
+  def test_cpu_sequence_lookup_ragged(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)
+    mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=feature_config, optimizer=optimizer)
+    features = self._get_ragged_tensors()[2:3]
+    result = mid_level(features, weights=None)
+
+    sparse_ver = features[0].to_sparse()
+    golden = self._numpy_sequence_lookup(
+        mid_level.embedding_tables[self.table_user].numpy(),
+        sparse_ver.indices.numpy(),
+        sparse_ver.values.numpy(),
+        self.data_batch_size,
+        feature_config[0].max_sequence_length,
+        self.table_user.dim)
+
+    self.assertAllClose(result[0], golden)
+
+  def test_cpu_high_dimensional_lookup_ragged(self):
+    feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+        table=self.table_user, name='friends', output_shape=[2, 2]),)
+    optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+    mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=feature_config, optimizer=optimizer)
+    features = self._get_ragged_tensors()[2:3]
+    result = mid_level(features, weights=None)
+
+    self.assertAllClose(result[0].shape, (2, 2, 2))
+
+  def test_cpu_high_dimensional_sequence_lookup_ragged(self):
+    # Prod of output shape is a factor of the data batch size.
+    # The divide result will be the sequence length.
+    feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+        table=self.table_user, name='friends', output_shape=[2, 4]),)
+    optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+    mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=feature_config, optimizer=optimizer)
+    features = self._get_ragged_tensors()[2:3]
+    result = mid_level(features, weights=None)
+    self.assertAllClose(result[0].shape, (2, 4, 2))
+
+  def test_cpu_high_dimensional_invalid_lookup_ragged(self):
+    # Prod of output shape is not a factor of the data batch size.
+    # An error will be raised in this case.
+    feature_config = (tpu_embedding_v2_utils.FeatureConfig(
+        table=self.table_user, name='friends', output_shape=[3]),)
+    optimizer = tpu_embedding_v2_utils.SGD(learning_rate=0.1)
+    mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=feature_config, optimizer=optimizer)
+    features = self._get_ragged_tensors()[2:3]
+    with self.assertRaisesRegex(
+        ValueError,
+        'Output shape set in the FeatureConfig should be the factor'):
+      mid_level(features, weights=None)
+
+  def test_cpu_no_optimizer(self):
+    feature_config = (
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=self.table_video, name='watched', max_sequence_length=2),)
+    mid_level = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=feature_config, optimizer=None)
+    # Build the layer manually to create the variables. Normally calling enqueue
+    # would do this.
+    mid_level.build()
+    self.assertEqual(
+        list(mid_level._variables[self.table_video.name].keys()),
+        ['parameters'])
+
+  def test_cpu_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)
+    embedding_one = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=feature_config, optimizer=optimizer)
+    embedding_two = tpu_embedding_for_serving.TPUEmbeddingForServing(
+        feature_config=feature_config, optimizer=optimizer)
+
+    # Both of the tpu embedding tables should be able to build on cpu.
+    embedding_one.build()
+    embedding_two.build()
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tpu_embedding_v1.py

@@ -0,0 +1,445 @@
+# 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.
+# ==============================================================================
+"""Mid level API for TPU Embeddings without Embedding Accelerator."""
+
+from typing import Any, Dict, Iterable, Optional, Text, Union
+
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import tpu_strategy
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops import variables as tf_variables
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.tpu import tpu_embedding_base
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu import tpu_replication
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("tpu.experimental.embedding.TPUEmbeddingV0")
+class TPUEmbeddingV0(tpu_embedding_base.TPUEmbeddingBase):
+  """The TPUEmbedding mid level API running on TPU without Embedding accelerator.
+
+  NOTE: This mid level API is not intended for large embedding table lookup.
+  Embedding tables will be replicated across devices rather than sharding
+  across them. To do large embedding table lookup, please use the
+  `tpu.experimental.embedding.TPUEmbedding` class. This class is an alternative
+  way to do embedding lookups when the TPU doesn't support any version of
+  embedding feature. See
+  `tpu.experimental.tpu_hardware_feature.embedding_feature` for a detailed
+  explanation.
+
+  This class has to be created under the `TPUStrategy`, Otherwise a RuntimeError
+  will be raised.
+  ```python
+  strategy = tf.distribute.TPUStrategy(...)
+  with strategy.scope():
+    embedding = tf.tpu.experimental.embedding.TPUEmbeddingV0(
+        feature_config=feature_config,
+        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
+  ```
+  When creating a distributed dataset that is to be passed to the lookup
+  operation a special input option must be specified:
+
+  ```python
+  distributed_dataset = (
+      strategy.distribute_datasets_from_function(
+          dataset_fn=...,
+          options=tf.distribute.InputOptions(
+              experimental_fetch_to_device=False))
+  dataset_iterator = iter(distributed_dataset)
+  ```
+
+  Below is an example of a training and evaluation step:
+
+  ```python
+  optimizer = tf.keras.optimizers.SGD(0.1)
+
+  @tf.function
+  def training_step(dataset_iterator, num_steps):
+    def tpu_step(embedding_features):
+      with tf.GradientTape() as tape:
+        tape.watch(embedding.embedding_table.values())
+        activation = embedding(embedding_features)
+        model_output = model(activations)
+        loss = ...  # some function of labels and model_output
+
+      embedding_gradients = tape.gradient(loss,
+                                          embedding.embedding_table.values())
+      optimizer.apply_gradients(list(zip(gradients,
+                                mid_level_api.embedding_tables.values())))
+      # Insert your model gradient and optimizer application here
+
+    for _ in tf.range(num_steps):
+      strategy.run(tpu_step, args=(next(dataset_iterator), ))
+
+  @tf.function
+  def evalution_step(dataset_iterator, num_steps):
+    def tpu_step(embedding_features):
+      activations = embedding(embedding_features)
+      model_output = model(activations)
+      # Insert your evaluation code here.
+
+    for _ in tf.range(num_steps):
+      strategy.run(tpu_step, args=(next(dataset_iterator), ))
+  ```
+
+  NOTE: The optimizer used here is a Keras optimizer. In order to make the slot
+  variable creation stay consistent between Keras optimizers and
+  embedding optimizers, the `slot_variable_creation_fn` argument of the
+  embedding optimizers has to be passed with the Keras `add_slot` function. Also
+  note that the slot names might be slightly different between them.
+
+  ```python
+  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.1)
+
+  def slot_variable_creation_fn(table, slot_names, slot_initializers):
+      slots = {}
+      for slot, initializer in zip(slot_names, slot_initializers):
+        slots[slot] = optimizer.add_slot(table, slot, initializer)
+      return slots
+
+  embedding_optimizer = tf.experimental.embedding.Adagrad(
+      learning_rate=0.1,
+      slot_variable_creation_fn=slot_variable_creation_fn)
+
+  # Use the embedding optimizer to create mid level api and keras optimizer to
+  # apply gradients.
+  ```
+  """
+
+  def __init__(
+      self,
+      feature_config: Union[tpu_embedding_v2_utils.FeatureConfig, Iterable],  # pylint:disable=g-bare-generic
+      optimizer: Optional[tpu_embedding_v2_utils._Optimizer]):  # pylint:disable=protected-access
+    super(TPUEmbeddingV0, self).__init__(feature_config, optimizer)
+    self._strategy = distribute_lib.get_strategy()
+    if not isinstance(self._strategy,
+                      (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV2)):
+      raise RuntimeError(
+          "TPUEmbeddingV0 should be created under TPUStrategy but found {}."
+          .format(self._strategy))
+    self._built = False
+
+  @property
+  def embedding_tables(
+      self) -> Dict[tpu_embedding_v2_utils.TableConfig, tf_variables.Variable]:
+    """Returns a dict of embedding tables, keyed by `TableConfig`."""
+    self._maybe_build()
+    # Only return the tables and not the slot variables.
+    return {
+        table: self._variables[table.name]["parameters"]
+        for table in self._table_config
+    }
+
+  def _create_variables_and_slots(
+      self) -> Dict[Text, Dict[Text, tf_variables.Variable]]:
+    """Create variables for TPU embeddings.
+
+    Note that this will always ensure that the variable is created under the
+    TPUStrategy.
+
+    Returns:
+      A dict of dicts. The outer dict is keyed by the table names and the inner
+      dicts are keyed by 'parameters' and the slot variable names.
+    """
+    variables = {}
+    for table in self._table_config:
+      # created TPUDistributedVariable.
+      variables[table.name] = self._create_variables(table, trainable=True)
+    return variables
+
+  def _maybe_build(self):
+    if not self._built:
+      # This can be called while tracing a function, so we wrap the
+      # initialization code with init_scope so it runs eagerly, this means that
+      # it will not be included in the function graph generated by tracing so
+      # that we can be sure that we only initialize the TPU for embeddings
+      # exactly once.
+      with ops.init_scope():
+        self.build()
+
+  def _apply_combiner_to_embeddings(
+      self,
+      embeddings: tensor.Tensor,
+      weight: tensor.Tensor,
+      combiner: Optional[Text] = None) -> tensor.Tensor:
+    """Apply the combiner to the embedding look up result on second to last axis.
+
+    Args:
+      embeddings: A Tensor of the embedding lookup result.
+      weight: A Tensor of weight which has the same shape of the embeddings.
+      combiner: One of "mean", "sum", "sqrtn". Defaults to "mean".
+
+    Raises:
+      ValueError: If the combiner is not one of 'mean', 'sqrtn' or 'sum'.
+    Returns:
+      A Tensor.
+    """
+    if combiner is None:
+      combiner = "mean"
+    if combiner == "sum":
+      embeddings = math_ops.reduce_sum(embeddings, axis=-2)
+    elif combiner == "mean":
+      embeddings = math_ops.reduce_sum(embeddings, axis=-2)
+      weight_sum = math_ops.reduce_sum(weight, axis=-2)
+      embeddings = math_ops.div_no_nan(embeddings, weight_sum)
+    elif combiner == "sqrtn":
+      embeddings = math_ops.reduce_sum(embeddings, axis=-2)
+      weight_squared = math_ops.pow(weight, 2)
+      weight_sum = math_ops.reduce_sum(weight_squared, axis=-2)
+      weight_sum_sqrt = math_ops.sqrt(weight_sum)
+      embeddings = math_ops.div_no_nan(embeddings, weight_sum_sqrt)
+    else:
+      raise ValueError(
+          f"combiner must be one of 'mean', 'sqrtn' or 'sum', got {combiner}")
+    return embeddings
+
+  def _pad_or_truncate_with_sequence_length(
+      self, embeddings: tensor.Tensor, sequence_length: int
+  ) -> tensor.Tensor:
+    """Pad or truncate the embedding lookup result based on the sequence length.
+
+    Args:
+      embeddings: A rank 3 Tensor of the embedding lookup result.
+      sequence_length: number of the max sequence length set in the feature
+        config.
+
+    Returns:
+      A Tensor with second last axis padded or truncated.
+    """
+    original_sequence_length = embeddings.shape[1]
+    if original_sequence_length > sequence_length:
+      embeddings = array_ops.slice(
+          embeddings, begin=[0, 0, 0], size=[-1, sequence_length, -1])
+    else:
+      embeddings = array_ops.pad(
+          embeddings,
+          paddings=[[0, 0], [0, sequence_length - original_sequence_length],
+                    [0, 0]])
+    return embeddings
+
+  def embedding_lookup(self,
+                       features: Any,
+                       weights: Optional[Any] = None) -> Any:
+    """Apply embedding lookup on TPUs using Tensorcore.
+
+    Note that all the sparse and ragged tensors will be converted to dense
+    tensors on CPU and then passed to the TPU to do embedding look up. Large
+    embedding lookup is not supported by this API, use the TPUEmbedding mid
+    level api instead.
+
+    Args:
+      features: a nested structure of Tensors, SparseTensors or RaggedTensors.
+      weights: a nested structure of Tensors, SparseTensors or RaggedTensors or
+        None for no weights. If not None, structure must match that of inputs,
+        but entries are allowed to be None.
+
+    Returns:
+      A nested structure of Tensors with the same structure as inputs.
+    """
+    if not self._built:
+      self.build()
+    nest.assert_same_structure(features, self._feature_config)
+
+    flat_inputs = nest.flatten(features)
+    flat_weights = [None] * len(flat_inputs)
+    if weights is not None:
+      nest.assert_same_structure(features, weights)
+      flat_weights = nest.flatten(weights)
+    flat_features = nest.flatten_with_joined_string_paths(self._feature_config)
+
+    outputs = []
+    for inp, weight, (path, feature) in zip(flat_inputs, flat_weights,
+                                            flat_features):
+      table = self.embedding_tables[feature.table]
+
+      if weight is not None:
+        if isinstance(inp, tensor.Tensor):
+          raise ValueError(
+              "Weight specified for {}, but input is dense.".format(path))
+        elif type(weight) is not type(inp):
+          raise ValueError(
+              "Weight for {} is of type {} but it does not match type of the "
+              "input which is {}.".format(path, type(weight), type(inp)))
+        elif feature.max_sequence_length > 0:
+          raise ValueError("Weight specified for {}, but this is a sequence "
+                           "feature.".format(path))
+
+      if isinstance(inp, tensor.Tensor):
+        if feature.max_sequence_length > 0:
+          raise ValueError(
+              "Feature {} is a sequence feature but a dense tensor "
+              "was passed.".format(path))
+        outputs.append(embedding_ops.embedding_lookup_v2(table, inp))
+
+      elif isinstance(inp, sparse_tensor.SparseTensor):
+        outputs.append(
+            self._embedding_lookup_for_sparse_tensor(inp, weight, table,
+                                                     feature))
+      elif isinstance(inp, ragged_tensor.RaggedTensor):
+        outputs.append(
+            self._embedding_lookup_for_ragged_tensor(inp, weight, table,
+                                                     feature))
+      else:
+        raise ValueError("Input {} is type {}. Tensor, SparseTensor or "
+                         "RaggedTensor expected.".format(path, type(inp)))
+    return nest.pack_sequence_as(self._feature_config, outputs)
+
+  def _embedding_lookup_for_sparse_tensor(
+      self, inp: sparse_tensor.SparseTensor,
+      weight: Optional[sparse_tensor.SparseTensor],
+      table: tf_variables.Variable,
+      feature: tpu_embedding_v2_utils.FeatureConfig) -> tensor.Tensor:
+    """Embedding lookup for sparse tensor based on its feature config.
+
+    Args:
+      inp: a single SparseTensor input.
+      weight: None or SparseTensor which has the same shape of the input.
+      table: a table variable.
+      feature: a feature config.
+
+    Returns:
+      Embedding lookup result.
+    """
+
+    # This computation needs to placed outside of tpu as the size of the
+    # indices and values can change for different batch which can cause
+    # the program to re-compile.
+    def sparse_to_dense_computation(inp, weight):
+      if weight is None:
+        weight = sparse_tensor.SparseTensor(
+            inp.indices,
+            array_ops.ones_like(inp.values, dtype=dtypes.float32),
+            dense_shape=inp.dense_shape)
+      # Pad the sparse tensor to be dense tensor.
+      inp = sparse_ops.sparse_tensor_to_dense(inp)
+      weight = sparse_ops.sparse_tensor_to_dense(weight)
+      return inp, weight
+
+    inp, weight = tpu_replication.outside_compilation(
+        sparse_to_dense_computation, inp=inp, weight=weight)
+
+    embeddings = embedding_ops.embedding_lookup_v2(table, inp)
+    weight = array_ops.expand_dims(weight, -1)
+    embeddings *= weight
+    if not feature.output_shape and feature.max_sequence_length > 0:
+      embeddings = self._pad_or_truncate_with_sequence_length(
+          embeddings, feature.max_sequence_length)
+    else:
+      embeddings = self._apply_combiner_to_embeddings(embeddings, weight,
+                                                      feature.table.combiner)
+    return embeddings
+
+  def _embedding_lookup_for_ragged_tensor(
+      self, inp: ragged_tensor.RaggedTensor,
+      weight: Optional[ragged_tensor.RaggedTensor],
+      table: tf_variables.Variable,
+      feature: tpu_embedding_v2_utils.FeatureConfig) -> tensor.Tensor:
+    """Embedding lookup for ragged tensor based on its feature config.
+
+    Args:
+      inp: a single rank 2 RaggedTensor input.
+      weight: None or RaggedTensor which has the same shape of the input.
+      table: a table variable.
+      feature: a feature config.
+
+    Returns:
+      Embedding lookup result.
+
+    Raises:
+      ValueError: if input ragged tensor is not rank 2 or output shape set in
+      the feature config doesn't match with the first dim size of the input.
+    """
+    if inp.shape.rank != 2:
+      raise ValueError(
+          "Only rank 2 ragged tensor is supported, but got rank {}".format(
+              inp.shape.rank))
+    batch_size = inp.shape[0]
+
+    # This computation needs to placed outside of tpu as the size of the row
+    # splits and values can change for different batch which can cause
+    # the program to re-compile.
+    def ragged_to_dense_outside_compilation(inp, weight, batch_size, feature):
+      if weight is None:
+        weight = ragged_tensor.RaggedTensor.from_row_splits(
+            array_ops.ones_like(inp.values, dtype=dtypes.float32),
+            inp.row_splits)
+      if not feature.output_shape and feature.max_sequence_length > 0:
+        inp = inp.to_tensor(shape=(batch_size, feature.max_sequence_length))
+        # Ignore weight if it is a sequence feature.
+        weight = array_ops.ones_like(inp, dtype=dtypes.float32)
+      elif feature.output_shape:
+        # Eagerly run the following op as the result as to be a number in
+        # order to use it as part of the output shape.
+        with ops.init_scope():
+          output_batch_size = math_ops.reduce_prod(feature.output_shape).numpy()
+        # If the output batch size matches the data batch size, treat it as
+        # normal ragged input.
+        if output_batch_size == batch_size:
+          inp, weight = inp.to_tensor(), weight.to_tensor()
+        # If the data batch size is a factor of the output batch size, the
+        # divide result will be the sequence length. Ignore the weights and
+        # combiner.
+        elif (
+            output_batch_size > batch_size
+            and output_batch_size % batch_size == 0
+        ):
+          # Pad or truncate in the sequence dimension
+          seq_length = output_batch_size // batch_size
+          inp = inp.to_tensor(shape=(batch_size, seq_length))
+          # Ignore weight if it is a sequence feature.
+          weight = array_ops.ones_like(inp, dtype=dtypes.float32)
+        else:
+          raise ValueError(
+              "Output shape set in the FeatureConfig should be the factor of "
+              "the input data batch size. But instead got output shape {}, "
+              "input data batch size {}".format(feature.output_shape,
+                                                batch_size))
+      else:
+        inp, weight = inp.to_tensor(), weight.to_tensor()
+      return inp, weight
+
+    inp, weight = tpu_replication.outside_compilation(
+        ragged_to_dense_outside_compilation,
+        inp=inp,
+        weight=weight,
+        batch_size=batch_size,
+        feature=feature)
+
+    embeddings = embedding_ops.embedding_lookup_v2(table, inp)
+    weight = array_ops.expand_dims(weight, -1)
+    embeddings *= weight
+
+    if feature.output_shape:
+      with ops.init_scope():
+        output_batch_size = math_ops.reduce_prod(feature.output_shape).numpy()
+      if output_batch_size == batch_size:
+        embeddings = self._apply_combiner_to_embeddings(embeddings, weight,
+                                                        feature.table.combiner)
+      embeddings = array_ops.reshape(
+          embeddings, shape=feature.output_shape + [feature.table.dim])
+    else:
+      if feature.max_sequence_length == 0:
+        embeddings = self._apply_combiner_to_embeddings(embeddings, weight,
+                                                        feature.table.combiner)
+    return embeddings

tensorflow/python/tpu/tpu_embedding_v2_utils_test.py

@@ -0,0 +1,237 @@
+# 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 utils on TPU."""
+
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf.tpu import tpu_embedding_configuration_pb2
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.eager import def_function
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+
+
+class TPUEmbeddingOptimizerTest(parameterized.TestCase, test.TestCase):
+
+  @parameterized.parameters(tpu_embedding_v2_utils.Adagrad,
+                            tpu_embedding_v2_utils.Adam,
+                            tpu_embedding_v2_utils.FTRL)
+  def test_grad_clip_with_accumulation_off(self, optimizer):
+    with self.assertRaisesRegex(ValueError, 'accumulation'):
+      optimizer(use_gradient_accumulation=False, clipvalue=0.)
+    with self.assertRaisesRegex(ValueError, 'accumulation'):
+      optimizer(use_gradient_accumulation=False, clipvalue=(None, 1.))
+
+  @parameterized.parameters(tpu_embedding_v2_utils.SGD,
+                            tpu_embedding_v2_utils.Adagrad,
+                            tpu_embedding_v2_utils.Adam,
+                            tpu_embedding_v2_utils.FTRL)
+  def test_grad_clip_with_tuple(self, optimizer):
+    opt = optimizer(clipvalue=(-1., 1.))
+    self.assertEqual(-1., opt.clip_gradient_min)
+    self.assertEqual(1., opt.clip_gradient_max)
+
+  @parameterized.parameters(tpu_embedding_v2_utils.SGD,
+                            tpu_embedding_v2_utils.Adagrad,
+                            tpu_embedding_v2_utils.Adam,
+                            tpu_embedding_v2_utils.FTRL)
+  def test_grad_clip_with_single_value(self, optimizer):
+    opt = optimizer(clipvalue=1.)
+    self.assertEqual(-1., opt.clip_gradient_min)
+    self.assertEqual(1., opt.clip_gradient_max)
+
+  @parameterized.parameters(tpu_embedding_v2_utils.SGD,
+                            tpu_embedding_v2_utils.Adagrad,
+                            tpu_embedding_v2_utils.Adam,
+                            tpu_embedding_v2_utils.FTRL)
+  def test_grad_clip_with_tuple_and_none(self, optimizer):
+    opt = optimizer(clipvalue=(None, 1))
+    self.assertIsNone(opt.clip_gradient_min)
+    self.assertEqual(1., opt.clip_gradient_max)
+
+  @parameterized.parameters(tpu_embedding_v2_utils.SGD,
+                            tpu_embedding_v2_utils.Adagrad,
+                            tpu_embedding_v2_utils.Adam,
+                            tpu_embedding_v2_utils.FTRL)
+  def test_equal_and_hash_function(self, optimizer):
+    opt1 = optimizer(0.1)
+    opt2 = optimizer(0.1)
+    opt3 = optimizer(0.2)
+    self.assertEqual(opt1, opt2)
+    self.assertEqual(hash(opt1), hash(opt2))
+    self.assertNotEqual(opt1, opt3)
+    self.assertNotEqual(hash(opt1), hash(opt3))
+
+
+class TPUEmbeddingCustomCombinerTest(parameterized.TestCase, test.TestCase):
+
+  def get_sum_combiner(self):
+
+    @def_function.function
+    def sum_combiner(valency, vectors):
+      max_valency = vectors.shape[0]
+      valid_mask = array_ops.range(max_valency) < valency
+      vectors_masked = array_ops.where(
+          array_ops.expand_dims(valid_mask, axis=-1),
+          vectors,
+          array_ops.zeros_like(vectors),
+      )
+      return math_ops.reduce_sum(vectors_masked, axis=0)
+
+    return sum_combiner
+
+  def get_positional_weight_combiner(self):
+
+    @def_function.function
+    def positional_weight_combiner(valency, vectors, weights):
+      max_valency = vectors.shape[0]
+      valid_mask = array_ops.range(max_valency) < valency
+      vectors_masked = array_ops.where(
+          array_ops.expand_dims(valid_mask, axis=-1),
+          vectors,
+          array_ops.zeros_like(vectors),
+      )
+      return math_ops.matvec(vectors_masked, weights, transpose_a=True)
+
+    return positional_weight_combiner
+
+  def test_zero_num_weights_combiner_has_no_slots(self):
+    combiner = tpu_embedding_v2_utils.CustomCombiner(
+        self.get_sum_combiner(),
+        max_valency=16,
+        num_weights=0,
+    )
+    self.assertEmpty(combiner._slot_names())
+    self.assertEmpty(combiner._slot_initializers())
+
+  def test_name_starts_with_custom_combiner(self):
+    combiner = tpu_embedding_v2_utils.CustomCombiner(
+        self.get_sum_combiner(),
+        max_valency=16,
+    )
+    self.assertStartsWith(str(combiner), 'custom_combiner')
+
+  def test_non_zero_weights_requires_initializer(self):
+    with self.assertRaisesRegex(ValueError, '`initializer` must be set'):
+      tpu_embedding_v2_utils.CustomCombiner(
+          self.get_positional_weight_combiner(),
+          max_valency=16,
+          num_weights=16,
+      )
+
+  def test_non_zero_weights_has_one_slot_variable(self):
+    combiner = tpu_embedding_v2_utils.CustomCombiner(
+        self.get_positional_weight_combiner(),
+        max_valency=16,
+        num_weights=16,
+        initializer=init_ops_v2.zeros_initializer,
+    )
+    self.assertLen(combiner._slot_names(), 1)
+    self.assertLen(combiner._slot_initializers(), 1)
+    self.assertStartsWith(combiner._slot_names()[0], 'custom_combiner')
+
+
+class ConfigTest(test.TestCase):
+
+  def test_table_config_repr(self):
+    table = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=2, dim=4,
+        combiner='sum', name='table')
+
+    self.assertEqual(
+        repr(table),
+        'TableConfig(vocabulary_size=2, dim=4, initializer=None, '
+        'optimizer=None, combiner=\'sum\', name=\'table\', '
+        'quantization_config=None)')
+
+  def test_feature_config_repr(self):
+    table = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=2, dim=4, initializer=None,
+        combiner='sum', name='table')
+
+    feature_config = tpu_embedding_v2_utils.FeatureConfig(
+        table=table, output_shape=[16, 4], name='feature')
+
+    self.assertEqual(
+        repr(feature_config),
+        'FeatureConfig(table=TableConfig(vocabulary_size=2, dim=4, '
+        'initializer=None, optimizer=None, combiner=\'sum\', '
+        'name=\'table\', quantization_config=None), max_sequence_length=0, '
+        'validate_weights_and_indices=True, output_shape=TensorShape([16, 4]), '
+        'name=\'feature\')')
+
+  def test_quantization_config_num_buckets(self):
+    with self.assertRaisesRegex(ValueError, 'num_buckets'):
+      tpu_embedding_v2_utils.QuantizationConfig(0, -1, 1)
+
+  def test_quantization_config_repr(self):
+    quantization_config = tpu_embedding_v2_utils.QuantizationConfig(
+        num_buckets=10, lower=-1.0, upper=1.0)
+
+    self.assertEqual(
+        repr(quantization_config),
+        'QuantizationConfig(num_buckets=10, lower=-1.0, upper=1.0)')
+
+
+class TPUEmbeddingConfigurationTest(test.TestCase):
+
+  def test_no_truncate(self):
+    truncate_length = 14937  # Experimentally maximum string length loggable.
+
+    config = tpu_embedding_configuration_pb2.TPUEmbeddingConfiguration()
+    for i in range(500):
+      td = config.table_descriptor.add()
+      td.name = 'table_{}'.format(i)
+      td.vocabulary_size = i
+    config.num_hosts = 2
+    config.num_tensor_cores = 4
+    config.batch_size_per_tensor_core = 128
+
+    self.assertGreater(
+        len(str(config)), truncate_length,
+        'Test sanity check: generated config should be of truncating length.')
+
+    with self.assertLogs() as logs:
+      tpu_embedding_v2_utils.log_tpu_embedding_configuration(config)
+
+    self.assertIn('table_499', ''.join(logs.output))
+    for line in logs.output:
+      self.assertLess(
+          len(line), truncate_length,
+          'Logging function lines should not be of truncating length.')
+
+
+class TPUEmbeddingUtilityFunctionTest(test.TestCase):
+
+  def test_sort_device_spec_strings(self):
+    device_spec_strings = []
+    for task in [2, 3, 0, 1]:  # Intentionally permuted
+      for device in range(8):
+        device_spec_strings.append(
+            f'/job:trainer/replica:0/task:{task}/device:TPU:{device}'
+        )
+    sorted_specs = tpu_embedding_v2_utils._sort_device_spec_strings(
+        device_spec_strings
+    )
+
+    self.assertEqual(sorted_specs, sorted(device_spec_strings))
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tpu_embedding_v3_additional_test.py

@@ -0,0 +1,501 @@
+# Copyright 2024 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.
+# ==============================================================================
+"""Additional multi/single worker tests for tpu_embedding_v3."""
+
+import os
+
+from absl import flags
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.checkpoint import checkpoint as tf_checkpoint
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.distribute import distribute_lib
+from tensorflow.python.distribute import tpu_strategy
+from tensorflow.python.distribute import values as values_lib
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import remote
+from tensorflow.python.framework import config
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework.constant_op import constant as tf_constant
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+from tensorflow.python.saved_model import load as tf_load
+from tensorflow.python.saved_model import save as tf_save
+from tensorflow.python.tpu import device_assignment as device_assignment_lib
+from tensorflow.python.tpu import tpu_embedding_for_serving
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu import tpu_embedding_v3
+from tensorflow.python.tpu import tpu_embedding_v3_utils
+
+
+_TPU = flags.DEFINE_string('tpu', None, 'The TPU to use for TPUStrategy.')
+# pylint: disable=g-long-lambda
+
+
+RowIdInitializer = tpu_embedding_v2_utils.RowIdInitializer
+shard_initializer = tpu_embedding_v3_utils.shard_initializer
+
+
+def get_replica_values(per_replica_or_tensor):
+  if isinstance(per_replica_or_tensor, values_lib.PerReplica):
+    return per_replica_or_tensor.values
+  else:
+    return [per_replica_or_tensor]
+
+
+def pad_to_shape_initializer(init_mat):
+  """An initializer that pads init_mat out to the given shape."""
+  return lambda shape, dtype: array_ops.pad(
+      init_mat,
+      [
+          [0, shape[0] - init_mat.shape[0]],
+          [0, shape[1] - init_mat.shape[1]],
+      ],
+      'CONSTANT',
+  )
+
+
+class TPUEmbeddingLayerV2Test(parameterized.TestCase, test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.vocabulary_size = 128
+    self.embedding_dim = 8
+
+    resolver = tpu_cluster_resolver.TPUClusterResolver(tpu=_TPU.value)
+    if _TPU.value is None:
+      remote.connect_to_cluster(resolver)
+    tpu_cluster_resolver.initialize_tpu_system(resolver)
+
+    # FIXME(b/303466959): Remove this device assignment after TPUStrategy
+    # can follow the actual device ordering under SC.
+    topology = tpu_cluster_resolver.initialize_tpu_system(resolver)
+    tpu_metadata = resolver.get_tpu_system_metadata()
+
+    device_assignment = device_assignment_lib.DeviceAssignment.build(
+        topology, num_replicas=tpu_metadata.num_cores
+    )
+    self._strategy = tpu_strategy.TPUStrategyV2(
+        resolver, experimental_device_assignment=device_assignment
+    )
+
+    self.addCleanup(tpu_cluster_resolver.shutdown_tpu_system, resolver)
+    self.sharded_row_id_initializer = lambda offset: shard_initializer(
+        self._strategy, RowIdInitializer(offset)
+    )
+
+    self.sharded_pad_to_shape_initializer = lambda init_mat: shard_initializer(
+        self._strategy, pad_to_shape_initializer(init_mat)
+    )
+
+    self.table_video = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=self.vocabulary_size,
+        dim=self.embedding_dim,
+        initializer=self.sharded_row_id_initializer(0),
+        combiner='sum',
+        name='video',
+    )
+    self.table_user = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=self.vocabulary_size,
+        dim=self.embedding_dim,
+        initializer=self.sharded_row_id_initializer(1000),
+        combiner='sum',
+        name='user',
+    )
+    self.feature_video = tpu_embedding_v2_utils.FeatureConfig(
+        table=self.table_video,
+        name='video',
+        output_shape=[self.vocabulary_size],
+    )
+    self.feature_user = tpu_embedding_v2_utils.FeatureConfig(
+        table=self.table_user, name='user', output_shape=[self.vocabulary_size]
+    )
+
+    self.assertEqual(
+        self._strategy.extended._tpu_devices.shape, (tpu_metadata.num_cores, 1)
+    )
+
+  def testSingleTableInitializeAndLookup(self):
+    # This test sets up devices to lookup the entire table.
+
+    feature_config = [self.feature_video]
+
+    strategy = self._strategy
+
+    with strategy.scope():
+      embedding_layer = tpu_embedding_v3.TPUEmbeddingV2(
+          feature_config,
+          tpu_embedding_v2_utils.Adagrad(),
+          pipeline_execution_with_tensor_core=True,
+      )
+
+      @def_function.function
+      def train_step(features):
+        def train_step_fn(features):
+          return embedding_layer(features)[0]
+
+        return strategy.run(train_step_fn, args=(features,))
+
+    def value_fn(ctx):
+      del ctx  # unused
+      return [
+          sparse_tensor.SparseTensor(
+              indices=[[i, 0] for i in range(0, self.vocabulary_size)],
+              values=np.arange(0, self.vocabulary_size),
+              dense_shape=[self.vocabulary_size, 1],
+          ),
+      ]
+
+    features = strategy.experimental_distribute_values_from_function(value_fn)
+
+    [embeddings] = train_step(features)
+
+    expected = RowIdInitializer(0)(
+        shape=(self.vocabulary_size, self.embedding_dim), dtype=dtypes.float32
+    )
+    for replica in get_replica_values(embeddings):
+      self.assertAllClose(replica, expected)
+
+  def testStackedTableInitializeAndLookup(self):
+    # This test sets up devices to lookup the entire table.
+    feature_config = [self.feature_video, self.feature_user]
+
+    strategy = self._strategy
+
+    def value_fn(ctx):
+      del ctx  # unused
+      return [
+          sparse_tensor.SparseTensor(
+              indices=[[i, 0] for i in range(0, self.vocabulary_size)],
+              values=np.arange(0, self.vocabulary_size),
+              dense_shape=[self.vocabulary_size, 1],
+          ),
+          sparse_tensor.SparseTensor(
+              indices=[[i, 0] for i in range(0, self.vocabulary_size)],
+              values=np.arange(0, self.vocabulary_size),
+              dense_shape=[self.vocabulary_size, 1],
+          ),
+      ]
+
+    features = strategy.experimental_distribute_values_from_function(value_fn)
+
+    with strategy.scope():
+      embedding_layer = tpu_embedding_v3.TPUEmbeddingV2(
+          feature_config,
+          tpu_embedding_v2_utils.Adagrad(),
+          pipeline_execution_with_tensor_core=True,
+      )
+
+      @def_function.function
+      def train_step(features):
+        def train_step_fn(features):
+          return embedding_layer(features)[0]
+
+        return strategy.run(train_step_fn, args=(features,))
+
+    [embeddings_video, embeddings_user] = train_step(features)
+
+    expected_video = RowIdInitializer(0)(
+        shape=(self.vocabulary_size, self.embedding_dim), dtype=dtypes.float32
+    )
+    expected_user = RowIdInitializer(1000)(
+        shape=(self.vocabulary_size, self.embedding_dim), dtype=dtypes.float32
+    )
+    for replica_video, replica_user in zip(
+        get_replica_values(embeddings_video),
+        get_replica_values(embeddings_user),
+    ):
+      self.assertAllClose(replica_video, expected_video)
+      self.assertAllClose(replica_user, expected_user)
+
+  def testTwoTablesStackedHaveCorrectInitialValues(self):
+    table1_initial_value = np.arange(
+        start=100, stop=120, dtype=np.float32
+    ).reshape([10, 2])
+    table1 = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=10,
+        dim=2,
+        initializer=self.sharded_pad_to_shape_initializer(table1_initial_value),
+        combiner='sum',
+        name='table1',
+    )
+    table2_initial_value = np.arange(
+        start=200, stop=240, dtype=np.float32
+    ).reshape([20, 2])
+    table2 = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=20,
+        dim=2,
+        initializer=self.sharded_pad_to_shape_initializer(table2_initial_value),
+        combiner='sum',
+        name='table2',
+    )
+    feature_configs = [
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=table1, name='feature1', output_shape=[16]
+        ),
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=table2, name='feature2', output_shape=[16]
+        ),
+    ]
+
+    with self._strategy.scope():
+      mid_level_api = tpu_embedding_v3.TPUEmbeddingV2(
+          feature_config=feature_configs,
+          optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1.0),
+          pipeline_execution_with_tensor_core=True,
+      )
+      # The two tables should be stacked into the same variable.
+      self.assertLen(mid_level_api.embedding_tables, 1)
+
+  def testCpuRestoreForNoStackedTables(self):
+    table1_initial_value = np.arange(
+        start=100, stop=120, dtype=np.float32
+    ).reshape([10, 2])
+    table1 = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=12,
+        dim=2,
+        initializer=self.sharded_pad_to_shape_initializer(table1_initial_value),
+        combiner='sum',
+        name='table1',
+    )
+    table2_initial_value = np.arange(
+        start=200, stop=380, dtype=np.float32
+    ).reshape([20, 9])
+    table2 = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=20,
+        dim=9,  # to ensure stacking does not occur
+        initializer=self.sharded_pad_to_shape_initializer(table2_initial_value),
+        combiner='sum',
+        name='table2',
+    )
+    feature_configs = [
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=table1, name='feature1', output_shape=[16]
+        ),
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=table2, name='feature2', output_shape=[16]
+        ),
+    ]
+
+    with self._strategy.scope():
+      mid_level_api = tpu_embedding_v3.TPUEmbeddingV2(
+          feature_config=feature_configs,
+          optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1.0),
+      )
+      # The two tables should be *not* be stacked.
+      self.assertLen(mid_level_api.embedding_tables, 2)
+      # Save v3 embedding
+      checkpoint = tf_checkpoint.Checkpoint(mid_level_api)
+      checkpoint_prefix = os.path.join(self.create_tempdir().full_path, 'ckpt')
+      checkpoint_path = checkpoint.save(checkpoint_prefix)
+
+    # Restore in serving embedding
+    with distribute_lib.get_strategy().scope():
+      serving_embedding = tpu_embedding_for_serving.TPUEmbeddingForServing(
+          feature_config=feature_configs,
+          optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1.0),
+      )
+      checkpoint_for_restore = tf_checkpoint.Checkpoint(serving_embedding)
+      checkpoint_for_restore.restore(checkpoint_path)
+      serving_embedding.build()
+    # Check that 2 tables exist in serving.
+    self.assertLen(serving_embedding.embedding_tables, 2)
+    look_for_row_idx_1 = [
+        sparse_tensor.SparseTensor(
+            indices=[[0, 0]], values=[1], dense_shape=[1, 1]
+        ),
+        sparse_tensor.SparseTensor(
+            indices=[[0, 0]], values=[1], dense_shape=[1, 1]
+        ),
+    ]
+    row_lookup = serving_embedding(look_for_row_idx_1)
+    self.assertAllEqual(
+        row_lookup[0],
+        tf_constant(
+            [
+                102.0,
+                103.0,
+            ],
+            shape=(1, 2),
+        ),
+    )
+    self.assertAllEqual(
+        row_lookup[1],
+        tf_constant(
+            [209.0, 210.0, 211.0, 212.0, 213.0, 214.0, 215.0, 216.0, 217.0],
+            shape=(1, 9),
+        ),
+    )
+
+  def testCpuRestoreForStackedTables(self):
+    table1_initial_value = np.arange(
+        start=100, stop=120, dtype=np.float32
+    ).reshape([10, 2])
+    table1 = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=12,
+        dim=2,
+        initializer=self.sharded_pad_to_shape_initializer(table1_initial_value),
+        combiner='sum',
+        name='table1',
+    )
+    table2_initial_value = np.arange(
+        start=200, stop=240, dtype=np.float32
+    ).reshape([20, 2])
+    table2 = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=20,
+        dim=2,
+        initializer=self.sharded_pad_to_shape_initializer(table2_initial_value),
+        combiner='sum',
+        name='table2',
+    )
+    feature_configs = [
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=table1, name='feature1', output_shape=[16]
+        ),
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=table2, name='feature2', output_shape=[16]
+        ),
+    ]
+
+    with self._strategy.scope():
+      mid_level_api = tpu_embedding_v3.TPUEmbeddingV2(
+          feature_config=feature_configs,
+          optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1.0),
+      )
+      # The two tables should be stacked into the same variable.
+      self.assertLen(mid_level_api.embedding_tables, 1)
+      # Save v3 embedding
+      checkpoint = tf_checkpoint.Checkpoint(mid_level_api)
+      checkpoint_prefix = os.path.join(self.create_tempdir().full_path, 'ckpt')
+      checkpoint_path = checkpoint.save(checkpoint_prefix)
+
+    # Restore in serving embedding
+    with distribute_lib.get_strategy().scope():
+      serving_embedding = tpu_embedding_for_serving.TPUEmbeddingForServing(
+          feature_config=feature_configs,
+          optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1.0),
+      )
+      checkpoint_for_restore = tf_checkpoint.Checkpoint(serving_embedding)
+      checkpoint_for_restore.restore(checkpoint_path)
+      serving_embedding.build()
+    # Check that unstacking happens on restore
+    self.assertLen(serving_embedding.embedding_tables, 2)
+    look_for_row_idx_1 = [
+        sparse_tensor.SparseTensor(
+            indices=[[0, 0]], values=[1], dense_shape=[1, 1]
+        ),
+        sparse_tensor.SparseTensor(
+            indices=[[0, 0]], values=[1], dense_shape=[1, 1]
+        ),
+    ]
+    row_lookup = serving_embedding(look_for_row_idx_1)
+    self.assertAllEqual(
+        row_lookup[0],
+        tf_constant([102.0, 103.0], shape=(1, 2)),
+    )
+    self.assertAllEqual(
+        row_lookup[1],
+        tf_constant([202.0, 203.0], shape=(1, 2)),
+    )
+    saved_model_path = os.path.join(
+        self.create_tempdir().full_path, 'saved_model'
+    )
+    tf_save.save(
+        serving_embedding, saved_model_path
+    )
+    loaded_embedding = tf_load.load(saved_model_path)
+    self.assertLen(loaded_embedding._variables, 2)
+
+  def testUnshardedToTpuRestore(self):
+    table1 = tpu_embedding_v2_utils.TableConfig(
+        vocabulary_size=25,
+        dim=6,
+        initializer=self.sharded_row_id_initializer(0),
+        combiner='sum',
+        name='table1',
+    )
+    feature_configs = [
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=table1, name='feature1', output_shape=[16]
+        ),
+    ]
+
+    with distribute_lib.get_strategy().scope():
+      cpu_embedding = tpu_embedding_for_serving.TPUEmbeddingForServing(
+          feature_config=feature_configs,
+          optimizer=tpu_embedding_v2_utils.Adagrad(0.1),
+      )
+
+      self.assertLen(cpu_embedding.embedding_tables, 1)
+      # Save unsharded embedding
+      checkpoint = tf_checkpoint.Checkpoint(cpu_embedding)
+      checkpoint_prefix = os.path.join(self.create_tempdir().full_path, 'ckpt')
+      checkpoint_path = checkpoint.save(checkpoint_prefix)
+
+    # Restore in TPU embedding
+    strategy = self._strategy
+    with strategy.scope():
+      mid_level_api = tpu_embedding_v3.TPUEmbeddingV2(
+          feature_config=feature_configs,
+          optimizer=tpu_embedding_v2_utils.Adagrad(0.1),
+      )
+      checkpoint_for_restore = tf_checkpoint.Checkpoint(mid_level_api)
+      checkpoint_for_restore.restore(checkpoint_path)
+      mid_level_api.build()
+
+    replicas, cores_per_replica = strategy.extended._tpu_devices.shape
+    total_sc_shards = (
+        replicas * cores_per_replica * mid_level_api._num_sc_per_chip
+    )
+    def get_padded_vocab_size(total_sc_shards, vocab_size):
+      # multiple of (8 * total_sc_shards) which is greater than or equal to
+      # vocab size
+      result = 0
+      while result < vocab_size:
+        result += 8 * total_sc_shards
+      return result
+
+    padded_vocab = get_padded_vocab_size(
+        total_sc_shards, table1.vocabulary_size
+    )
+    unsharded_full_value = cpu_embedding._variables['table1']['parameters']
+    shard_shape = [padded_vocab // total_sc_shards, 8]
+    offset = 0
+    ordered_devices = []
+    for devices in strategy.extended._tpu_devices:  # pylint: disable=protected-access
+      ordered_devices.extend(devices)
+    for device in ordered_devices:
+      partition = []
+      for _ in range(mid_level_api._num_sc_per_chip):
+        sh = unsharded_full_value[offset::total_sc_shards, :]
+        padded_sh = pad_to_shape_initializer(sh)(shard_shape, dtypes.float32)
+        partition.append(padded_sh)
+        offset += 1
+      # Check value at each partition
+      self.assertAllEqual(
+          mid_level_api._variables['table1']['parameters'].read_from_device(
+              device
+          ),
+          array_ops.concat(partition, axis=0),
+      )
+
+
+if __name__ == '__main__':
+  v2_compat.enable_v2_behavior()
+  config.enable_mlir_bridge()
+  test.main()

tensorflow/python/tpu/tpu_embedding_v3_checkpoint_adapter.py

@@ -0,0 +1,591 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Checkpoint adapter for TPUEmbedding."""
+
+import collections
+import time
+from typing import Mapping, Optional, Sequence
+
+from absl import logging
+
+from tensorflow.core.tpu.kernels import sparse_core_layout_pb2
+from tensorflow.python.checkpoint import checkpoint_adapter
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.tpu import tpu_embedding_v3_utils
+from tensorflow.python.trackable import base as trackable_base
+from tensorflow.python.training import py_checkpoint_reader
+from tensorflow.python.util.protobuf import compare
+
+
+def _parse_shard_info_str(
+    spec: str,
+) -> tuple[list[int], trackable_base.ShardInfo]:
+  """Parses shape and shard_info string."""
+
+  shape = [int(x) for x in spec.split()[:-1]]
+  slices = spec.split()[-1].split(":")
+  offset = [int(x.split(",")[0]) for x in slices]
+  shard_shape = [int(x.split(",")[1]) for x in slices]
+  return shape, trackable_base.ShardInfo(offset=offset, shape=shard_shape)
+
+
+def _shard_info_str(shape, shard_info) -> str:
+  """Created shape and shard_info string."""
+
+  full_shape_str = " ".join("%d" % d for d in shape) + " "
+  slice_spec = ":".join(
+      "%d,%d" % (o, s) for o, s in zip(shard_info.offset, shard_info.shape)
+  )
+  return full_shape_str + slice_spec
+
+
+def _shard_from_cpu_to_sc(
+    feature_values: tensor.Tensor,
+    shape_and_slice: str,
+    to_shard_layout: Sequence[sparse_core_layout_pb2.SparseCoreTableLayout],
+) -> tensor.Tensor:
+  """Shards the feature tables from CPU to SparseCore."""
+
+  def pad_value(value, variable_shape, table_shape):
+    return array_ops.pad(
+        value,
+        [
+            [0, variable_shape[0] - table_shape[0]],
+            [0, variable_shape[1] - table_shape[1]],
+        ],
+        "CONSTANT",
+    )
+
+  var_full_shape, shard_info = _parse_shard_info_str(shape_and_slice)
+  if shard_info.offset > var_full_shape:
+    raise ValueError(
+        "Invalid shard offset: {}. Offset should be less than the full shape"
+        " of the variable: {}".format(
+            shard_info.offset,
+            var_full_shape,
+        )
+    )
+  num_sc_per_partition = (
+      to_shard_layout[0].num_sparse_cores // to_shard_layout[0].num_partitions
+  )
+
+  total_rows_per_sc = to_shard_layout[0].total_rows_per_sparse_core_shard
+  total_rows_per_partition = total_rows_per_sc * num_sc_per_partition
+  full_values = {}
+  if (shard_info.shape[0] % total_rows_per_partition) != 0:
+    raise ValueError(
+        "Invalid shard shape: {}. Number of rows in input shard slice should"
+        " be multiple of number of rows in a partition({})".format(
+            shard_info.shape,
+            total_rows_per_partition,
+        )
+    )
+  # From the shard info, get the row offsets corresponding to the slice
+  # being looked up.
+  required_shard_offsets = range(
+      shard_info.offset[0],
+      shard_info.offset[0] + shard_info.shape[0],
+      total_rows_per_partition,
+  )
+  output_shards = []
+  for required_shard_offset in required_shard_offsets:
+    sharded_tensors = []
+    for i in range(num_sc_per_partition):
+      shard_idx = (required_shard_offset // total_rows_per_sc) + i
+      for table_idx, layout in enumerate(to_shard_layout):
+        if table_idx not in full_values:
+          full_values[table_idx] = pad_value(
+              feature_values[table_idx],
+              layout.unsharded_padded_shape,
+              layout.unsharded_shape,
+          )
+
+        table_value = full_values[table_idx]
+        # Apply rotation to get this table's shard index
+        table_shard_offset = (
+            shard_idx
+            + (layout.num_sparse_cores - layout.sparse_core_shard_rotation)
+        ) % layout.num_sparse_cores
+        sharded_tensors.append(
+            table_value[
+                table_shard_offset :: layout.num_sparse_cores,
+                :,
+            ]
+        )
+    output_shards.append(array_ops.concat(sharded_tensors, axis=0))
+    logging.vlog(
+        1,
+        "_shard_from_cpu_to_sc: last output_shards.shape: %s",
+        output_shards[-1].shape,
+    )
+  return array_ops.concat(output_shards, axis=0)
+
+
+def _unshard_from_sc_to_cpu(
+    stacked_table: tensor.Tensor,
+    from_shard_layouts: Sequence[sparse_core_layout_pb2.SparseCoreTableLayout],
+) -> Sequence[tensor.Tensor]:
+  """Undo the shard the feature tables into SparseCore stacked table.
+
+  Args:
+    stacked_table: The value of a SparseCore stacked and sharded table.
+    from_shard_layouts: The target layouts for the target hardware.
+
+  Returns:
+    The unsharded feature tables.
+  """
+  logging.vlog(
+      1,
+      "To unshuffle_from_sc_to_cpu on stacked_table.shape: %s",
+      stacked_table.shape,
+  )
+  ret_tensors = []
+
+  for layout in from_shard_layouts:
+    padded_table = tpu_embedding_v3_utils.unshuffle_from_sc_to_cpu(
+        stacked_table,
+        num_sparse_cores=layout.num_sparse_cores,
+        offset_in_shard=layout.sparse_core_shard_row_offset,
+        size_in_shard=layout.unsharded_padded_shape[0]
+        // layout.num_sparse_cores,
+        shard_rotation=layout.sparse_core_shard_rotation,
+    )
+
+    orig_table = tpu_embedding_v3_utils.remove_padding_from_sc(
+        padded_table, layout.unsharded_shape
+    )
+
+    logging.vlog(
+        1, "orig_tensors.shape[%s]: %s", layout.table_name, orig_table.shape
+    )
+    ret_tensors.append(orig_table)
+
+  return ret_tensors
+
+
+class EmbeddingUnshardToShardCallback(checkpoint_adapter.ReshardCallback):
+  """Reshard callback for embeddings."""
+
+  def __init__(
+      self,
+      object_local_name: str,
+      checkpoint_local_names: Sequence[str],
+      to_shard_layout: Optional[
+          Sequence[sparse_core_layout_pb2.SparseCoreTableLayout]
+      ] = None,
+      to_unshard_layout: Optional[
+          Sequence[sparse_core_layout_pb2.SparseCoreTableLayout]
+      ] = None,
+  ):
+    """Initializes  Reshard callback.
+
+    Args:
+      object_local_name:  The local name of the object being restored.
+      checkpoint_local_names: The local names of the checkpoint positions that
+        need to be read.
+      to_shard_layout: (Optional) Target layouts as specified in the embedding
+        being restored.
+      to_unshard_layout: (Optional) Layouts as stored in checkpoint being
+        restored from.
+    """
+    self._object_local_name = object_local_name
+    self._checkpoint_local_names = checkpoint_local_names
+    self._to_shard_layout = to_shard_layout
+    self._to_unshard_layout = to_unshard_layout
+    self._main_checkpoint_name = checkpoint_local_names[0]
+
+  def object_name(self) -> str:
+    return self._object_local_name
+
+  def update_restore_inputs(
+      self, checkpoint_key: str, shape_and_slice_spec: str
+  ) -> tuple[Sequence[str], Sequence[str]]:
+    """Updates checkpoint key and slice spec acorrding to the resharding plan.
+
+    Args:
+      checkpoint_key: The input checkpoint key to be read.
+      shape_and_slice_spec: The shape and slice spec of the checkpoint key to be
+        read.
+
+    Returns:
+      A tuple of (keys, slices) that should be passed to restore_v2 inorder to
+      reshard according to the resharding plan. The restored tensors from
+      restore_v2 op will usually be passed to reshard method of this class to
+      get the final resharded value.
+    """
+    keys = []
+    slices = []
+    # TODO(b/398016624): Make this a vlog this log after bug is fixed.
+    logging.info(
+        "Updating restore v2 inputs for %s: %s",
+        checkpoint_key,
+        shape_and_slice_spec,
+    )
+    for i, layout in enumerate(self._to_shard_layout):
+      sub_checkpoint_key = checkpoint_key.replace(
+          self._main_checkpoint_name, self._checkpoint_local_names[i]
+      )
+      # For resharding later, we need to read the full value here.
+      # TODO(b/398016624): Make this a vlog this log after bug is fixed.
+      logging.info(
+          "Will read sub key %s: %s",
+          sub_checkpoint_key,
+          layout.unsharded_shape,
+      )
+      keys.append(sub_checkpoint_key)
+      slices.append(
+          _shard_info_str(
+              layout.unsharded_shape,
+              trackable_base.ShardInfo(
+                  offset=[0, 0], shape=layout.unsharded_shape
+              ),
+          )
+      )
+    return (keys, slices)
+
+  def reshard(
+      self, checkpoint_values: tensor.Tensor, shape_and_slice: str
+  ) -> tensor.Tensor:
+    """Reshards the checkpoint values according to the resharding plan.
+
+    Args:
+      checkpoint_values: The checkpoint values to be resharded.
+      shape_and_slice: The shape and slice spec to be returned after resharding.
+
+    Returns:
+      The resharded tensor slice.
+    """
+    return _shard_from_cpu_to_sc(
+        checkpoint_values, shape_and_slice, self._to_shard_layout
+    )
+
+
+class EmbeddingReshardCallback(checkpoint_adapter.ReshardCallback):
+  """Reshard callback for embeddings."""
+
+  def __init__(
+      self,
+      object_local_name: str,
+      from_shard_layouts: Mapping[
+          str, Sequence[sparse_core_layout_pb2.SparseCoreTableLayout]
+      ],
+      to_shard_layouts: Sequence[sparse_core_layout_pb2.SparseCoreTableLayout],
+  ):
+    """Initializes  Reshard callback.
+
+    Args:
+      object_local_name:  The local name of the object being restored.
+      from_shard_layouts: A dictionary in stacked table name to a list of its
+        consituent table layouts.  The layouts are coming from the checkpoint
+        being restored.
+      to_shard_layouts: a list of target layouts that will be resharded to.
+    """
+    logging.info("Creating EmbeddingReshardCallback for %s", object_local_name)
+    self._object_local_name = object_local_name
+    self._from_shard_layouts = from_shard_layouts
+    self._to_shard_layouts = to_shard_layouts
+
+  def object_name(self) -> str:
+    return self._object_local_name
+
+  def update_restore_inputs(
+      self, checkpoint_key: str, shape_and_slice_spec: str
+  ) -> tuple[Sequence[str], Sequence[str]]:
+    """Return the full shape of the stacked that is passed into restore_v2.
+
+    This shape information is required by the restore_v2 process to ensure it
+    loads the complete tensor from the checkpoint. The full tensor is required
+    to perform resharding operations.
+
+    Args:
+      checkpoint_key: The input checkpoint key to be read.
+      shape_and_slice_spec: The shape and slice spec of the checkpoint key to be
+        read.
+
+    Returns:
+      A tuple of (keys, slices) that should be passed to restore_v2 in order to
+      reshard according to the resharding plan. The restored tensors from
+      restore_v2 op will usually be passed to reshard method of this class to
+      get the final resharded value.
+    """
+    keys = []
+    slices = []
+    for stacked_name, table_layouts in self._from_shard_layouts.items():
+      key = checkpoint_key.replace(self._object_local_name, stacked_name)
+      keys.append(key)
+
+      # use the first layout get the full shape of the stacked table
+      first_layout = table_layouts[0]
+      full_vocab_size = (
+          first_layout.total_rows_per_sparse_core_shard
+          * first_layout.num_sparse_cores
+      )
+      stack_dim = first_layout.unsharded_padded_shape[1]
+      full_shape = [full_vocab_size, stack_dim]
+      slices.append(
+          _shard_info_str(
+              full_shape,
+              trackable_base.ShardInfo(offset=[0, 0], shape=full_shape),
+          )
+      )
+
+    logging.info(
+        "Updating restore v2 inputs for %s[%s]:%s to stacked_tables: [%s],"
+        " slices: [%s]",
+        checkpoint_key,
+        self._object_local_name,
+        shape_and_slice_spec,
+        ", ".join(keys),
+        ", ".join(slices),
+    )
+
+    return (keys, slices)
+
+  def reshard(
+      self,
+      checkpoint_values: Sequence[tensor.Tensor],
+      shape_and_slice: str,
+  ) -> tensor.Tensor:
+    # unshard
+    stime = time.time()
+    logging.info(
+        "EmbeddingReshardCallback: starting to reshard [%s],"
+        " from checkpoint_value with shapes: %s",
+        self._object_local_name,
+        ", ".join([str(t.shape) for t in checkpoint_values]),
+    )
+
+    unsharded_tables = dict()
+
+    for stacked_table, layouts in zip(
+        checkpoint_values,
+        list(self._from_shard_layouts.values()),
+    ):
+      logging.info(
+          "Unshard sc_to_cpu stacked_table: %s, shape: %s, no. of constituent"
+          " tables: %d",
+          layouts[0].stacked_table_name,
+          stacked_table.shape,
+          len(layouts),
+      )
+
+      unsharded_tensors = _unshard_from_sc_to_cpu(stacked_table, layouts)
+      for unshared_tensor, layout in zip(unsharded_tensors, layouts):
+        unsharded_tables[layout.table_name] = unshared_tensor
+
+    required_tables = [
+        unsharded_tables[layout.table_name] for layout in self._to_shard_layouts
+    ]
+    ret = _shard_from_cpu_to_sc(
+        required_tables, shape_and_slice, self._to_shard_layouts
+    )
+
+    etime = time.time()
+    logging.info(
+        "EmbeddingReshardCallback: reshard [%s] took %s",
+        self._object_local_name,
+        etime - stime,
+    )
+    return ret
+
+
+def _reorg_layouts(
+    layouts: Sequence[sparse_core_layout_pb2.SparseCoreTableLayout],
+) -> Mapping[str, Sequence[sparse_core_layout_pb2.SparseCoreTableLayout]]:
+  """Reorg the layouts to be in the order of the logical table.
+
+    Building a Dict[StackedTableName, SortedList[TableLayout]]
+
+  Args:
+    layouts: The layouts to be reorged.
+
+  Returns:
+    A dict of stacked table name to sorted list of table layouts.
+  """
+  stacked_name_to_table_names = collections.defaultdict(list)
+  for layout in layouts:
+    stacked_name_to_table_names[layout.stacked_table_name].append(layout)
+  for stacked_name in stacked_name_to_table_names.keys():
+    sorted_layouts = sorted(
+        stacked_name_to_table_names[stacked_name],
+        key=lambda layout: layout.sparse_core_shard_row_offset,
+    )
+    stacked_name_to_table_names[stacked_name] = sorted_layouts
+
+  return stacked_name_to_table_names
+
+
+class TpuEmbeddingV3CheckpointAdapter(
+    checkpoint_adapter.AbstractCheckpointAdapter
+):
+  """Adapter for TPU Embedding V3 to handle checkpoint resharding."""
+
+  def __init__(
+      self,
+      layouts: Optional[sparse_core_layout_pb2.SparseCoreTableLayouts] = None,
+  ):
+    """An adapter for TPUEmbeddingV3 checkpoints.
+
+    Constructs an adapter for TPUEmbeddingV3 to handle layout changes. between
+    checkpoint values and embedding object being restored.
+
+    Args:
+     layouts: The target layouts required.
+    """
+    self._checkpoint_layouts = {}
+    self._checkpoint_to_reshard_callback = {}
+    if layouts:
+      for layout in layouts.tables:
+        self._checkpoint_layouts[layout.table_name] = layout
+
+  @classmethod
+  def create_from_checkpoint(cls, save_path: str):
+    reader = py_checkpoint_reader.NewCheckpointReader(save_path)
+    sparsecore_layouts_str = None
+    for name in reader.get_variable_to_dtype_map():
+      if tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY in name:
+        sparsecore_layouts_str = reader.get_tensor(name)
+        break
+    if sparsecore_layouts_str is None:
+      return cls(None)
+    layouts = sparse_core_layout_pb2.SparseCoreTableLayouts()
+    layouts.ParseFromString(sparsecore_layouts_str)
+    logging.info("Loaded layouts from checkpoint: %s", layouts)
+    return cls(layouts)
+
+  def initialize_reshard_callbacks(
+      self,
+      embedding_layouts: Optional[
+          Mapping[str, sparse_core_layout_pb2.SparseCoreTableLayout]
+      ] = None,
+  ):
+    if not self._checkpoint_layouts and embedding_layouts:
+      # From Unsharded to Sharded
+      stacked_name_to_table_names = collections.defaultdict(list)
+      for layout in embedding_layouts.values():
+        stacked_name_to_table_names[layout.stacked_table_name].append(layout)
+      for stacked_name, layouts in stacked_name_to_table_names.items():
+        # Make the first table name as the key for checkpoint position
+        # The sorting here is by the position of the logical table in the shard
+        sorted_layouts = sorted(
+            layouts, key=lambda layout: layout.sparse_core_shard_row_offset
+        )
+        logging.info("Creating resharding plan for %s", stacked_name)
+        self._checkpoint_to_reshard_callback[sorted_layouts[0].table_name] = (
+            EmbeddingUnshardToShardCallback(
+                stacked_name,
+                [l.table_name for l in sorted_layouts],
+                sorted_layouts,
+                None,
+            )
+        )
+      return
+    if not embedding_layouts:
+      # TODO(b/326644306): From sharded to unsharded
+      raise NotImplementedError("Sharded to Unsharded is not implemented yet.")
+    # Reshard to different SC Layout
+    from_layouts = _reorg_layouts(list(self._checkpoint_layouts.values()))
+    to_layouts = _reorg_layouts(list(embedding_layouts.values()))
+    for stacked_name, table_layouts in to_layouts.items():
+      # look for required stacked tables
+      required_stacked_tables = dict()
+      for table_layout in table_layouts:
+        for from_stacked_name, from_table_layouts in from_layouts.items():
+          if table_layout.table_name in {
+              layout.table_name for layout in from_table_layouts
+          }:
+            required_stacked_tables[from_stacked_name] = from_table_layouts
+
+      logging.info(
+          "Creating resharding plan for %s, required stacked_tables: %s",
+          stacked_name,
+          ", ".join(required_stacked_tables.keys()),
+      )
+      self._checkpoint_to_reshard_callback[stacked_name] = (
+          EmbeddingReshardCallback(
+              object_local_name=stacked_name,
+              from_shard_layouts=required_stacked_tables,
+              to_shard_layouts=to_layouts[stacked_name],
+          )
+      )
+
+  def is_layouts_same(self, embedding_layouts) -> bool:
+    """Returns True if the all the embedding and checkpoint layouts are the same.
+
+    Args:
+      embedding_layouts: dict of layouts for embedding tables.
+
+    Raises: ValueError if the embedding layouts and checkpoint layouts do not
+      have the same keys.
+    Returns: Bool representing if the embedding layouts match the layouts in
+      checkpoint.
+    """
+    if self._checkpoint_layouts.keys() != embedding_layouts.keys():
+      raise ValueError(
+          "Layouts in checkpoint and embedding must have the same keys. found"
+          " {} and {}".format(
+              self._checkpoint_layouts.keys(), embedding_layouts.keys()
+          )
+      )
+
+    for key, layout in self._checkpoint_layouts.items():
+      if not compare.ProtoEq(layout, embedding_layouts[key]):
+        logging.info(
+            "Layouts do not match for %s this will require resharding; %s"
+            " vs %s",
+            key,
+            layout,
+            embedding_layouts[key],
+        )
+        return False
+    return True
+
+  def is_applicable(self, trackable: trackable_base.Trackable) -> bool:
+    # issubclass(trackable, TPUEmbeddingBase) adds circular deps, hence using
+    #  a workaround to select the applicable embedding implementations.
+    allowed_class_names = [".TPUEmbeddingV2Plus", ".TPUEmbeddingV2"]
+    if not any(x in str(type(trackable)) for x in allowed_class_names):
+      return False
+    embedding_layouts = None
+    if hasattr(trackable, "embedding_layouts"):
+      embedding_layouts = trackable.embedding_layouts
+    # Neither checkpoint not target embedding has layout, no resharding needed.
+    if not self._checkpoint_layouts and not embedding_layouts:
+      logging.info("No resharding needed, no layouts")
+      return False
+    # Only if both checkpoint and embedding have layouts and they match,
+    # no resharding needed.
+    if (
+        self._checkpoint_layouts
+        and embedding_layouts
+        and self.is_layouts_same(embedding_layouts)
+    ):
+      logging.info("No resharding needed; layouts match")
+      return False
+    # Else we need to reshard.
+    self.initialize_reshard_callbacks(embedding_layouts)
+    return True
+
+  def get_reshard_callback(
+      self, name: str
+  ) -> Optional[checkpoint_adapter.ReshardCallback]:
+    if name in self._checkpoint_to_reshard_callback:
+      return self._checkpoint_to_reshard_callback[name]
+    # Check if this is slot variable
+    var_name = name.split("/")[0]
+    if var_name in self._checkpoint_to_reshard_callback:
+      return self._checkpoint_to_reshard_callback[var_name]
+    return None

tensorflow/python/tpu/tpu_embedding_v3_checkpoint_adapter_test.py

@@ -0,0 +1,675 @@
+# Copyright 2024 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_embedding_v3_checkpoint_adapter."""
+
+from tensorflow.core.tpu.kernels import sparse_core_layout_pb2
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v3_checkpoint_adapter
+
+tf_constant = constant_op.constant
+
+
+def create_layout(
+    tables_name: str,
+    stacked_table_name: str,
+    num_sparse_cores: int,
+    num_partitions: int,
+    unsharded_shape: tuple[int, int],
+    unsharded_padded_shape: tuple[int, int],
+    row_offset: int,
+    shard_rotation: int,
+    total_rows_per_sparse_core_shard=None,
+) -> sparse_core_layout_pb2.SparseCoreTableLayout():
+  layout = sparse_core_layout_pb2.SparseCoreTableLayout()
+  layout.table_name = tables_name
+  layout.stacked_table_name = stacked_table_name
+  layout.num_sparse_cores = num_sparse_cores
+  layout.num_partitions = num_partitions
+  layout.total_rows_per_sparse_core_shard = (
+      (unsharded_padded_shape[0] // num_sparse_cores)
+      if total_rows_per_sparse_core_shard is None
+      else total_rows_per_sparse_core_shard
+  )
+  layout.unsharded_shape.extend(unsharded_shape)
+  layout.unsharded_padded_shape.extend(unsharded_padded_shape)
+  layout.sparse_core_shard_row_offset = row_offset
+  layout.sparse_core_shard_rotation = shard_rotation
+  return layout
+
+
+class TpuEmbeddingV3CheckpointAdapterTest(test.TestCase):
+
+  def test_adapt_unsharded_to_sharded_simple(self):
+    adapter = (
+        tpu_embedding_v3_checkpoint_adapter.TpuEmbeddingV3CheckpointAdapter(
+            None
+        )
+    )
+    layout = create_layout(
+        tables_name="some_feature",
+        stacked_table_name="some_feature",
+        num_sparse_cores=8,
+        num_partitions=2,
+        unsharded_shape=(20, 4),
+        unsharded_padded_shape=(24, 8),
+        row_offset=0,
+        shard_rotation=8,
+    )
+    t = math_ops.range(start=0.0, limit=20.0, delta=1)[
+        :, None
+    ] * array_ops.ones((20, 4))
+    adapter.initialize_reshard_callbacks({"some_feature": layout})
+    callback = adapter.get_reshard_callback("some_feature")
+    # Check partition index 1 (second parition)
+    self.assertAllEqual(
+        callback.reshard([t], "128 8 8,12:0,8"),
+        tf_constant([
+            [2, 2, 2, 2, 0, 0, 0, 0],
+            [10, 10, 10, 10, 0, 0, 0, 0],
+            [18, 18, 18, 18, 0, 0, 0, 0],
+            [3, 3, 3, 3, 0, 0, 0, 0],
+            [11, 11, 11, 11, 0, 0, 0, 0],
+            [19, 19, 19, 19, 0, 0, 0, 0],
+            [4, 4, 4, 4, 0, 0, 0, 0],
+            [12, 12, 12, 12, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0, 0, 0, 0],
+            [5, 5, 5, 5, 0, 0, 0, 0],
+            [13, 13, 13, 13, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0, 0, 0, 0],
+        ]),
+    )
+
+  def test_adapt_unsharded_to_sharded_stacked(self):
+    adapter = (
+        tpu_embedding_v3_checkpoint_adapter.TpuEmbeddingV3CheckpointAdapter(
+            None
+        )
+    )
+    layouts = {
+        "two": create_layout(
+            tables_name="two",
+            stacked_table_name="one_two",
+            num_sparse_cores=8,
+            num_partitions=4,
+            unsharded_shape=(32, 4),
+            unsharded_padded_shape=(32, 8),
+            row_offset=3,
+            shard_rotation=1,
+            total_rows_per_sparse_core_shard=7,
+        ),
+        "one": create_layout(
+            tables_name="one",
+            stacked_table_name="one_two",
+            num_sparse_cores=8,
+            num_partitions=4,
+            unsharded_shape=(20, 4),
+            unsharded_padded_shape=(24, 8),
+            row_offset=0,
+            shard_rotation=0,
+            total_rows_per_sparse_core_shard=7,
+        ),
+    }
+    one_t = math_ops.range(start=0.0, limit=20.0, delta=1)[
+        :, None
+    ] * array_ops.ones((20, 4))
+    two_t = math_ops.range(start=50.0, limit=82.0, delta=1)[
+        :, None
+    ] * array_ops.ones((32, 4))
+    adapter.initialize_reshard_callbacks(layouts)
+    callback = adapter.get_reshard_callback("one")
+    self.assertEqual(callback.object_name(), "one_two")
+    updated_keys, updated_slices = callback.update_restore_inputs(
+        "path/to/embedding/one/in/checkpoint", "56 8 14,28:0,8"
+    )
+    self.assertAllEqual(
+        updated_keys,
+        [
+            "path/to/embedding/one/in/checkpoint",
+            "path/to/embedding/two/in/checkpoint",
+        ],
+    )
+    self.assertAllEqual(
+        updated_slices,
+        ["20 4 0,20:0,4", "32 4 0,32:0,4"],
+    )
+    actual = callback.reshard([one_t, two_t], "56 8 14,14:0,8")
+    self.assertAllEqual(
+        actual,
+        tf_constant([
+            # table one shard 2
+            [2, 2, 2, 2, 0, 0, 0, 0],
+            [10, 10, 10, 10, 0, 0, 0, 0],
+            [18, 18, 18, 18, 0, 0, 0, 0],
+            # table two shard 2
+            [51, 51, 51, 51, 0, 0, 0, 0],
+            [59, 59, 59, 59, 0, 0, 0, 0],
+            [67, 67, 67, 67, 0, 0, 0, 0],
+            [75, 75, 75, 75, 0, 0, 0, 0],
+            # table one shard 3
+            [3, 3, 3, 3, 0, 0, 0, 0],
+            [11, 11, 11, 11, 0, 0, 0, 0],
+            [19, 19, 19, 19, 0, 0, 0, 0],
+            # table two shard 3
+            [52, 52, 52, 52, 0, 0, 0, 0],
+            [60, 60, 60, 60, 0, 0, 0, 0],
+            [68, 68, 68, 68, 0, 0, 0, 0],
+            [76, 76, 76, 76, 0, 0, 0, 0],
+        ]),
+    )
+    # Check that full resharding works.
+    actual_full = callback.reshard([one_t, two_t], "56 8 0,56:0,8")
+    self.assertAllEqual(
+        actual_full,
+        tf_constant(
+            [
+                # table one shard 0
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                [8, 8, 8, 8, 0, 0, 0, 0],
+                [16, 16, 16, 16, 0, 0, 0, 0],
+                # table two shard 0
+                [57, 57, 57, 57, 0, 0, 0, 0],
+                [65, 65, 65, 65, 0, 0, 0, 0],
+                [73, 73, 73, 73, 0, 0, 0, 0],
+                [81, 81, 81, 81, 0, 0, 0, 0],
+                # table one shard 1
+                [1, 1, 1, 1, 0, 0, 0, 0],
+                [9, 9, 9, 9, 0, 0, 0, 0],
+                [17, 17, 17, 17, 0, 0, 0, 0],
+                # table two shard 1
+                [50, 50, 50, 50, 0, 0, 0, 0],
+                [58, 58, 58, 58, 0, 0, 0, 0],
+                [66, 66, 66, 66, 0, 0, 0, 0],
+                [74, 74, 74, 74, 0, 0, 0, 0],
+                # table one shard 2
+                [2, 2, 2, 2, 0, 0, 0, 0],
+                [10, 10, 10, 10, 0, 0, 0, 0],
+                [18, 18, 18, 18, 0, 0, 0, 0],
+                # table two shard 2
+                [51, 51, 51, 51, 0, 0, 0, 0],
+                [59, 59, 59, 59, 0, 0, 0, 0],
+                [67, 67, 67, 67, 0, 0, 0, 0],
+                [75, 75, 75, 75, 0, 0, 0, 0],
+                # table one shard 3
+                [3, 3, 3, 3, 0, 0, 0, 0],
+                [11, 11, 11, 11, 0, 0, 0, 0],
+                [19, 19, 19, 19, 0, 0, 0, 0],
+                # table two shard 3
+                [52, 52, 52, 52, 0, 0, 0, 0],
+                [60, 60, 60, 60, 0, 0, 0, 0],
+                [68, 68, 68, 68, 0, 0, 0, 0],
+                [76, 76, 76, 76, 0, 0, 0, 0],
+                # table one shard 4
+                [4, 4, 4, 4, 0, 0, 0, 0],
+                [12, 12, 12, 12, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                # table two shard 4
+                [53, 53, 53, 53, 0, 0, 0, 0],
+                [61, 61, 61, 61, 0, 0, 0, 0],
+                [69, 69, 69, 69, 0, 0, 0, 0],
+                [77, 77, 77, 77, 0, 0, 0, 0],
+                # table one shard 5
+                [5, 5, 5, 5, 0, 0, 0, 0],
+                [13, 13, 13, 13, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                # table two shard 5
+                [54, 54, 54, 54, 0, 0, 0, 0],
+                [62, 62, 62, 62, 0, 0, 0, 0],
+                [70, 70, 70, 70, 0, 0, 0, 0],
+                [78, 78, 78, 78, 0, 0, 0, 0],
+                # table one shard 6
+                [6, 6, 6, 6, 0, 0, 0, 0],
+                [14, 14, 14, 14, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                # table two shard 6
+                [55, 55, 55, 55, 0, 0, 0, 0],
+                [63, 63, 63, 63, 0, 0, 0, 0],
+                [71, 71, 71, 71, 0, 0, 0, 0],
+                [79, 79, 79, 79, 0, 0, 0, 0],
+                # table one shard 7
+                [7, 7, 7, 7, 0, 0, 0, 0],
+                [15, 15, 15, 15, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0, 0, 0, 0],
+                # table two shard 7
+                [56, 56, 56, 56, 0, 0, 0, 0],
+                [64, 64, 64, 64, 0, 0, 0, 0],
+                [72, 72, 72, 72, 0, 0, 0, 0],
+                [80, 80, 80, 80, 0, 0, 0, 0],
+            ],
+            dtype=dtypes.float32,
+        ),
+    )
+    self.assertAllEqual(callback._checkpoint_local_names, ["one", "two"])
+    self.assertAllEqual(
+        [l.table_name for l in callback._to_shard_layout],
+        ["one", "two"],
+    )
+
+  def test_adapt_sharded_to_unsharded_simple(self):
+    pass
+
+  def test_adapt_sharded_to_unsharded_stacked(self):
+    pass
+
+  def test_is_layouts_same_works(self):
+    layout = create_layout(
+        tables_name="some_feature",
+        stacked_table_name="some_feature",
+        num_sparse_cores=8,
+        num_partitions=8,
+        unsharded_shape=(100, 4),
+        unsharded_padded_shape=(128, 8),
+        row_offset=0,
+        shard_rotation=0,
+    )
+    layouts = sparse_core_layout_pb2.SparseCoreTableLayouts()
+    layouts.tables.append(layout)
+    adapter = (
+        tpu_embedding_v3_checkpoint_adapter.TpuEmbeddingV3CheckpointAdapter(
+            layouts
+        )
+    )
+    self.assertTrue(adapter.is_layouts_same({layout.table_name: layout}))
+    layout.num_sparse_cores = 3
+    self.assertFalse(adapter.is_layouts_same({layout.table_name: layout}))
+
+  def test_adapt_to_different_sharded_stacked(self):
+
+    source_layouts = {
+        "one": create_layout(
+            tables_name="one",
+            stacked_table_name="one_two_three",
+            num_sparse_cores=4,
+            num_partitions=2,
+            unsharded_shape=(6, 5),
+            unsharded_padded_shape=(8, 8),
+            row_offset=0,
+            shard_rotation=0,
+            total_rows_per_sparse_core_shard=6,
+        ),
+        "two": create_layout(
+            tables_name="two",
+            stacked_table_name="one_two_three",
+            num_sparse_cores=4,
+            num_partitions=2,
+            unsharded_shape=(7, 4),
+            unsharded_padded_shape=(8, 8),
+            row_offset=2,
+            shard_rotation=1,
+            total_rows_per_sparse_core_shard=6,
+        ),
+        "three": create_layout(
+            tables_name="three",
+            stacked_table_name="one_two_three",
+            num_sparse_cores=4,
+            num_partitions=2,
+            unsharded_shape=(15, 3),
+            unsharded_padded_shape=(16, 8),
+            row_offset=4,
+            shard_rotation=2,
+            total_rows_per_sparse_core_shard=6,
+        ),
+    }
+    src_layouts_pb = sparse_core_layout_pb2.SparseCoreTableLayouts()
+    src_layouts_pb.tables.extend(source_layouts.values())
+
+    sc_to_sc_adapter = (
+        tpu_embedding_v3_checkpoint_adapter.TpuEmbeddingV3CheckpointAdapter(
+            layouts=src_layouts_pb
+        )
+    )
+
+    target_layouts = {
+        "one": create_layout(
+            tables_name="one",
+            stacked_table_name="one_two_three",
+            num_sparse_cores=8,
+            num_partitions=4,
+            unsharded_shape=(6, 5),
+            unsharded_padded_shape=(8, 8),
+            row_offset=0,
+            shard_rotation=0,
+            total_rows_per_sparse_core_shard=4,
+        ),
+        "two": create_layout(
+            tables_name="two",
+            stacked_table_name="one_two_three",
+            num_sparse_cores=8,
+            num_partitions=4,
+            unsharded_shape=(7, 4),
+            unsharded_padded_shape=(8, 8),
+            row_offset=1,
+            shard_rotation=1,
+            total_rows_per_sparse_core_shard=4,
+        ),
+        "three": create_layout(
+            tables_name="three",
+            stacked_table_name="one_two_three",
+            num_sparse_cores=8,
+            num_partitions=4,
+            unsharded_shape=(15, 3),
+            unsharded_padded_shape=(16, 8),
+            row_offset=2,
+            shard_rotation=2,
+            total_rows_per_sparse_core_shard=4,
+        ),
+    }
+
+    # this take a mapping[str, sparse_core_layout_pb2.SparseCoreTableLayout]
+    sc_to_sc_adapter.initialize_reshard_callbacks(target_layouts)
+    callback = sc_to_sc_adapter.get_reshard_callback("one_two_three")
+    self.assertEqual(callback.object_name(), "one_two_three")
+    updated_keys, updated_slices = callback.update_restore_inputs(
+        "path/to/embedding/one_two_three/in/checkpoint", "24 8 6,12:0,8"
+    )
+    self.assertAllEqual(
+        updated_keys,
+        [
+            "path/to/embedding/one_two_three/in/checkpoint",
+        ],
+    )
+    self.assertAllEqual(
+        updated_slices,
+        ["24 8 0,24:0,8"],
+    )
+
+    one_two_three = tf_constant([
+        # table one shard 0
+        [0, 0, 0, 0, 0, 0, 0, 0],
+        [4, 4, 4, 4, 4, 0, 0, 0],
+        [13, 13, 13, 13, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0],
+        [102, 102, 102, 0, 0, 0, 0, 0],
+        [106, 106, 106, 0, 0, 0, 0, 0],
+        [110, 110, 110, 0, 0, 0, 0, 0],
+        [114, 114, 114, 0, 0, 0, 0, 0],
+        # table one shard 1
+        [1, 1, 1, 1, 1, 0, 0, 0],
+        [5, 5, 5, 5, 5, 0, 0, 0],
+        [10, 10, 10, 10, 0, 0, 0, 0],
+        [14, 14, 14, 14, 0, 0, 0, 0],
+        [103, 103, 103, 0, 0, 0, 0, 0],
+        [107, 107, 107, 0, 0, 0, 0, 0],
+        [111, 111, 111, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0],
+        # table one shard 2
+        [2, 2, 2, 2, 2, 0, 0, 0],
+        [6, 6, 6, 6, 6, 0, 0, 0],
+        [11, 11, 11, 11, 0, 0, 0, 0],
+        [15, 15, 15, 15, 0, 0, 0, 0],
+        [100, 100, 100, 0, 0, 0, 0, 0],
+        [104, 104, 104, 0, 0, 0, 0, 0],
+        [108, 108, 108, 0, 0, 0, 0, 0],
+        [112, 112, 112, 0, 0, 0, 0, 0],
+        # table one shard 3
+        [3, 3, 3, 3, 3, 0, 0, 0],
+        [7, 7, 7, 7, 7, 0, 0, 0],
+        [12, 12, 12, 12, 0, 0, 0, 0],
+        [16, 16, 16, 16, 0, 0, 0, 0],
+        [101, 101, 101, 0, 0, 0, 0, 0],
+        [105, 105, 105, 0, 0, 0, 0, 0],
+        [109, 109, 109, 0, 0, 0, 0, 0],
+        [113, 113, 113, 0, 0, 0, 0, 0],
+    ])
+
+    self.assertAllEqual(
+        tf_constant([
+            #  shard 2
+            [2, 2, 2, 2, 2, 0, 0, 0],
+            [11, 11, 11, 11, 0, 0, 0, 0],
+            [100, 100, 100, 0, 0, 0, 0, 0],
+            [108, 108, 108, 0, 0, 0, 0, 0],
+            # shard 3
+            [3, 3, 3, 3, 3, 0, 0, 0],
+            [12, 12, 12, 12, 0, 0, 0, 0],
+            [101, 101, 101, 0, 0, 0, 0, 0],
+            [109, 109, 109, 0, 0, 0, 0, 0],
+            # shard 4
+            [4, 4, 4, 4, 4, 0, 0, 0],
+            [13, 13, 13, 13, 0, 0, 0, 0],
+            [102, 102, 102, 0, 0, 0, 0, 0],
+            [110, 110, 110, 0, 0, 0, 0, 0],
+            # shard 5
+            [5, 5, 5, 5, 5, 0, 0, 0],
+            [14, 14, 14, 14, 0, 0, 0, 0],
+            [103, 103, 103, 0, 0, 0, 0, 0],
+            [111, 111, 111, 0, 0, 0, 0, 0],
+        ]),
+        callback.reshard([one_two_three], "32 8 8,16:0,8"),
+    )
+
+  def test_adapt_to_different_sc_table_stacking(self):
+
+    source_layouts = {
+        "one": create_layout(
+            tables_name="one",
+            stacked_table_name="one_two",
+            num_sparse_cores=4,
+            num_partitions=2,
+            unsharded_shape=(6, 5),
+            unsharded_padded_shape=(8, 8),
+            row_offset=0,
+            shard_rotation=0,
+            total_rows_per_sparse_core_shard=4,
+        ),
+        "two": create_layout(
+            tables_name="two",
+            stacked_table_name="one_two",
+            num_sparse_cores=4,
+            num_partitions=2,
+            unsharded_shape=(7, 4),
+            unsharded_padded_shape=(8, 8),
+            row_offset=2,
+            shard_rotation=1,
+            total_rows_per_sparse_core_shard=4,
+        ),
+        "three": create_layout(
+            tables_name="three",
+            stacked_table_name="three",
+            num_sparse_cores=4,
+            num_partitions=2,
+            unsharded_shape=(15, 3),
+            unsharded_padded_shape=(16, 8),
+            row_offset=0,
+            shard_rotation=0,
+            total_rows_per_sparse_core_shard=4,
+        ),
+    }
+    src_layouts_pb = sparse_core_layout_pb2.SparseCoreTableLayouts()
+    src_layouts_pb.tables.extend(source_layouts.values())
+
+    sc_to_sc_adapter = (
+        tpu_embedding_v3_checkpoint_adapter.TpuEmbeddingV3CheckpointAdapter(
+            layouts=src_layouts_pb
+        )
+    )
+
+    target_layouts = {
+        "one": create_layout(
+            tables_name="one",
+            stacked_table_name="one",
+            num_sparse_cores=8,
+            num_partitions=4,
+            unsharded_shape=(6, 5),
+            unsharded_padded_shape=(8, 8),
+            row_offset=0,
+            shard_rotation=0,
+            total_rows_per_sparse_core_shard=1,
+        ),
+        "two": create_layout(
+            tables_name="two",
+            stacked_table_name="two_three",
+            num_sparse_cores=8,
+            num_partitions=4,
+            unsharded_shape=(7, 4),
+            unsharded_padded_shape=(8, 8),
+            row_offset=0,
+            shard_rotation=0,
+            total_rows_per_sparse_core_shard=3,
+        ),
+        "three": create_layout(
+            tables_name="three",
+            stacked_table_name="two_three",
+            num_sparse_cores=8,
+            num_partitions=4,
+            unsharded_shape=(15, 3),
+            unsharded_padded_shape=(16, 8),
+            row_offset=1,
+            shard_rotation=1,
+            total_rows_per_sparse_core_shard=3,
+        ),
+    }
+
+    # this take a mapping[str, sparse_core_layout_pb2.SparseCoreTableLayout]
+    sc_to_sc_adapter.initialize_reshard_callbacks(target_layouts)
+
+    src_one_two = tf_constant([
+        # shard 0
+        [0, 0, 0, 0, 0, 0, 0, 0],
+        [4, 4, 4, 4, 4, 0, 0, 0],
+        [13, 13, 13, 13, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0],
+        # shard 1
+        [1, 1, 1, 1, 1, 0, 0, 0],
+        [5, 5, 5, 5, 5, 0, 0, 0],
+        [10, 10, 10, 10, 0, 0, 0, 0],
+        [14, 14, 14, 14, 0, 0, 0, 0],
+        # shard 2
+        [2, 2, 2, 2, 2, 0, 0, 0],
+        [6, 6, 6, 6, 6, 0, 0, 0],
+        [11, 11, 11, 11, 0, 0, 0, 0],
+        [15, 15, 15, 15, 0, 0, 0, 0],
+        # shard 3
+        [3, 3, 3, 3, 3, 0, 0, 0],
+        [7, 7, 7, 7, 7, 0, 0, 0],
+        [12, 12, 12, 12, 0, 0, 0, 0],
+        [16, 16, 16, 16, 0, 0, 0, 0],
+    ])
+
+    src_three = tf_constant([
+        # shard 0
+        [100, 100, 100, 0, 0, 0, 0, 0],
+        [104, 104, 104, 0, 0, 0, 0, 0],
+        [108, 108, 108, 0, 0, 0, 0, 0],
+        [112, 112, 112, 0, 0, 0, 0, 0],
+        # shard 1
+        [101, 101, 101, 0, 0, 0, 0, 0],
+        [105, 105, 105, 0, 0, 0, 0, 0],
+        [109, 109, 109, 0, 0, 0, 0, 0],
+        [113, 113, 113, 0, 0, 0, 0, 0],
+        # shard 2
+        [102, 102, 102, 0, 0, 0, 0, 0],
+        [106, 106, 106, 0, 0, 0, 0, 0],
+        [110, 110, 110, 0, 0, 0, 0, 0],
+        [114, 114, 114, 0, 0, 0, 0, 0],
+        # shard 3
+        [103, 103, 103, 0, 0, 0, 0, 0],
+        [107, 107, 107, 0, 0, 0, 0, 0],
+        [111, 111, 111, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0, 0, 0],
+    ])
+
+    with self.subTest("one"):
+      callback = sc_to_sc_adapter.get_reshard_callback("one")
+      self.assertEqual(callback.object_name(), "one")
+      updated_keys, updated_slices = callback.update_restore_inputs(
+          "path/to/embedding/one/in/checkpoint", "8 8 2,3:0,8"
+      )
+
+      self.assertAllEqual(
+          updated_keys,
+          [
+              "path/to/embedding/one_two/in/checkpoint",
+          ],
+      )
+      self.assertAllEqual(
+          updated_slices,
+          ["16 8 0,16:0,8"],
+      )
+
+      self.assertAllEqual(
+          tf_constant([
+              [2, 2, 2, 2, 2, 0, 0, 0],
+              [3, 3, 3, 3, 3, 0, 0, 0],
+              [4, 4, 4, 4, 4, 0, 0, 0],
+              [5, 5, 5, 5, 5, 0, 0, 0],
+          ]),
+          callback.reshard([src_one_two], "8 8 2,4:0,8"),
+      )
+
+    self.assertAllEqual(
+        tf_constant([
+            [4, 4, 4, 4, 4, 0, 0, 0],
+            [5, 5, 5, 5, 5, 0, 0, 0],
+            [0, 0, 0, 0, 0, 0, 0, 0],
+            [0, 0, 0, 0, 0, 0, 0, 0],
+        ]),
+        callback.reshard([src_one_two], "8 8 4,4:0,8"),
+    )
+
+    with self.subTest("two_three"):
+      callback = sc_to_sc_adapter.get_reshard_callback("two_three")
+      self.assertEqual(callback.object_name(), "two_three")
+      updated_keys, updated_slices = callback.update_restore_inputs(
+          "path/to/embedding/two_three/in/checkpoint", "24 8 8,6:0,8"
+      )
+      self.assertAllEqual(
+          updated_keys,
+          [
+              "path/to/embedding/one_two/in/checkpoint",
+              "path/to/embedding/three/in/checkpoint",
+          ],
+      )
+      self.assertAllEqual(
+          updated_slices,
+          ["16 8 0,16:0,8", "16 8 0,16:0,8"],
+      )
+
+    self.assertAllEqual(
+        tf_constant([
+            #  shard 2
+            [12, 12, 12, 12, 0, 0, 0, 0],
+            [101, 101, 101, 0, 0, 0, 0, 0],
+            [109, 109, 109, 0, 0, 0, 0, 0],
+            #  shard 3
+            [13, 13, 13, 13, 0, 0, 0, 0],
+            [102, 102, 102, 0, 0, 0, 0, 0],
+            [110, 110, 110, 0, 0, 0, 0, 0],
+        ]),
+        callback.reshard([src_one_two, src_three], "24 8 6,6:0,8"),
+    )
+
+    self.assertAllEqual(
+        tf_constant([
+            #  shard 6
+            [16, 16, 16, 16, 0, 0, 0, 0],
+            [105, 105, 105, 0, 0, 0, 0, 0],
+            [113, 113, 113, 0, 0, 0, 0, 0],
+            #  shard 7
+            [0, 0, 0, 0, 0, 0, 0, 0],
+            [106, 106, 106, 0, 0, 0, 0, 0],
+            [114, 114, 114, 0, 0, 0, 0, 0],
+        ]),
+        callback.reshard([src_one_two, src_three], "24 8 18,6:0,8"),
+    )
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tpu_embedding_v3_checkpoint_test.py

@@ -0,0 +1,124 @@
+# Copyright 2025 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_embedding_v3_checkpoint."""
+import os
+
+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 import tpu_strategy
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.framework import config
+from tensorflow.python.ops import init_ops_v2
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v2_utils
+from tensorflow.python.tpu import tpu_embedding_v3
+
+
+class TPUEmbeddingV3CheckpointTest(parameterized.TestCase, test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.vocabulary_size = 16384
+    self.embedding_dim = 128
+
+  def test_checkpoint_save_and_restore(self):
+    feature_config_1 = (
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=tpu_embedding_v2_utils.TableConfig(
+                vocabulary_size=self.vocabulary_size,
+                dim=self.embedding_dim,
+                initializer=init_ops_v2.Constant(1.0),
+                optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1),
+                combiner="sum",
+                name="video"),
+            name="watched",
+            output_shape=[16]))
+
+    feature_config_2 = (
+        tpu_embedding_v2_utils.FeatureConfig(
+            table=tpu_embedding_v2_utils.TableConfig(
+                vocabulary_size=self.vocabulary_size,
+                dim=self.embedding_dim,
+                initializer=init_ops_v2.Constant(2.0),  # different initializer
+                optimizer=tpu_embedding_v2_utils.SGD(learning_rate=1),
+                combiner="sum",
+                name="video"),
+            name="watched",
+            output_shape=[16]))
+
+    resolver = tpu_cluster_resolver.TPUClusterResolver(tpu="")
+    tpu_cluster_resolver.initialize_tpu_system(resolver)
+    strategy = tpu_strategy.TPUStrategy(resolver)
+
+    with strategy.scope():
+      model1 = tpu_embedding_v3.TPUEmbeddingV2(
+          feature_config=feature_config_1,
+          optimizer=tpu_embedding_v2_utils.SGD())
+      model1.build()
+
+      # Check saving from inside scope works.
+      checkpoint = util.Checkpoint(model=model1)
+      checkpoint.save(self._get_tmpdir("restore", "save"))
+
+    # Check the variable created by model1
+    expected_shard_shape = (self.vocabulary_size //
+                            strategy.num_replicas_in_sync, self.embedding_dim)
+    self.assertIsInstance(model1._variables["video"]["parameters"],
+                          tpu_embedding_v3.TPUEmbeddingShardedVariable)
+    self.assertLen(model1._variables["video"]["parameters"].values,
+                   strategy.num_replicas_in_sync)
+    self.assertEqual(model1._variables["video"]["parameters"].values[0].shape,
+                     expected_shard_shape)
+    self.assertAllEqual(
+        model1._variables["video"]["parameters"].values[0].numpy(),
+        np.ones(expected_shard_shape) * 1.0)
+
+    with strategy.scope():
+      model2 = tpu_embedding_v3.TPUEmbeddingV2(
+          feature_config=feature_config_2,
+          optimizer=tpu_embedding_v2_utils.SGD())
+
+      def fail_initializer(*args, **kwargs):
+        del args, kwargs
+        self.fail("initializer should not be called when restoring")
+
+      assert model2._batch_initialize_tables
+      model2._batch_initialize_tables = fail_initializer
+
+    checkpoint = util.Checkpoint(model=model2)
+
+    # Load from checkpoint
+    checkpoint.restore(self._get_tmpdir("restore", "save-1"))
+    model2.build()
+
+    # Check the variable restored by model2
+    self.assertAllEqual(
+        model2._variables["video"]["parameters"].values[0].numpy(),
+        np.ones(expected_shard_shape) * 1.0)
+
+  def _get_tmpdir(self, name, subdir=""):
+    segments = [os.environ.get("TEST_TMPDIR", "/tmp"), name] + (
+        [subdir] if subdir else []
+    )
+    return os.path.join(*segments)
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  config.enable_mlir_bridge()
+  test.main()

tensorflow/python/tpu/tpu_embedding_v3_cpu_ops_test.py

@@ -0,0 +1,782 @@
+# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import functools
+import itertools
+import math
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import bitwise_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.platform import test
+from tensorflow.python.tpu.ops import gen_xla_ops as xla_ops
+
+
+def _get_combiner_scale_contribution(x, combiner):
+  if combiner == "sum":
+    return 1.0
+  elif combiner == "mean":
+    return x
+  else:
+    return x * x
+
+
+def _get_combiner_scale_transform(x, combiner):
+  if combiner == "sum":
+    return 1.0
+  elif combiner == "mean":
+    return 0 if x == 0 else (1 / x)
+  else:
+    return 0 if x == 0 else (1 / math_ops.sqrt(x))
+
+
+def convert_input_to_coo_tensor(
+    indices_or_row_splits, values, weight, sample_count, combiner
+):
+  if indices_or_row_splits.shape.rank >= 2:
+    row_ids_before_dedup = indices_or_row_splits[:, 0]
+  elif indices_or_row_splits.shape.rank == 1:
+    row_ids_before_dedup = []
+    current_row_id = -1
+    for i, _ in enumerate(values):
+      while i == indices_or_row_splits[current_row_id + 1]:
+        current_row_id += 1
+      row_ids_before_dedup.append(current_row_id)
+  else:
+    row_ids_before_dedup = math_ops.range(0, values.shape[0])
+
+  col_ids = []
+  row_ids = []
+  gains = []
+  gains_rescale = [0] * sample_count
+  for i in range(values.shape[0]):
+    gains_rescale[row_ids_before_dedup[i]] += _get_combiner_scale_contribution(
+        weight[i], combiner
+    )
+    if (
+        row_ids
+        and row_ids[-1] == row_ids_before_dedup[i]
+        and col_ids[-1] == values[i]
+    ):
+      gains[-1] += weight[i]
+    else:
+      row_ids.append(row_ids_before_dedup[i])
+      col_ids.append(values[i])
+      gains.append(weight[i])
+
+  for i in range(sample_count):
+    gains_rescale[i] = _get_combiner_scale_transform(gains_rescale[i], combiner)
+  for i in range(len(row_ids)):
+    gains[i] *= gains_rescale[row_ids[i]]
+
+  row_ids = ops.convert_to_tensor(row_ids, dtype=dtypes.int32)
+  col_ids = ops.convert_to_tensor(col_ids, dtype=dtypes.int32)
+  gains = ops.convert_to_tensor(gains, dtype=dtypes.float32)
+
+  return row_ids, col_ids, gains
+
+
+def _compute_sparse_core_stats(row_ids, col_ids, num_sc_shards):
+  max_ids = np.zeros(num_sc_shards)
+  max_unique_ids = np.zeros(num_sc_shards)
+  previous_col_id = -1
+  previous_row_id = -1
+  for col_id, row_id in sorted(zip(col_ids, row_ids)):
+    if col_id != previous_col_id:
+      max_ids[col_id % num_sc_shards] += 1
+      max_unique_ids[col_id % num_sc_shards] += 1
+    else:
+      if previous_row_id != row_id:
+        max_ids[col_id % num_sc_shards] += 1
+    previous_col_id = col_id
+    previous_row_id = row_id
+  return max(max_ids), max(max_unique_ids)
+
+
+def _convert_coo_tensor_to_csr_with_physical_replica(
+    row_ids,
+    col_ids,
+    gains,
+    splits,
+    sample_count,
+    num_replica,
+    max_minibatches_per_sc,
+    max_ids_per_chip_per_sample,
+    table_vocab_size,
+):
+  num_sc_per_replica = 4
+  num_physical_replica = num_replica * num_sc_per_replica
+  assert (
+      sample_count % num_sc_per_replica == 0
+  ), f"sample count should be multiply of 4 instead got {sample_count}"
+
+  per_sc_sample_count = sample_count // num_sc_per_replica
+
+  splits = splits.numpy()
+  if splits.size > 1:
+    splits = functools.reduce(lambda x, y: x | y, splits)
+
+  max_division_level = 6
+  max_divisions = 1 << max_division_level
+
+  division_size = (table_vocab_size + max_divisions - 1) // max_divisions
+
+  bucket_splits = []
+  current_index = 0
+  while splits > 0:
+    if splits % 2 == 1:
+      split_level = int(current_index + 1).bit_length() - 1
+      split_offset = current_index + 1 - (1 << split_level)
+      split_size = 1 << (max_division_level - 1 - split_level)
+      bucket_splits.append(split_size + split_offset * split_size * 2)
+    splits >>= 1
+    current_index += 1
+
+  bucket_splits.sort()
+
+  num_minibatch_per_sc = len(bucket_splits) + 1
+
+  embedding_lookup_inputs = []
+  for row_id, col_id, gain in zip(row_ids, col_ids, gains):
+    embedding_lookup_inputs.append(
+        (col_id % num_physical_replica, col_id, row_id, gain)
+    )
+  # sort based on replica id first, then col_id.
+  embedding_lookup_inputs.sort()
+
+  total_minibatches = num_minibatch_per_sc * num_sc_per_replica
+
+  assert num_minibatch_per_sc <= max_minibatches_per_sc, (
+      f"Get {num_minibatch_per_sc} minibatches per sparse core, but the"
+      " number of max minibatches per sparse core is"
+      f" {max_minibatches_per_sc}"
+  )
+
+  minibatches = [[] for _ in range(total_minibatches)]
+
+  def calculate_minibatch_id(col_id):
+    for i, bucket_split in enumerate(bucket_splits):
+      if bucket_split * division_size > col_id:
+        return i
+    return len(bucket_splits)
+
+  for embedding_lookup_input in embedding_lookup_inputs:
+    sc_id = embedding_lookup_input[2] // per_sc_sample_count
+    minibatch_id = calculate_minibatch_id(embedding_lookup_input[1])
+
+    minibatches[sc_id * num_minibatch_per_sc + minibatch_id].append(
+        embedding_lookup_input
+    )
+
+  def round_up_to(x, round_value):
+    return x + -x % round_value
+
+  max_ids_per_chip = max_ids_per_chip_per_sample * sample_count
+
+  padded_row_pointers_size = round_up_to(num_physical_replica, 8)
+
+  total_row_pinters_size = padded_row_pointers_size * (
+      max_minibatches_per_sc * num_sc_per_replica
+  )
+
+  row_pointers = np.full(total_row_pinters_size, 8, dtype=np.int32)
+  sorted_sample_ids = np.full(max_ids_per_chip, 8, dtype=np.int32)
+  sorted_token_ids = np.full(max_ids_per_chip, 8, dtype=np.int32)
+  sorted_gains = np.full(max_ids_per_chip, 8, dtype=np.float32)
+
+  id_index = 0
+  row_pointers_index = 0
+  for minibatch in minibatches:
+    index = 0
+    for replica_id in range(num_physical_replica):
+      while index < len(minibatch) and replica_id == minibatch[index][0]:
+        sorted_token_ids[id_index] = minibatch[index][1] // num_physical_replica
+        sorted_sample_ids[id_index] = minibatch[index][2] % per_sc_sample_count
+        sorted_gains[id_index] = minibatch[index][3]
+        index += 1
+        id_index += 1
+      row_pointers[row_pointers_index] = id_index
+      id_index = round_up_to(id_index, 8)
+      row_pointers_index += 1
+
+    for i in range(
+        row_pointers_index,
+        round_up_to(row_pointers_index, padded_row_pointers_size),
+    ):
+      row_pointers[i] = id_index
+    row_pointers_index = round_up_to(
+        row_pointers_index, padded_row_pointers_size
+    )
+
+  row_pointers_unpadded_size = total_minibatches * padded_row_pointers_size
+  ids_unpadded_size = id_index
+
+  return (
+      row_pointers,
+      sorted_sample_ids,
+      sorted_token_ids,
+      sorted_gains,
+      np.array(row_pointers_unpadded_size, dtype=np.int32),
+      np.array(ids_unpadded_size, dtype=np.int32),
+      np.array(num_minibatch_per_sc, dtype=np.int32),
+  )
+
+
+class TpuEmbeddingV3CPUOpsTest(parameterized.TestCase, test.TestCase):
+
+  @parameterized.parameters(
+      *list(
+          itertools.product(
+              [16, 32],
+              [1024, 2048],
+              ["sum", "mean", "sqrtn"],
+              [0, 56, 1600000],
+              [0, 12, 20],
+          )
+      )
+  )
+  def test_convert_to_list_of_sparse_core_coo_tensor(
+      self, sample_count, token_count, combiner, col_offset, col_shift
+  ):
+    sparse_feature = sparse_ops.sparse_reorder(
+        sparse_tensor.SparseTensor(
+            indices=[
+                [i % sample_count, i]
+                for i in np.random.randint(low=0, high=1024, size=token_count)
+            ],
+            values=np.random.randint(low=0, high=1024, size=token_count),
+            dense_shape=[sample_count, 1024],
+        )
+    )
+
+    num_sc_per_chip = 4
+    num_chip = 128
+
+    row_offset = sample_count
+    num_sc_shards = num_sc_per_chip * num_chip
+    stacked_table_sample_count = sample_count * 4
+
+    num_sc_shards_bit = int(math.log2(num_sc_shards))
+    num_sc_shards_bit_mod = (1 << num_sc_shards_bit) - 1
+    num_sc_shards_bit_mod_inv = bitwise_ops.invert(num_sc_shards_bit_mod)
+
+    row_ids, col_ids, gains = convert_input_to_coo_tensor(
+        indices_or_row_splits=sparse_feature.indices,
+        values=sparse_feature.values,
+        weight=np.ones(shape=token_count),
+        sample_count=sample_count,
+        combiner=combiner,
+    )
+
+    golden_row_ids = (
+        row_ids % (sample_count // num_sc_per_chip)
+        + int(row_offset // num_sc_per_chip)
+        + int(stacked_table_sample_count // num_sc_per_chip)
+        * (row_ids // (sample_count // num_sc_per_chip))
+    )
+    golden_col_ids = (
+        bitwise_ops.bitwise_and(col_ids + col_shift, num_sc_shards_bit_mod)
+        + bitwise_ops.bitwise_and(col_ids, num_sc_shards_bit_mod_inv)
+        + col_offset
+    )
+
+    row_ids_list, col_ids_list, gains_list = (
+        xla_ops.convert_to_list_of_sparse_core_coo_tensors(
+            indices_or_row_splits=math_ops.cast(
+                sparse_feature.indices, dtype=dtypes.int32
+            ),
+            values=math_ops.cast(sparse_feature.values, dtype=dtypes.int32),
+            weights=1.0,
+            sample_count=sample_count,
+            combiner=combiner,
+            num_sc_per_chip=4,
+            row_offset=row_offset,
+            col_offset=col_offset,
+            col_shift=col_shift,
+            num_sc_shards=num_sc_shards,
+            stacked_table_sample_count=stacked_table_sample_count,
+        )
+    )
+
+    self.assertAllClose(golden_row_ids, array_ops.concat(row_ids_list, axis=0))
+    self.assertAllClose(golden_col_ids, array_ops.concat(col_ids_list, axis=0))
+    self.assertAllClose(gains, array_ops.concat(gains_list, axis=0))
+
+  def test_convert_to_list_of_sparse_core_coo_tensors(self):
+    sample_count = 16
+    token_count = 1024
+    combiner = "sum"
+    sparse_feature = sparse_ops.sparse_reorder(
+        sparse_tensor.SparseTensor(
+            indices=[[i % sample_count, i] for i in np.arange(token_count)],
+            values=np.arange(token_count),
+            dense_shape=[sample_count, 1024],
+        )
+    )
+
+    row_ids_list, col_ids_list, gains_list = (
+        xla_ops.convert_to_list_of_sparse_core_coo_tensors(
+            indices_or_row_splits=math_ops.cast(
+                sparse_feature.indices, dtype=dtypes.int32
+            ),
+            values=math_ops.cast(sparse_feature.values, dtype=dtypes.int32),
+            weights=1.0,
+            sample_count=sample_count,
+            combiner=combiner,
+            num_sc_per_chip=4,
+            row_offset=0,
+            col_offset=0,
+            col_shift=0,
+            num_sc_shards=16,
+            stacked_table_sample_count=sample_count,
+        )
+    )
+
+    sorted_row_ids_list = []
+    sorted_col_ids_list = []
+    sorted_gains_list = []
+    id_counts_list = []
+    for i in range(4):
+      (
+          sorted_row_ids,
+          sorted_col_ids,
+          sorted_gains,
+          id_counts,
+      ) = xla_ops.sort_list_of_sparse_core_coo_tensors(
+          row_ids_list=[row_ids_list[i]],
+          col_ids_list=[col_ids_list[i]],
+          gains_list=[gains_list[i]],
+          sample_count_list=[sample_count // 4],
+          col_offset_list=[0],
+          num_replica=4,
+          table_vocab_size=16384,
+          feature_width=16,
+          num_sc_per_chip=4,
+          max_ids_per_sparse_core=256,
+          max_unique_ids_per_sparse_core=256,
+          table_name="table",
+      )
+      sorted_row_ids_list.append(sorted_row_ids)
+      sorted_col_ids_list.append(sorted_col_ids)
+      sorted_gains_list.append(sorted_gains)
+      id_counts_list.append(id_counts)
+
+    (
+        row_pointers,
+        sorted_sample_ids,
+        sorted_token_ids,
+        sorted_gains,
+        row_pointers_unpadded_size,
+        ids_unpadded_size,
+        num_minibatches_per_sc,
+    ) = xla_ops.convert_to_sparse_core_csr_wrapped_coo_tensor(
+        sorted_row_ids_list=sorted_row_ids_list,
+        sorted_col_ids_list=sorted_col_ids_list,
+        sorted_gains_list=sorted_gains_list,
+        id_counts_list=id_counts_list,
+        splits=constant_op.constant(0, dtype=dtypes.int64),
+        sample_count_per_sc=sample_count // 4,
+        max_minibatches_per_sc=4,
+        max_ids_per_chip_per_sample=64,
+        table_vocab_size=16384,
+        feature_width=16,
+        num_replica=4,
+        allow_id_dropping=False,
+        table_name="table",
+    )
+
+    (
+        golden_row_pointers,
+        golden_sorted_sample_ids,
+        golden_sorted_token_ids,
+        golden_sorted_gains,
+        golden_row_pointers_unpadded_size,
+        golden_ids_unpadded_size,
+        golden_num_minibatches_per_sc,
+    ) = _convert_coo_tensor_to_csr_with_physical_replica(
+        row_ids=array_ops.concat(row_ids_list, axis=0),
+        col_ids=array_ops.concat(col_ids_list, axis=0),
+        gains=array_ops.concat(gains_list, axis=0),
+        splits=constant_op.constant(0, dtype=dtypes.int64),
+        sample_count=sample_count,
+        num_replica=4,
+        max_minibatches_per_sc=4,
+        max_ids_per_chip_per_sample=64,
+        table_vocab_size=16384,
+    )
+
+    self.assertAllClose(
+        golden_row_pointers[:golden_row_pointers_unpadded_size],
+        row_pointers[:row_pointers_unpadded_size],
+    )
+    self.assertAllClose(
+        golden_sorted_sample_ids[:golden_ids_unpadded_size],
+        sorted_sample_ids[:ids_unpadded_size],
+    )
+    self.assertAllClose(
+        golden_sorted_token_ids[:golden_ids_unpadded_size],
+        sorted_token_ids[:ids_unpadded_size],
+    )
+    self.assertAllClose(
+        golden_sorted_gains[:golden_ids_unpadded_size],
+        sorted_gains[:ids_unpadded_size],
+    )
+    self.assertEqual(golden_num_minibatches_per_sc, num_minibatches_per_sc)
+
+  def test_get_stats_from_list_of_sparse_core_coo_tensors(self):
+    sample_count = 16
+    token_count = 1024
+    combiner = "sum"
+    sparse_feature = sparse_ops.sparse_reorder(
+        sparse_tensor.SparseTensor(
+            indices=[
+                [i % sample_count, i]
+                for i in np.random.randint(low=0, high=1024, size=token_count)
+            ],
+            values=np.random.randint(low=0, high=1024, size=token_count),
+            dense_shape=[sample_count, 1024],
+        )
+    )
+    max_ids_golden = 0
+    max_unique_ids_golden = 0
+
+    for i in range(4):
+      sparse_feature_slice = sparse_ops.sparse_slice(
+          sparse_feature,
+          [i * sample_count // 4, 0],
+          [sample_count // 4, 1024],
+      )
+      max_ids_per_sparse_core, max_uniques_per_sparse_core = (
+          _compute_sparse_core_stats(
+              sparse_feature_slice.indices[:, 0],
+              sparse_feature_slice.values,
+              num_sc_shards=16,
+          )
+      )
+      max_ids_golden = max(max_ids_golden, max_ids_per_sparse_core)
+      max_unique_ids_golden = max(
+          max_unique_ids_golden, max_uniques_per_sparse_core
+      )
+
+    row_ids_list, col_ids_list, gains_list = (
+        xla_ops.convert_to_list_of_sparse_core_coo_tensors(
+            indices_or_row_splits=math_ops.cast(
+                sparse_feature.indices, dtype=dtypes.int32
+            ),
+            values=math_ops.cast(sparse_feature.values, dtype=dtypes.int32),
+            weights=1.0,
+            sample_count=sample_count,
+            combiner=combiner,
+            num_sc_per_chip=4,
+            row_offset=0,
+            col_offset=0,
+            col_shift=0,
+            num_sc_shards=16,
+            stacked_table_sample_count=sample_count,
+        )
+    )
+
+    max_ids = 0
+    max_uniques = 0
+    for i in range(4):
+      max_ids_per_sparse_core, max_unique_ids_per_sparse_core = (
+          xla_ops.get_stats_from_list_of_sparse_core_coo_tensors(
+              row_ids_list=[row_ids_list[i]],
+              col_ids_list=[col_ids_list[i]],
+              gains_list=[gains_list[i]],
+              sample_count_list=[sample_count // 4],
+              col_offset_list=[0],
+              num_replica=4,
+              table_vocab_size=16384,
+              feature_width=16,
+              num_sc_per_chip=4,
+              table_name="table",
+          )
+      )
+      max_ids = max(max_ids, max_ids_per_sparse_core)
+      max_uniques = max(max_uniques, max_unique_ids_per_sparse_core)
+
+    self.assertEqual(max_ids, max_ids_golden)
+    self.assertEqual(max_uniques, max_unique_ids_golden)
+
+  def test_sort_list_of_sparse_core_coo_tensors(self):
+    sample_count = 16
+    token_count = 1024
+    combiner = "sum"
+    num_chips = 4
+    sparse_feature = sparse_ops.sparse_reorder(
+        sparse_tensor.SparseTensor(
+            indices=[[i % sample_count, i] for i in np.arange(token_count)],
+            values=np.arange(token_count),
+            dense_shape=[sample_count, 1024],
+        )
+    )
+
+    row_ids_list, col_ids_list, gains_list = (
+        xla_ops.convert_to_list_of_sparse_core_coo_tensors(
+            indices_or_row_splits=math_ops.cast(
+                sparse_feature.indices, dtype=dtypes.int32
+            ),
+            values=math_ops.cast(sparse_feature.values, dtype=dtypes.int32),
+            weights=1.0,
+            sample_count=sample_count,
+            combiner=combiner,
+            num_sc_per_chip=4,
+            row_offset=0,
+            col_offset=0,
+            col_shift=0,
+            num_sc_shards=num_chips * 4,
+            stacked_table_sample_count=sample_count,
+        )
+    )
+
+    for i in range(4):
+      (
+          sorted_row_ids,
+          sorted_col_ids,
+          sorted_gains,
+          _,
+      ) = xla_ops.sort_list_of_sparse_core_coo_tensors(
+          row_ids_list=[row_ids_list[i]],
+          col_ids_list=[col_ids_list[i]],
+          gains_list=[gains_list[i]],
+          sample_count_list=[sample_count // 4],
+          col_offset_list=[0],
+          num_replica=num_chips,
+          table_vocab_size=16384,
+          feature_width=16,
+          num_sc_per_chip=4,
+          max_ids_per_sparse_core=256,
+          max_unique_ids_per_sparse_core=256,
+          table_name="table",
+      )
+
+      embedding_lookup_inputs = []
+      for row_id, col_id, gain in zip(
+          row_ids_list[i], col_ids_list[i], gains_list[i]
+      ):
+        embedding_lookup_inputs.append((col_id % 16, col_id, row_id, gain))
+      # sort based on replica id first, then col_id.
+      embedding_lookup_inputs.sort()
+
+      self.assertAllClose(
+          sorted_row_ids,
+          [inp[2] % (sample_count // 4) for inp in embedding_lookup_inputs],
+      )
+      self.assertAllClose(
+          sorted_col_ids,
+          [inp[1] // (num_chips * 4) for inp in embedding_lookup_inputs],
+      )
+      self.assertAllClose(
+          sorted_gains, [inp[3] for inp in embedding_lookup_inputs]
+      )
+
+  def test_id_dropping_with_convert_to_list_of_sparse_core_coo_tensors(self):
+    sample_count = 16
+    token_count = 1024
+    combiner = "sum"
+    sparse_feature = sparse_ops.sparse_reorder(
+        sparse_tensor.SparseTensor(
+            indices=[[i % sample_count, i] for i in np.arange(token_count)],
+            values=np.arange(token_count),
+            dense_shape=[sample_count, 1024],
+        )
+    )
+
+    row_ids_list, col_ids_list, gains_list = (
+        xla_ops.convert_to_list_of_sparse_core_coo_tensors(
+            indices_or_row_splits=math_ops.cast(
+                sparse_feature.indices, dtype=dtypes.int32
+            ),
+            values=math_ops.cast(sparse_feature.values, dtype=dtypes.int32),
+            weights=1.0,
+            sample_count=sample_count,
+            combiner=combiner,
+            num_sc_per_chip=4,
+            row_offset=0,
+            col_offset=0,
+            col_shift=0,
+            num_sc_shards=16,
+            stacked_table_sample_count=sample_count,
+        )
+    )
+
+    sorted_row_ids_list = []
+    sorted_col_ids_list = []
+    sorted_gains_list = []
+    id_counts_list = []
+    for i in range(4):
+      (
+          sorted_row_ids,
+          sorted_col_ids,
+          sorted_gains,
+          id_counts,
+      ) = xla_ops.sort_list_of_sparse_core_coo_tensors(
+          row_ids_list=[row_ids_list[i]],
+          col_ids_list=[col_ids_list[i]],
+          gains_list=[gains_list[i]],
+          sample_count_list=[sample_count // 4],
+          col_offset_list=[0],
+          num_replica=4,
+          table_vocab_size=16384,
+          feature_width=16,
+          num_sc_per_chip=4,
+          max_ids_per_sparse_core=256,
+          max_unique_ids_per_sparse_core=256,
+          table_name="table",
+      )
+      sorted_row_ids_list.append(sorted_row_ids)
+      sorted_col_ids_list.append(sorted_col_ids)
+      sorted_gains_list.append(sorted_gains)
+      id_counts_list.append(id_counts)
+
+    # If not allow id dropping, the op will fail with very small
+    # max_ids_per_chip_per_sample.
+    with self.assertRaises(Exception):
+      xla_ops.convert_to_sparse_core_csr_wrapped_coo_tensor(
+          sorted_row_ids_list=sorted_row_ids_list,
+          sorted_col_ids_list=sorted_col_ids_list,
+          sorted_gains_list=sorted_gains_list,
+          id_counts_list=id_counts_list,
+          splits=constant_op.constant(0, dtype=dtypes.int64),
+          sample_count_per_sc=sample_count // 4,
+          max_minibatches_per_sc=4,
+          max_ids_per_chip_per_sample=8,
+          table_vocab_size=16384,
+          feature_width=16,
+          num_replica=4,
+          allow_id_dropping=False,
+          table_name="table",
+      )
+
+    # Allow id dropping, the op will succeed,
+    xla_ops.convert_to_sparse_core_csr_wrapped_coo_tensor(
+        sorted_row_ids_list=sorted_row_ids_list,
+        sorted_col_ids_list=sorted_col_ids_list,
+        sorted_gains_list=sorted_gains_list,
+        id_counts_list=id_counts_list,
+        splits=constant_op.constant(0, dtype=dtypes.int64),
+        sample_count_per_sc=sample_count // 4,
+        max_minibatches_per_sc=4,
+        max_ids_per_chip_per_sample=8,
+        table_vocab_size=16384,
+        feature_width=16,
+        num_replica=4,
+        allow_id_dropping=True,
+        table_name="table",
+    )
+
+  def test_invalid_id_counts_with_convert_to_sparse_core_csr_wrapped_coo_tensor(
+      self,
+  ):
+    sample_count = 16
+    token_count = 1024
+    combiner = "sum"
+    sparse_feature = sparse_ops.sparse_reorder(
+        sparse_tensor.SparseTensor(
+            indices=[[i % sample_count, i] for i in np.arange(token_count)],
+            values=np.arange(token_count),
+            dense_shape=[sample_count, 1024],
+        )
+    )
+
+    row_ids_list, col_ids_list, gains_list = (
+        xla_ops.convert_to_list_of_sparse_core_coo_tensors(
+            indices_or_row_splits=math_ops.cast(
+                sparse_feature.indices, dtype=dtypes.int32
+            ),
+            values=math_ops.cast(sparse_feature.values, dtype=dtypes.int32),
+            weights=1.0,
+            sample_count=sample_count,
+            combiner=combiner,
+            num_sc_per_chip=4,
+            row_offset=0,
+            col_offset=0,
+            col_shift=0,
+            num_sc_shards=16,
+            stacked_table_sample_count=sample_count,
+        )
+    )
+
+    sorted_row_ids_list = []
+    sorted_col_ids_list = []
+    sorted_gains_list = []
+    id_counts_list = []
+    for i in range(4):
+      (
+          sorted_row_ids,
+          sorted_col_ids,
+          sorted_gains,
+          id_counts,
+      ) = xla_ops.sort_list_of_sparse_core_coo_tensors(
+          row_ids_list=[row_ids_list[i]],
+          col_ids_list=[col_ids_list[i]],
+          gains_list=[gains_list[i]],
+          sample_count_list=[sample_count // 4],
+          col_offset_list=[0],
+          num_replica=4,
+          table_vocab_size=16384,
+          feature_width=16,
+          num_sc_per_chip=4,
+          max_ids_per_sparse_core=256,
+          max_unique_ids_per_sparse_core=256,
+          table_name="table",
+      )
+      sorted_row_ids_list.append(sorted_row_ids)
+      sorted_col_ids_list.append(sorted_col_ids)
+      sorted_gains_list.append(sorted_gains)
+
+      # Manually modify id_counts to trigger validation error
+      id_counts_size = array_ops.shape(id_counts)[0]
+      id_counts_before = array_ops.slice(id_counts, [0], [id_counts_size - 1])
+      last_element = array_ops.slice(id_counts, [id_counts_size - 1], [1])
+      new_last_element = last_element + 1
+      invalid_id_counts = array_ops.concat(
+          [id_counts_before, new_last_element], axis=0
+      )
+      id_counts_list.append(invalid_id_counts)
+
+    with self.assertRaisesRegex(
+        Exception,
+        "should not exceed the size of the corresponding sorted tensors",
+    ):
+      xla_ops.convert_to_sparse_core_csr_wrapped_coo_tensor(
+          sorted_row_ids_list=sorted_row_ids_list,
+          sorted_col_ids_list=sorted_col_ids_list,
+          sorted_gains_list=sorted_gains_list,
+          id_counts_list=id_counts_list,
+          splits=constant_op.constant(0, dtype=dtypes.int64),
+          sample_count_per_sc=sample_count // 4,
+          max_minibatches_per_sc=4,
+          max_ids_per_chip_per_sample=64,
+          table_vocab_size=16384,
+          feature_width=16,
+          num_replica=4,
+          allow_id_dropping=False,
+          table_name="table",
+      )
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tpu_embedding_v3_utils.py

@@ -0,0 +1,277 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Utils for Sparsecore Checkpoints."""
+
+import functools
+from typing import Any, Dict
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework.constant_op import constant as tf_constant
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_control_flow_ops
+from tensorflow.python.ops import manip_ops
+from tensorflow.python.ops import variables as tf_variables
+from tensorflow.python.trackable import base as trackable_base
+
+SPARSECORE_LAYOUTS_CHECKPOINT_KEY = "_sparse_core_table_layouts"
+
+
+def unshuffle_from_sc_to_cpu(
+    t: tensor.Tensor,
+    num_sparse_cores: int,
+    offset_in_shard: int,
+    size_in_shard: int,
+    shard_rotation: int = 0,
+) -> tensor.Tensor:
+  """Unshuffles the sparse core sharded embedding tables to unsharded.
+
+  This converts an input tensor respresenting stacked and sharded embedding
+  table into a specific embedding table variable by using the provided
+  metadata about the said table within the stacked, sharded embedding table.
+  Args:
+    t: The input stacked and sharded embedding table from sparsecore.
+    num_sparse_cores: The number of sparsecores, this determines the number of
+      shards that are present in the input t.
+    offset_in_shard: Offset within a shard where the queried table starts.
+    size_in_shard: size (number of rows) of this queried table within each shard
+      of the input t.
+    shard_rotation: The rotation of this table's shards.
+
+  Returns:
+    An embedding table which is part of the stacked embedding table t.
+  """
+  old_shape = t.shape
+  # The width of the table must be a multiple of number of SC devices. The
+  # tpu strategy does this round off at training time so we expect the
+  # checkpoints value to meet this requirement.
+  if t.shape[0] % num_sparse_cores != 0:
+    raise ValueError(
+        "The first dim of the table ({}) should be multiple of number of sparse"
+        " cores ({})".format(t.shape[0], num_sparse_cores)
+    )
+  # get shards in the input t
+  shards_t = array_ops.reshape(
+      t,
+      (
+          num_sparse_cores,
+          t.shape[0] // num_sparse_cores,
+          t.shape[1],
+      ),
+  )
+  # From each shard in t, get the part for just the queried table.
+  shards = shards_t[:, offset_in_shard : offset_in_shard + size_in_shard, :]
+  # This table's shards were rotated by `shard_rotation`, so we need to rotate
+  # the same amount in opposite direction
+  if shard_rotation:
+    shards = manip_ops.roll(shards, -shard_rotation, axis=0)
+  # Re-arrange (transpose and reshape) the shards to get the queried embedding
+  # table.
+  intermediate_tensor = array_ops.transpose(shards, (1, 0, 2))
+  new_shape = size_in_shard * num_sparse_cores, old_shape[1]
+  return array_ops.reshape(intermediate_tensor, new_shape)
+
+
+def remove_padding_from_sc(
+    value_in_checkpoint: tensor.Tensor, variable_shape: tuple[int, int]
+) -> tensor.Tensor:
+  """Removes padding, if any, from sparsecore checkpoint.
+
+  Args:
+    value_in_checkpoint: input tensor value, usually from checkpoint.
+    variable_shape: Expected shape of tensor after removing padding.
+
+  Returns:
+    A slice of the input tensor to match the variable_shape if the
+    variable shape is a valid slice if the input tensor.
+  """
+  checkpoint_value_shape = value_in_checkpoint.shape.as_list()
+  # If the checkpoint shape is at least the size of the variable, we conclude
+  # that the extra rows and cols must be padding.
+  is_init_value_padded = all(
+      [i >= j for i, j in zip(checkpoint_value_shape, variable_shape)]
+  )
+  if not is_init_value_padded:
+    return value_in_checkpoint
+  # checkpoint has padding so we can remove it.
+  begin = [0] * len(checkpoint_value_shape)
+  return array_ops.slice(value_in_checkpoint, begin=begin, size=variable_shape)
+
+
+def map_indices_in_shard(
+    num_sparse_cores: int,
+    offset_in_shard: int,
+    shard_rotation: int,
+    row_indices: tensor.Tensor,
+) -> tuple[tensor.Tensor, tensor.Tensor]:
+  """Maps a row of a given table to its sparse core shard and position.
+
+  Maps a given a row index of a logical table and its layout in sparse core,
+  returns the index of the shard where the row is placed and its relative
+  position within
+  that sparse core shard.
+  Args:
+    num_sparse_cores: The number of sparsecores, this determines the number of
+      shards present.
+    offset_in_shard: Offset within a shard where the queried table starts.
+    shard_rotation: The rotation of this table's shards.
+    row_indices: row indices of the embedding table being looked up.
+
+  Returns:
+    A Tuple representing shard_index and position of the row in that shard.
+  """
+  shard_index = (
+      (row_indices % num_sparse_cores) + shard_rotation
+  ) % num_sparse_cores
+  position_in_shard = offset_in_shard + row_indices // num_sparse_cores
+  return (shard_index, position_in_shard)
+
+
+class SparseCoreLayoutsTrackable(trackable_base.Trackable):
+  """Trackable for sparsecore layouts used in training."""
+
+  def __init__(self, proto_str_tensor: tensor.Tensor):
+    self.value = proto_str_tensor
+
+  def _serialize_to_tensors(self) -> Dict[str, tensor.Tensor]:
+    return {trackable_base.VARIABLE_VALUE_KEY: self.value}
+
+  def _restore_from_tensors(
+      self, restored_tensors: Dict[str, tensor.Tensor]
+  ) -> None:
+    # Do not restore the layouts proto from checkpoint, it should always match
+    # the embedding, which is set at build time.
+    gen_control_flow_ops.no_op()
+
+
+class SparseCoreStackedTableTrackable(trackable_base.Trackable):
+  """Trackable for stacked tables generated from sparse core."""
+
+  def __init__(self, stacked_layouts, table_to_config):
+    self.vars = {}
+    self._stacked_layouts = stacked_layouts
+    for table_layout in stacked_layouts:
+      variable_shape = tuple(table_layout.unsharded_shape)
+      self.vars[table_layout.table_name] = tf_variables.Variable(
+          name=table_layout.table_name,
+          initial_value=functools.partial(
+              table_to_config[table_layout.table_name].initializer,
+              variable_shape,
+              dtype=dtypes.float32,
+          ),
+          shape=variable_shape,
+          dtype=dtypes.float32,
+      )
+  # TODO(b/312743130): This is a workaround. During checkpoint restoration
+  # optimizer expects the trackable to provide a `_unique_id` or equivalent.
+  # Remove this when the bug is fixed.
+  @property
+  def _unique_id(self):
+    return self.vars[self._stacked_layouts[0].table_name]._unique_id
+
+  def _serialize_to_tensors(self) -> Any:
+    return {
+        # We need to export some variable here for restore to pick
+        # the checkpoint key the actual value is not important so 0 works
+        trackable_base.VARIABLE_VALUE_KEY: tf_constant(
+            0.0, dtype=dtypes.float32
+        ),
+    }
+
+  def _restore_from_tensors(self, restored_tensors: Dict[str, tensor.Tensor]):
+    def fn(restored_tensors):
+      value_from_checkpoint = restored_tensors[
+          trackable_base.VARIABLE_VALUE_KEY
+      ]
+      # Do unsharding to get the individual tables from the stacked table in
+      # checkpoint
+      for layout in self._stacked_layouts:
+        variable_shape = (
+            layout.unsharded_shape[0],
+            layout.unsharded_shape[1],
+        )
+        t_part = unshuffle_from_sc_to_cpu(
+            t=value_from_checkpoint,
+            num_sparse_cores=layout.num_sparse_cores,
+            offset_in_shard=layout.sparse_core_shard_row_offset,
+            size_in_shard=(
+                layout.unsharded_padded_shape[0] // layout.num_sparse_cores
+            ),
+            shard_rotation=layout.sparse_core_shard_rotation,
+        )
+        t_part = remove_padding_from_sc(t_part, variable_shape)
+        self.vars[layout.table_name].assign(t_part)
+
+    return fn(restored_tensors)
+
+  def get_var(self, name: str) -> tf_variables.Variable:
+    return self.vars[name]
+
+  def get_vars(self) -> Dict[str, tf_variables.Variable]:
+    return self.vars
+
+  def __repr__(self):
+    return "SparseCoreStackedTableTrackable({})".format(self.vars.keys())
+
+
+def shard_table(
+    num_sparse_cores: int,
+    table: tensor.Tensor,
+) -> tensor.Tensor:
+  """Convert a table to the internal layout.
+
+  Args:
+    num_sparse_cores: The total number of sparse cores.
+    table: The full table, unsharded.
+
+  Returns:
+    A tensor containing the sharded value for this table.
+  """
+  assert table.shape[0] % num_sparse_cores == 0
+
+  # Do the sparse core rotation:
+  tmp = array_ops.reshape(
+      table,
+      [-1, num_sparse_cores, table.shape[1]],
+  )
+  # The mod sharding across sparse cores.
+  tmp = array_ops.transpose(tmp, [1, 0, 2])
+  return array_ops.reshape(tmp, [-1, table.shape[1]])
+
+
+def shard_initializer(strategy, initializer) -> tensor.Tensor:
+  """Wraps an initializer to convert a table to the internal layout."""
+
+  num_devices = strategy.extended._tpu_devices.size  # pylint: disable=protected-access
+  num_sc_per_chip = (
+      strategy.extended.tpu_hardware_feature.num_embedding_devices_per_chip
+  )
+  num_scs = num_devices * num_sc_per_chip
+
+  @functools.wraps(initializer)
+  def wrapper(shape, dtype, shard_info=None):
+    # Initializes the whole table.
+    table = initializer(shape, dtype)
+    if shard_info is None:
+      return table
+
+    # Convert the table to the internal layout.
+    table = shard_table(num_scs, table)
+    # Pull out the shard of interest.
+    return table[
+        shard_info.offset[0] : shard_info.offset[0] + shard_info.shape[0],
+        shard_info.offset[1] : shard_info.offset[1] + shard_info.shape[1],
+    ]
+
+  return wrapper

tensorflow/python/tpu/tpu_embedding_v3_utils_test.py

@@ -0,0 +1,278 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Test for tpu_embedding_v3_utils."""
+
+import collections
+
+from absl.testing import parameterized
+
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.framework.constant_op import constant as tf_constant
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_embedding_v3_utils as v3_utils
+
+TestTable = collections.namedtuple("Table", ["vocab", "dim", "shift"])
+
+
+def create_test_table_shards(
+    table: TestTable, num_sc_shards: int, table_data_start=0
+):
+  t = array_ops.reshape(
+      math_ops.range(
+          start=table_data_start,
+          delta=1,
+          limit=table_data_start + table.vocab * table.dim,
+      ),
+      (table.vocab, table.dim),
+  )
+  shards = [t[i::num_sc_shards, :] for i in range(num_sc_shards)]
+  if table.shift:
+    shards = collections.deque(shards)
+    shards.rotate(table.shift)
+    return (t, list(shards))
+  else:
+    return (t, shards)
+
+
+class TpuEmbeddingV3UtilsTest(test.TestCase, parameterized.TestCase):
+
+  def test_unpadding(self):
+    self.assertAllEqual(
+        v3_utils.remove_padding_from_sc(
+            array_ops.ones((4, 5)), variable_shape=(3, 2)
+        ),
+        array_ops.ones((3, 2)),
+    )
+    x = array_ops.reshape(math_ops.range(12), (3, 4))
+    self.assertAllEqual(
+        v3_utils.remove_padding_from_sc(x, variable_shape=(2, 2)),
+        tf_constant([[0, 1], [4, 5]]),
+    )
+    self.assertAllEqual(
+        v3_utils.remove_padding_from_sc(x, variable_shape=(3, 5)),
+        x,
+    )
+
+  @parameterized.named_parameters(
+      ("one", 8, 4, 4), ("two", 27, 6, 3), ("three", 128, 8, 4)
+  )
+  def test_unshuffle_one_table_basic(self, vocab, dim, num_sc):
+    # input vocab should be multiple of num_sc
+    self.assertEqual(vocab % num_sc, 0)
+    x, shards = create_test_table_shards(
+        TestTable(vocab=vocab, dim=dim, shift=0), num_sc
+    )
+    x_sharded = array_ops.concat(shards, axis=0)
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=x_sharded,
+            num_sparse_cores=num_sc,
+            offset_in_shard=0,
+            size_in_shard=vocab // num_sc,
+            shard_rotation=0,
+        ),
+        x,
+    )
+
+  def test_unshuffle_stacking_basic(self):
+    num_sc = 4
+    ta = TestTable(vocab=12, dim=4, shift=0)
+    tb = TestTable(vocab=32, dim=4, shift=1)
+    x, x_shards = create_test_table_shards(ta, num_sc)
+    y, y_shards = create_test_table_shards(tb, num_sc)
+    stacked_shards = [
+        array_ops.concat([i, j], axis=0) for i, j in zip(x_shards, y_shards)
+    ]
+    stacked = array_ops.concat(stacked_shards, axis=0)
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=stacked,
+            num_sparse_cores=num_sc,
+            offset_in_shard=0,
+            size_in_shard=ta.vocab // num_sc,
+            shard_rotation=ta.shift,
+        ),
+        x,
+    )
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=stacked,
+            num_sparse_cores=num_sc,
+            offset_in_shard=ta.vocab // num_sc,
+            size_in_shard=tb.vocab // num_sc,
+            shard_rotation=tb.shift,
+        ),
+        y,
+    )
+
+  def test_unshuffle_stacking_many_tables(self):
+    num_sc = 4
+    tables = [
+        TestTable(vocab=12, dim=4, shift=0),
+        TestTable(vocab=32, dim=4, shift=1),
+        TestTable(vocab=32, dim=4, shift=2),
+        TestTable(vocab=32, dim=4, shift=3),
+        TestTable(vocab=32, dim=4, shift=4),
+        TestTable(vocab=32, dim=4, shift=5),
+    ]
+    u, u_shards = create_test_table_shards(tables[0], num_sc)
+    v, v_shards = create_test_table_shards(tables[1], num_sc)
+    w, w_shards = create_test_table_shards(tables[2], num_sc)
+    x, x_shards = create_test_table_shards(tables[3], num_sc)
+    y, y_shards = create_test_table_shards(tables[4], num_sc)
+    z, z_shards = create_test_table_shards(tables[5], num_sc)
+    stacked_shards = [
+        array_ops.concat([i, j, k, l, m, n], axis=0)
+        for i, j, k, l, m, n in zip(
+            u_shards, v_shards, w_shards, x_shards, y_shards, z_shards
+        )
+    ]
+    stacked = array_ops.concat(stacked_shards, axis=0)
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=stacked,
+            num_sparse_cores=num_sc,
+            offset_in_shard=0,
+            size_in_shard=tables[0].vocab // num_sc,
+            shard_rotation=tables[0].shift,
+        ),
+        u,
+    )
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=stacked,
+            num_sparse_cores=num_sc,
+            offset_in_shard=tables[0].vocab // num_sc,
+            size_in_shard=tables[1].vocab // num_sc,
+            shard_rotation=tables[1].shift,
+        ),
+        v,
+    )
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=stacked,
+            num_sparse_cores=num_sc,
+            offset_in_shard=(tables[0].vocab + tables[1].vocab) // num_sc,
+            size_in_shard=tables[2].vocab // num_sc,
+            shard_rotation=tables[2].shift,
+        ),
+        w,
+    )
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=stacked,
+            num_sparse_cores=num_sc,
+            offset_in_shard=(
+                tables[0].vocab + tables[1].vocab + tables[2].vocab
+            )
+            // num_sc,
+            size_in_shard=tables[3].vocab // num_sc,
+            shard_rotation=tables[3].shift,
+        ),
+        x,
+    )
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=stacked,
+            num_sparse_cores=num_sc,
+            offset_in_shard=(
+                tables[0].vocab
+                + tables[1].vocab
+                + tables[2].vocab
+                + tables[3].vocab
+            )
+            // num_sc,
+            size_in_shard=tables[4].vocab // num_sc,
+            shard_rotation=tables[4].shift,
+        ),
+        y,
+    )
+    self.assertAllEqual(
+        v3_utils.unshuffle_from_sc_to_cpu(
+            t=stacked,
+            num_sparse_cores=num_sc,
+            offset_in_shard=(
+                tables[0].vocab
+                + tables[1].vocab
+                + tables[2].vocab
+                + tables[3].vocab
+                + tables[4].vocab
+            )
+            // num_sc,
+            size_in_shard=tables[5].vocab // num_sc,
+            shard_rotation=tables[5].shift,
+        ),
+        z,
+    )
+
+  def test_index_mapping_one_table(self):
+    num_sc = 4
+    x, shards = create_test_table_shards(
+        TestTable(vocab=12, dim=4, shift=0), num_sc
+    )
+    indices = tf_constant([1, 2, 5, 7, 9])
+    shard_idx, position_in_shard = v3_utils.map_indices_in_shard(
+        num_sparse_cores=num_sc,
+        offset_in_shard=0,
+        shard_rotation=0,
+        row_indices=indices,
+    )
+    self.assertAllEqual(
+        shard_idx,
+        indices % num_sc,
+    )
+    self.assertAllEqual(
+        [x[i] for i in indices],
+        [shards[j][k] for j, k in zip(shard_idx, position_in_shard)],
+    )
+
+  def test_index_mapping_stacked_tables(self):
+    num_sc = 4
+    ta = TestTable(vocab=12, dim=4, shift=0)
+    tb = TestTable(vocab=32, dim=4, shift=1)
+    x, x_shards = create_test_table_shards(ta, num_sc)
+    y, y_shards = create_test_table_shards(tb, num_sc, table_data_start=100)
+    stacked_shards = [
+        array_ops.concat([i, j], axis=0) for i, j in zip(x_shards, y_shards)
+    ]
+    indices_ta = tf_constant([1, 2, 7, 9, 11])
+    shard_idx, position_in_shard = v3_utils.map_indices_in_shard(
+        num_sparse_cores=num_sc,
+        offset_in_shard=0,
+        shard_rotation=ta.shift,
+        row_indices=indices_ta,
+    )
+    self.assertAllEqual(
+        [x[i] for i in indices_ta],
+        [stacked_shards[j][k] for j, k in zip(shard_idx, position_in_shard)],
+    )
+    indices_tb = tf_constant([1, 2, 7, 9, 15, 27])
+    shard_idx, position_in_shard = v3_utils.map_indices_in_shard(
+        num_sparse_cores=num_sc,
+        offset_in_shard=ta.vocab // num_sc,
+        shard_rotation=tb.shift,
+        row_indices=indices_tb,
+    )
+    self.assertAllEqual(
+        [y[i] for i in indices_tb],
+        [stacked_shards[j][k] for j, k in zip(shard_idx, position_in_shard)],
+    )
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/python/tpu/tpu_feed.py

@@ -0,0 +1,956 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===================================================================
+
+"""Helper library for handling infeed between hosts and TPUs.
+"""
+
+import itertools
+
+import numpy as np
+
+from tensorflow.python.compiler.xla.experimental import xla_sharding
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.tpu import tpu_name_util
+from tensorflow.python.tpu import tpu_sharding
+from tensorflow.python.tpu.ops import tpu_ops
+
+from tensorflow.python.util import nest
+
+
+def partition_or_replicate_on_host(tensor, dims):
+  """Partitions or replicates the input tensor.
+
+    The ops inside this function are placed on the host side.
+
+  Args:
+    tensor: The input tensor which will be partitioned or replicated.
+    dims: A list of integer describes how to partition the input tensor.
+
+  Returns:
+    An iterator of `Tensor`s or a list of partitioned tensors.
+  """
+  if dims is None:
+    return itertools.repeat(tensor)
+  dims = np.array(dims)
+  output = [tensor]
+  shape_list = np.array(tensor.shape.as_list())
+  quotients, remainders = np.divmod(shape_list, dims)
+  for axis, (quotient, remainder, dim, original_size) in enumerate(
+      zip(quotients, remainders, dims, shape_list)):
+    if dim <= 1:
+      continue
+    if remainder > 0:
+      # For each dimension, when it cannot be evenly partitioned, XLA assumes
+      # tensors are partitioned in a greedy manner by using
+      # ceil_ratio(size/dim) first. E.g. 2D tensor with shape (5, 14) and dims
+      # are (2, 4). Since 5 % 2 = 1 and 14 % 4 = 2, [5, 14] =>
+      # [[(3, 4), (3, 4), (2, 4), (2, 2)],
+      # [(2, 4), (2, 4), (2, 4), (2, 2)]]
+      ceil_ratio = quotient + 1
+      num_full_slots, left_over = np.divmod(original_size, ceil_ratio)
+      num_or_size_splits = [ceil_ratio] * num_full_slots + [left_over]
+      if len(num_or_size_splits) < dim:
+        num_or_size_splits += [0] * (dim - len(num_or_size_splits))
+      new_output = []
+      for x in output:
+        new_output.append(
+            array_ops.split(
+                x, num_or_size_splits=num_or_size_splits, axis=axis))
+      output = new_output
+    else:
+      output = [array_ops.split(x, int(dim), axis=axis) for x in output]
+    output = nest.flatten(output)
+  return output
+
+
+def _tag_sharding_attribute_for_dequeued_tensor(tensor, dims):
+  """Tags appropriate XLA sharding attribute to the dequeued tensor.
+
+  The sharding attribute of the dequeued tensor will be a tuple.
+
+  Args:
+    tensor: The dequeued tensor on TPU.
+    dims: A list of integer describes how the tensor is partitioned.
+
+  Returns:
+    The same tensor with the xla_sharding attribute.
+  """
+  if dims is None:
+    return xla_sharding.replicate(tensor, assign_tuple_sharding=True)
+  elif np.prod(dims) == 1:
+    return xla_sharding.single_device(tensor, 0, assign_tuple_sharding=True)
+  else:
+    tile_assignment = np.arange(np.prod(dims)).reshape(dims)
+    return xla_sharding.tile(
+        tensor=tensor,
+        tile_assignment=tile_assignment,
+        assign_tuple_sharding=True)
+
+
+def tag_sharding_attribute_for_dequeued_tensors(dequeues, dims):
+  """Tags appropriate XLA sharding attribute to the dequeued tensors.
+
+  Args:
+    dequeues: A list of dequeued tensors on TPU.
+    dims: A list of integer describes how the tensor is partitioned.
+
+  Returns:
+    The same dequeues with appropriate xla_sharding attribute.
+  """
+  nest.assert_shallow_structure(dequeues, dims)
+  return nest.map_structure_up_to(
+      dequeues, _tag_sharding_attribute_for_dequeued_tensor, dequeues, dims)
+
+
+class InfeedQueue(object):
+  """A helper object to build a device infeed queue.
+
+  The InfeedQueue builds the host-side and device-side Ops to enqueue and
+  dequeue elements, respectively, and ensures that their types and
+  shapes match.
+  """
+
+  def __init__(self,
+               number_of_tuple_elements=None,
+               tuple_types=None,
+               tuple_shapes=None,
+               shard_dimensions=None,
+               number_of_partitions=None,
+               name=None):
+    """Creates a new InfeedQueue with the given configuration.
+
+    The configuration need not be fully specified at creation since it
+    can be modified subsequently by methods that set the values
+    explicitly or infer them from the shapes of inputs.
+
+    Args:
+      number_of_tuple_elements: the number of Tensors fed atomically through the
+        queue, must be present unless it can be inferred from other arguments.
+      tuple_types: if not None, a list of types of the elements of the queue.
+      tuple_shapes: if not None, a list of shapes of the elements of the queue.
+      shard_dimensions: if not None, a list of dimensions on which the
+        elements of the queue should be sharded during automatic
+        parallelization.
+      number_of_partitions: if > 1, the infeed dequeue shape will contain
+        the full shape that includes all partitions and add corresponding XLA
+        annotation on the infeed dequeue op. In this case, the infeed is still
+        data parallel that feeds per-core batch size to each core while the XLA
+        computation may be partitioned. As XLA requires infeed dequeue shape to
+        be per-replica shape, thus we need number_of_partitions here to
+        calculate the per-replica unpartitioned shape.
+      name: the name of the queue.
+
+    Raises:
+      ValueError: if number_of_tuple_elements <= 0; or
+        number_of_tuple_arguments, tuple_types, tuple_shapes, and
+        shard_dimensions are all None; or the length of tuple_types,
+        tuple_shapes, or shard_dimensions is not equal to
+        number_of_tuple_elements; or any element of shard_dimensions
+        can't be converted to a Dimension.
+      TypeError: if any element of tuple_types or tuple_shapes can't
+        be converted to a dtype or TensorShape, respectively.
+    """
+    self._frozen = False
+    self._generated_enqueue_ops = False
+    self._generated_dequeue_op = False
+    self._name = "InfeedQueue" if name is None else name
+    if number_of_partitions is None:
+      self._number_of_partitions = 1
+    else:
+      self._number_of_partitions = number_of_partitions
+    if number_of_tuple_elements is None:
+      if tuple_types is not None:
+        number_of_tuple_elements = len(tuple_types)
+      elif tuple_shapes is not None:
+        number_of_tuple_elements = len(tuple_shapes)
+      elif shard_dimensions is not None:
+        number_of_tuple_elements = len(shard_dimensions)
+      else:
+        raise ValueError(
+            "number of tuple elements cannot be inferred from InfeedQueue "
+            "constructor")
+    if number_of_tuple_elements <= 0:
+      raise ValueError(f"number_of_tuple_elements {number_of_tuple_elements} "
+                       "must be > 0")
+    # Make an empty sharding policy for each tuple element.
+    self._sharding_policies = [
+        tpu_sharding.ShardingPolicy() for _ in range(number_of_tuple_elements)
+    ]
+    if tuple_types is not None:
+      self.set_tuple_types(tuple_types)
+    else:
+      self._tuple_types = None
+    if tuple_shapes is not None:
+      self.set_tuple_shapes(tuple_shapes)
+    else:
+      self._tuple_shapes = None
+    if shard_dimensions is not None:
+      self.set_shard_dimensions(shard_dimensions)
+    self._validate()
+
+  def _validate(self):
+    """Checks that the configuration is self-consistent.
+
+    Raises:
+      ValueError: if the shapes and sharding policies don't match.
+    """
+    if self.tuple_shapes is not None:
+      for (policy, shape) in zip(self._sharding_policies, self._tuple_shapes):
+        # Raise an error if the policy is incompatible with the shape.
+        _ = policy.get_sharded_shape(shape)
+
+  @property
+  def number_of_tuple_elements(self):
+    """Returns the number of InfeedQueue tuple elements."""
+    return len(self._sharding_policies)
+
+  @property
+  def tuple_types(self):
+    """Returns the types of the InfeedQueue tuple elements."""
+    return self._tuple_types
+
+  def set_tuple_types(self, tuple_types):
+    """Sets the type of each element of the queue.
+
+    tuple_types must be a list of length
+    self.number_of_tuple_elements, and each element must be
+    convertible to a dtype.
+
+    Args:
+      tuple_types: the types of each queue element.
+
+    Raises:
+      ValueError: if tuple_types is not of length
+        self.number_of_tuple_elements.
+      TypeError: if an element of tuple_types cannot be converted to a
+        dtype.
+    """
+    if len(tuple_types) != self.number_of_tuple_elements:
+      raise ValueError(
+          f"tuple_types is {str(tuple_types)}, but must be a list of "
+          f"length {self.number_of_tuple_elements}"
+      )
+    if self._frozen:
+      for (frozen, updated) in zip(self._tuple_types, tuple_types):
+        if frozen != updated:
+          raise ValueError(
+              "Trying to update InfeedQueue with frozen configuration with an "
+              f"incompatible type. Frozen types are {str(self._tuple_types)}, "
+              f"updated types are {str(tuple_types)}")
+    else:
+      try:
+        self._tuple_types = [dtypes.as_dtype(t) for t in tuple_types]
+      except (TypeError) as e:
+        raise TypeError(
+            f"tuple_types is {str(tuple_types)}, but must be a list of "
+            f"elements each convertible to dtype: got error {str(e)}") from e
+
+  @property
+  def tuple_shapes(self):
+    """Returns the shapes of the InfeedQueue tuple elements."""
+    return self._tuple_shapes
+
+  def set_tuple_shapes(self, tuple_shapes):
+    """Sets the shape of each element of the queue.
+
+    tuple_shapes must be a list of length
+    self.number_of_tuple_elements, and each element must be
+    convertible to a TensorShape.
+
+    Args:
+      tuple_shapes: the shapes of each queue element.
+
+    Raises:
+      ValueError: if tuple_shapes is not of length
+        self.number_of_tuple_elements.
+      TypeError: if an element of tuple_shapes cannot be converted to
+        a TensorShape.
+    """
+    if len(tuple_shapes) != self.number_of_tuple_elements:
+      raise ValueError(
+          f"tuple_shapes is {str(tuple_shapes)}, but must be a list of "
+          f"length {self.number_of_tuple_elements}"
+      )
+    try:
+      tuple_shapes = [tensor_shape.as_shape(shape) for shape in tuple_shapes]
+    except (ValueError, TypeError) as e:
+      raise TypeError(
+          f"tuple_shapes is {str(tuple_shapes)}, but must be a list of "
+          "elements each convertible to TensorShape: got error "
+          f"{str(e)}") from e
+    if self._frozen:
+      for (frozen, updated) in zip(self._tuple_shapes, tuple_shapes):
+        if frozen != updated:
+          raise ValueError(
+              "Trying to update InfeedQueue with frozen configuration with an "
+              "incompatible shape. Frozen shapes are "
+              f"{str(self._tuple_shapes)}, updated shapes are "
+              f"{str(tuple_shapes)}")
+
+    else:
+      self._tuple_shapes = tuple_shapes
+    self._validate()
+
+  @property
+  def sharding_policies(self):
+    """Returns the sharding policies of the InfeedQueue tuple elements."""
+    return self._sharding_policies
+
+  @property
+  def shard_dimensions(self):
+    """Gets the shard dimension of each tuple element.
+
+    Returns:
+      A list of length number_of_tuple_elements, where each list entry
+      is the shard dimension of that tuple element or None if the
+      shard dimension has not been set.
+    """
+    # The number of shards is always the same for all the policies.
+    return [policy.shard_dimension for policy in self._sharding_policies]
+
+  def set_shard_dimensions(self, shard_dimensions):
+    """Sets the shard_dimension of each element of the queue.
+
+    shard_dimensions must be a list of length
+    self.number_of_tuple_elements, and each element must be
+    convertible to a Dimension compatible with self.tuple_shapes.
+
+    Args:
+      shard_dimensions: the dimensions of each queue element.
+
+    Raises:
+      ValueError: if shard_dimensions is not of length
+        self.number_of_tuple_elements; or an element of
+        shard_dimensions cannot be converted to a Dimension; or an
+        element of shard_dimensions is a Dimension that is out of
+        range for the corresponding tuple element shape.
+    """
+    if len(shard_dimensions) != self.number_of_tuple_elements:
+      raise ValueError(f"shard_dimensions is {str(shard_dimensions)}, but must "
+                       f"be a list of length {self.number_of_tuple_elements}")
+    for (policy, dimension) in zip(self._sharding_policies, shard_dimensions):
+      policy.set_shard_dimension(dimension)
+    self._validate()
+
+  @property
+  def number_of_shards(self):
+    """Gets the number of shards to use for the InfeedQueue.
+
+    Returns:
+      Number of shards or None if the number of shards has not been set.
+    """
+    # The number of shards is always the same for all the policies.
+    return self._sharding_policies[0].number_of_shards
+
+  def set_number_of_shards(self, number_of_shards):
+    """Sets the number of shards to use for the InfeedQueue.
+
+    Args:
+      number_of_shards: number of ways to shard the InfeedQueue.
+
+    Raises:
+      ValueError: if number_of_shards is not > 0; or the policies have
+        been frozen and number_of_shards was already set to something
+        else.
+    """
+    for policy in self._sharding_policies:
+      policy.set_number_of_shards(number_of_shards)
+      policy.set_number_of_partitions(self._number_of_partitions)
+    self._validate()
+
+  def set_configuration_from_input_tensors(self, input_tensors):
+    """Sets the shapes and types of the queue tuple elements.
+
+    input_tensors is a list of Tensors whose types and shapes are used
+    to set the queue configuration.
+
+    Args:
+      input_tensors: list of Tensors of the same types and shapes as
+        the desired queue Tuple.
+
+    Raises:
+      ValueError: if input_tensors is not a list of length
+        self.number_of_tuple_elements
+    """
+    if len(input_tensors) != self.number_of_tuple_elements:
+      raise ValueError(f"input_tensors is {str(input_tensors)}, but should be "
+                       f"a list of {self.number_of_tuple_elements} Tensors")
+    self.set_tuple_shapes([t.shape for t in input_tensors])
+    self.set_tuple_types([t.dtype for t in input_tensors])
+
+  def set_configuration_from_sharded_input_tensors(self, input_tensors):
+    """Sets the shapes and types of the queue tuple elements.
+
+    input_tensors is a list of lists of Tensors whose types and shapes are used
+    to set the queue configuration. The length of the outer list is the number
+    of shards required, and each inner list is the tuple of Tensors to use to
+    determine the types and shapes of the corresponding shard. This method
+    depends on the shard dimension, and calling it freezes the shard policy.
+
+    Args:
+      input_tensors: list of lists of Tensors. The outer list length corresponds
+        to the desired number of shards, and each inner list is the size
+        and shape of the desired configuration of the corresponding shard.
+
+    Raises:
+      ValueError: if any inner list is not a list of length
+        self.number_of_tuple_elements; or the inner lists do not combine to
+        form a consistent unsharded shape.
+      TypeError: if the types of the Tensors in the inner lists do not match.
+    """
+    if not self._frozen:
+      # Unset the tuple shapes in case the configuration becomes
+      # transiently inconsistent.
+      self._tuple_shapes = None
+    number_of_shards = len(input_tensors)
+    self.set_number_of_shards(number_of_shards)
+    for t in input_tensors:
+      if len(t) != self.number_of_tuple_elements:
+        raise ValueError(
+            f"input_tensors is {str(input_tensors)} but must be a list of "
+            "lists, where each inner list has length "
+            f"number_of_tuple_elements={self.number_of_tuple_elements}")
+    # Transpose the inputs to make a list of shard shapes for each tuple
+    # element.
+    sharded_shapes = [[t[i].shape
+                       for t in input_tensors]
+                      for i in range(self.number_of_tuple_elements)]
+    # For each tuple, get the unsharded shape using that tuple's policy.
+    unsharded_shapes = [
+        policy.get_unsharded_shape(s)
+        for (policy, s) in zip(self._sharding_policies, sharded_shapes)
+    ]
+    self.set_tuple_shapes(unsharded_shapes)
+    for i in range(1, self.number_of_shards):
+      for (t1, t2) in zip(input_tensors[0], input_tensors[i]):
+        if t1.dtype != t2.dtype:
+          raise TypeError(
+              "types of the tuple elements of input_tensors "
+              f"{str(input_tensors)} are not consistent")
+    self.set_tuple_types([t.dtype for t in input_tensors[0]])
+
+  def freeze(self):
+    """Freezes the InfeedQueue so it can no longer be modified.
+
+    The configuration is implicitly frozen before any host-side or
+    device-side Ops are generated. The configuration cannot be frozen
+    until the types and shapes of the tuple elements have been set.
+
+    Raises:
+      ValueError: if the types or shapes of the tuple elements have not been
+      set.
+    """
+    self._frozen = True
+    if self._tuple_types is None:
+      raise ValueError(
+          "Can't freeze an InfeedQueue without setting all tuple types.")
+    if self._tuple_shapes is None:
+      raise ValueError(
+          "Can't freeze an InfeedQueue without setting all tuple shapes.")
+    for shape in self._tuple_shapes:
+      if shape.dims is None:
+        raise ValueError(
+            "Can't freeze an InfeedQueue without setting all tuple shapes.")
+    for policy in self._sharding_policies:
+      policy.freeze()
+    self._validate()
+
+  def generate_dequeue_op(self, tpu_device=0):
+    """Generates the device-side Op to dequeue a tuple from the queue.
+
+    Implicitly freezes the queue configuration if it is not already
+    frozen, which will raise errors if the shapes and types have not
+    been fully specified.
+
+    Args:
+      tpu_device: The TPU device ordinal where the infeed instruction should be
+        placed. If None, no explicit placement will be performed, and it is up
+        to the user to call this API from within a proper TPU device scope.
+        The XLA code will fail if the TPU dequeue instruction is not bound to
+        any device.
+
+    Returns:
+      A list of Outputs corresponding to a shard of infeed dequeued
+      into XLA, suitable for use within a replicated block.
+
+    Raises:
+      ValueError: if the types or shapes of the tuple elements have not been
+      set; or if a dequeue op has already been generated.
+    """
+    self.freeze()
+    if self._generated_dequeue_op and not ops.inside_function():
+      raise ValueError("Can't generate two dequeue Ops from the same queue")
+    self._generated_dequeue_op = True
+    full_name = "%s/dequeue" % self._name
+    sharded_shapes = [
+        policy.get_unpartitioned_shape(policy.get_sharded_shape(shape))
+        for (shape, policy) in zip(self._tuple_shapes, self._sharding_policies)
+    ]
+    if tpu_device is not None:
+      with ops.device(tpu_name_util.core(tpu_device)):
+        dequeue_op = tpu_ops.infeed_dequeue_tuple(
+            dtypes=self._tuple_types, shapes=sharded_shapes, name=full_name)
+    else:
+      dequeue_op = tpu_ops.infeed_dequeue_tuple(
+          dtypes=self._tuple_types, shapes=sharded_shapes, name=full_name)
+    if self._number_of_partitions <= 1:
+      return dequeue_op
+    partitions = [
+        policy.get_unpartitioned_shape([1] * shape.ndims).as_list()
+        for (shape, policy) in zip(self._tuple_shapes, self._sharding_policies)
+    ]
+    return tag_sharding_attribute_for_dequeued_tensors(dequeue_op, partitions)
+
+  def _generate_enqueue_op(self,
+                           inputs,
+                           name_prefix,
+                           index,
+                           device=None,
+                           tpu_ordinal=-1):
+    """Generate a host-side Op to enqueue a tuple to the queue.
+
+    If device is None the inputs are all required to have the same
+    device specification, and the enqueue Op is colocated with
+    inputs[0]. Otherwise the enqueue Op is placed on 'device'.
+
+    Args:
+      inputs: a list of Tensors with the types and shapes of the tuple elements.
+      name_prefix: the base name for the Op.
+      index: the shard index, used to uniquify the Op name.
+      device: device to place the Op on, or None if it should be
+        colocated with the inputs.
+      tpu_ordinal: ordinal of the TPU device on the host to use for
+      infeed if device is a CPU device. Should be set to -1 if device
+      is a TPU device.
+
+    Returns:
+      An Op corresponding to a shard of infeed enqueued at the host,
+      suitable for use within a replicated block.
+
+    Raises:
+      ValueError: if device is None and inputs do not all have the
+        same device specification.
+    """
+    full_name = "%s/%d" % (name_prefix, index)
+    shapes = [t.shape for t in inputs]
+    if device is None:
+      devices = [t.device for t in inputs]
+      for i in range(1, self.number_of_tuple_elements):
+        if devices[0] != devices[i]:
+          raise ValueError(
+              f"input devices for shard {index} are {str(devices)}, but should "
+              "all be the same")
+      with ops.colocate_with(inputs[0]):
+        return tpu_ops.infeed_enqueue_tuple(
+            inputs=inputs,
+            shapes=shapes,
+            name=full_name,
+            device_ordinal=tpu_ordinal)
+    else:
+      with ops.device(device):
+        return tpu_ops.infeed_enqueue_tuple(
+            inputs=inputs,
+            shapes=shapes,
+            name=full_name,
+            device_ordinal=tpu_ordinal)
+
+  def generate_enqueue_ops(self,
+                           sharded_inputs,
+                           tpu_ordinal_function=None,
+                           placement_function=None):
+    """Generates the host-side Ops to enqueue the shards of a tuple.
+
+    sharded_inputs is a list, one for each shard, of lists of
+    Tensors. sharded_inputs[i] is the tuple of Tensors to use to feed
+    shard i of the queue. Returns the host-side Ops that must be run to
+    enqueue the sharded tuple. The Op for shard i is colocated with the inputs
+    for shard i.
+
+    Implicitly freezes the queue configuration if it is not already
+    frozen. If the configuration has already been frozen, and is not
+    compatible with the types and shapes of sharded_inputs, an error
+    will be raised.
+
+    Args:
+      sharded_inputs: a list of lists of Tensors. The length of the outer list
+        determines the number of shards. Each inner list indicates the types
+        and shapes of the tuples in the corresponding shard.
+      tpu_ordinal_function: if not None, a function that takes the
+        shard index as input and returns the ordinal of the TPU device
+        the shard's infeed should be placed on. tpu_ordinal_function must be
+        set if the inputs are placed on CPU devices.
+      placement_function: if not None, a function that takes the shard index as
+        input and returns the host device where the enqueue op should be placed
+        on.
+
+    Returns:
+      A list of host-side Ops, one for each shard, that when executed together
+      will enqueue a full-size element of infeed.
+
+    Raises:
+      ValueError: if the queue configuration has previously been frozen and the
+        shapes of the elements of sharded_inputs are not compatible with the
+        frozen configuration; or if the shapes of the elements of sharded_inputs
+        don't form a consistent unsharded tuple; or if the elements of a tuple
+        have different device constraints.
+      TypeError: if the queue configuration has previously been frozen and the
+        types of the elements of sharded_inputs are not compatible with the
+        frozen configuration; or if the types of the elements of sharded_inputs
+        don't form a consistent unsharded tuple.
+    """
+    self.set_configuration_from_sharded_input_tensors(sharded_inputs)
+    self.freeze()
+    if self._generated_enqueue_ops and not ops.inside_function():
+      raise ValueError("Can't generate two enqueue Ops from the same queue")
+    self._generated_enqueue_ops = True
+    if tpu_ordinal_function is None:
+      tpu_ordinal_function = lambda index: -1
+    name_prefix = "%s/enqueue" % self._name
+    return [
+        self._generate_enqueue_op(
+            shard,
+            name_prefix,
+            index,
+            tpu_ordinal=tpu_ordinal_function(index),
+            device=placement_function(index) if placement_function else None)
+        for (shard, index) in zip(sharded_inputs, range(self.number_of_shards))
+    ]
+
+  # TODO(misard) Generalize this to the case of systems that don't
+  # have 8 devices per host, and figure out what to do with
+  # model-parallelism.
+  def _default_placement_function(self, index):
+    return "/task:%d/device:CPU:0" % (index / 8)
+
+  def _default_ordinal_function(self, index):
+    return index % 8
+
+  # TODO(b/36470756) remove this from tutorials once we have a better story
+  # for automatic placement of input pipelines.
+  def split_inputs_and_generate_enqueue_ops(self,
+                                            inputs,
+                                            device_assignment=None,
+                                            placement_function=None,
+                                            tpu_ordinal_function=None):
+    """POORLY-PERFORMING ON MULTI-HOST SYSTEMS.
+
+    Generates the host-side Ops to enqueue a tuple.
+
+    This method performs poorly because it takes an entire input on a single
+    host, splits it, and distributes it to all of the cores. It is present only
+    to simplify tutorial examples.
+
+    inputs is a list of Tensors to use to feed the queue. Each input is split
+    into self.number_of_shards shards. Returns an Op for each shard to enqueue
+    the shard. The Op for shard i is placed on device placement_function(i).
+
+    Implicitly freezes the queue configuration if it is not already
+    frozen. If the configuration has already been frozen, and is not
+    compatible with the types and shapes of inputs, an error
+    will be raised.
+
+    Args:
+      inputs: a list of Tensors which indicates the types and shapes of the
+        queue tuple.
+     device_assignment: if not `None`, a TPU `DeviceAssignment`. If
+        device_assignment is not `None`, but `placement_function` and
+        `ordinal_function` are None, then `device_assignment` will be used to
+        place infeeds on the first k TPU shards, where k is the number of shards
+        in the queue. If all three are `None`, then default placement and
+        ordinal functions are used.
+      placement_function: if not None, a function that takes the shard
+        index as input and returns a device string indicating which
+        device the shard's infeed should be placed on. If placement_function
+        and tpu_ordinal_function are None, inputs are sharded round-robin
+        across the devices in the system.
+      tpu_ordinal_function: if not None, a function that takes the
+        shard index as input and returns the ordinal of the TPU device
+        the shard's infeed should be placed on. If placement_function
+        and tpu_ordinal_function are None, inputs are sharded round-robin
+        across the devices in the system.
+
+    Returns:
+      A list of host-side Ops, one for each shard, that when executed together
+      will enqueue a full-size element of infeed.
+
+    Raises:
+      ValueError: if the queue configuration has previously been frozen and the
+        shapes of the elements of inputs are not compatible with the frozen
+        configuration.
+      TypeError: if the queue configuration has previously been frozen and the
+        types of the elements of inputs are not compatible with the frozen
+        configuration.
+    """
+    if device_assignment is None:
+      if placement_function is None:
+        placement_function = self._default_placement_function
+      if tpu_ordinal_function is None:
+        tpu_ordinal_function = self._default_ordinal_function
+    else:
+
+      def _placement_function_from_map(index):
+        return device_assignment.host_device(replica=index)
+
+      def _ordinal_function_from_map(index):
+        return device_assignment.tpu_ordinal(replica=index)
+
+      if placement_function is None:
+        placement_function = _placement_function_from_map
+      if tpu_ordinal_function is None:
+        tpu_ordinal_function = _ordinal_function_from_map
+    self.set_configuration_from_input_tensors(inputs)
+    self.freeze()
+    if self._generated_enqueue_ops and not ops.inside_function():
+      raise ValueError("Can't generate two enqueue Ops from the same queue")
+    self._generated_enqueue_ops = True
+    split_name_prefix = "%s/split" % self._name
+    if self.number_of_shards == 1:
+      transposed_sharded_inputs = [[inp] for inp in inputs]
+    else:
+
+      def split_fn(inp, num_shards, axis, name):
+        with ops.colocate_with(inp):
+          return array_ops.split(inp, num_shards, axis=axis, name=name)
+
+      transposed_sharded_inputs = [
+          split_fn(
+              inp,
+              self.number_of_shards,
+              axis=policy.shard_dimension,
+              name="%s/%d" % (split_name_prefix, index))
+          for (inp, policy, index) in zip(inputs, self._sharding_policies,
+                                          range(self.number_of_tuple_elements))
+      ]
+    sharded_inputs = [[shard[i]
+                       for shard in transposed_sharded_inputs]
+                      for i in range(self.number_of_shards)]
+    name_prefix = "%s/enqueue" % self._name
+    return [
+        self._generate_enqueue_op(
+            shard,
+            name_prefix,
+            index,
+            device=placement_function(index),
+            tpu_ordinal=tpu_ordinal_function(index))
+        for (shard, index) in zip(sharded_inputs, range(self.number_of_shards))
+    ]
+
+
+class _PartitionedInfeedQueue(InfeedQueue):
+  """A helper object to build a device infeed queue with input partition.
+
+  Args:
+    number_of_tuple_elements: the number of Tensors fed atomically through the
+      queue, must be present unless it can be inferred from other arguments.
+    device_assignment: A TPU `DeviceAssignment` which is used to place all the
+      partitions to different TPU infeed queues.
+    host_id: The id of the host machine.
+    input_partition_dims: A nested list/tuple of integers. Each inner
+      list/tuple describes how to partition the corresponding input tensor.
+    tuple_types: If not None, a list of types of the elements of the queue.
+    tuple_shapes: If not None, a list of shapes of the elements of the queue.
+    name: The name of the queue.
+  """
+
+  def __init__(self,
+               number_of_tuple_elements,
+               device_assignment,
+               host_id,
+               input_partition_dims=None,
+               tuple_types=None,
+               tuple_shapes=None,
+               name=None):
+    super(_PartitionedInfeedQueue, self).__init__(
+        number_of_tuple_elements=number_of_tuple_elements,
+        tuple_types=tuple_types,
+        tuple_shapes=None,
+        shard_dimensions=None,
+        name="PartitionedInfeedQueue" if name is None else name)
+    self._input_partition_dims = input_partition_dims
+    self._host_id = host_id
+    self._device_assignment = device_assignment
+
+  def generate_dequeue_op(self, tpu_device=0):
+    """Generate TPU dequeue ops.
+
+    Args:
+      tpu_device: The TPU device ordinal where the infeed instruction should be
+        placed.
+
+    Returns:
+      A list of Outputs corresponding to a partition of infeed dequeued
+      into XLA, suitable for use within a replicated block.
+
+    Raises:
+      ValueError: if the types or shapes of the tuple elements have not been
+      set; or if a dequeue op has already been generated.
+    """
+    self.freeze()
+    if self._generated_dequeue_op and not ops.inside_function():
+      raise ValueError("Can't generate two dequeue Ops from the same queue")
+    self._generated_dequeue_op = True
+    full_name = "%s/dequeue" % self._name
+    sharded_shapes = [
+        policy.get_sharded_shape(shape)
+        for (shape, policy) in zip(self._tuple_shapes, self._sharding_policies)
+    ]
+    with ops.device(tpu_name_util.core(tpu_device)):
+      values = tpu_ops.infeed_dequeue_tuple(
+          dtypes=self._tuple_types, shapes=sharded_shapes, name=full_name)
+    return tag_sharding_attribute_for_dequeued_tensors(
+        values, self._input_partition_dims)
+
+  def generate_enqueue_ops(self, sharded_inputs):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
+    """Generates the host-side Ops to enqueue the partitioned inputs.
+
+    sharded_inputs is a list, one for each replica, of lists of
+    Tensors. sharded_inputs[i] is the tuple of Tensors to use to feed
+    replica i.
+    sharded_inputs[i][j] is partitioned by self._input_partition_dims[j].
+
+    For example, if sharded_inputs[i][j] is a 2-D Tensor:
+    [[A, B, C, D],
+     [E ,F, G, H]]
+    self._input_partition_dims[j] is [2, 4].
+
+    sharded_inputs[i][j] will be partitioned and flattened into:
+    [A, B, C, D, E, F, G, H] and fed into the logical core ids:
+    [0, 1, 2, 3, 4, 5, 6, 7] respectively.
+
+    Args:
+      sharded_inputs: a list of lists of Tensors. The length of the
+        outer list determines the number of shards. Each inner list indicates
+        the types and shapes of the tuples in the corresponding shard.
+
+    Returns:
+      A list of host-side Ops, one for each shard, that when executed together
+      will enqueue a full-size element of infeed.
+
+    Raises:
+      ValueError: if the queue configuration has previously been frozen and the
+        shapes of the elements of sharded_inputs are not compatible with the
+        frozen configuration; or if the shapes of the elements of sharded_inputs
+        don't form a consistent unsharded tuple; or if the elements of a tuple
+        have different device constraints; or if the partition dims are invalid.
+      TypeError: if the queue configuration has previously been frozen and the
+        types of the elements of sharded_inputs are not compatible with the
+        frozen configuration; or if the types of the elements of sharded_inputs
+        don't form a consistent unsharded tuple.
+    """
+    self.set_configuration_from_sharded_input_tensors(sharded_inputs)
+    number_of_replicas = len(sharded_inputs)
+    number_of_tuple_elements = len(sharded_inputs[0])
+
+    assert len(self._input_partition_dims) == number_of_tuple_elements
+    enqueue_ops = []
+
+    for replica_index in range(number_of_replicas):
+      flattened_inputs = sharded_inputs[replica_index]
+      inputs_part_dims_flat = nest.flatten_up_to(flattened_inputs,
+                                                 self._input_partition_dims)
+      inputs_parted_iters = [
+          iter(self._check_dims_and_partition_or_replicate_on_host(x, dims))
+          for x, dims in zip(sharded_inputs[replica_index],
+                             inputs_part_dims_flat)
+      ]
+
+      # Find the replica_id of the host's logical core 0.
+      # The self._host_id is guaranteed to contain the logical core 0,
+      # even when num_cores_per_replica > num_cores_per_host -- the function
+      # caller makes sure that this host_id will must be receiving data (calls
+      # input_fn).
+      replica_id = self._device_assignment.lookup_replicas(
+          task_id=self._host_id, logical_core=0)[replica_index]
+      for logical_core in range(self._device_assignment.num_cores_per_replica):
+        # Places different partitions to different logic cores.
+        # Since there can be multiple hosts per replica, we need to find
+        # the actual host (device) of this logical core.
+        device = self._device_assignment.host_device(
+            replica=replica_id, logical_core=logical_core)
+
+        with ops.device(device):
+          ordinal = self._device_assignment.tpu_ordinal(
+              replica=replica_id, logical_core=logical_core)
+          infeed_inputs = []
+          for it in inputs_parted_iters:
+            input_for_device = next(it, None)
+            if input_for_device is not None:
+              infeed_inputs.append(input_for_device)
+
+          if infeed_inputs:
+            enqueue_ops.append(
+                tpu_ops.infeed_enqueue_tuple(
+                    inputs=infeed_inputs,
+                    shapes=[x.shape for x in infeed_inputs],
+                    name="enqueue/replica_{0}/input_{1}".format(
+                        replica_index, logical_core),
+                    device_ordinal=ordinal))
+    return enqueue_ops
+
+  def _check_input_partition_dims(self, tensor, dims):
+    """Checks that input partition dims are valid for the `Tensor`.
+
+    Args:
+      tensor: Input tensor for partitioning.
+      dims: A list of integer describes how to partition the input tensor.
+
+    Raises:
+      ValueError: If the tensor can't be partitioned by dims or the
+        num_cores_per_replica doesn't match the number of
+        partitions(dims.prod()).
+    """
+    # No partitioning specified, so don't perform further checks.
+    if dims is None:
+      return
+
+    dims = np.array(dims)
+
+    if (dims < 1).any():
+      raise ValueError("All input partition dims must be >= 1.")
+
+    # No partitioning, so don't perform further checks.
+    if dims.prod() == 1:
+      return
+
+    if dims.prod() != self._device_assignment.num_cores_per_replica:
+      raise ValueError(
+          "The product of each input partition dim should equal to "
+          "num_cores_per_replica. (dim = {}, num_cores_per_replica "
+          "= {})".format(dims, self._device_assignment.num_cores_per_replica))
+    if dims.shape[0] != tensor.shape.ndims:
+      raise ValueError(
+          "Input partition dims must have the same number of dimensions "
+          "as the `Tensor` to be partitioned. (tensor shape = {}, input "
+          "partition dims = {}).".format(tensor.shape.as_list(), dims))
+
+    tensor.shape.assert_is_fully_defined()
+
+  def _check_dims_and_partition_or_replicate_on_host(self, tensor, dims):
+    """Checks dims and partitions or replicates the input tensor.
+
+      The ops inside this function are placed on the host side.
+
+    Args:
+      tensor: The input tensor which will be partitioned or replicated.
+      dims: A list of integer describes how to partition the input tensor.
+
+    Returns:
+      An iterator of `Tensor`s or a list of partitioned tensors.
+    """
+    self._check_input_partition_dims(tensor, dims)
+    return partition_or_replicate_on_host(tensor, dims)

tensorflow/python/tpu/tpu_function.py

@@ -0,0 +1,64 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Helper library for functions used during TPU compilation."""
+
+import contextlib
+import threading
+
+
+class TpuContext(threading.local):
+  """A context object holding state about the TPU computation being built."""
+
+  def __init__(self):
+    """Creates a new TpuContext."""
+    self._number_of_shards = None
+
+  @property
+  def number_of_shards(self):
+    return self._number_of_shards
+
+  def set_number_of_shards(self, number_of_shards):
+    self._number_of_shards = number_of_shards
+
+
+# The Tpu context holds the number of shards when a sharded computation is
+# being built, or None if no computation is being built.
+_current_tpu_context = TpuContext()
+
+
+@contextlib.contextmanager
+def tpu_shard_context(number_of_shards):
+  """A context manager setting current number of shards."""
+  if _current_tpu_context.number_of_shards is not None:
+    raise NotImplementedError("tpu_shard_context cannot be nested")
+  try:
+    _current_tpu_context.set_number_of_shards(number_of_shards)
+    yield
+  finally:
+    _current_tpu_context.set_number_of_shards(None)
+
+
+def get_tpu_context():
+  return _current_tpu_context
+
+
+# Decorator function for tpu computation func that was passed to tpu.rewrite()
+# if there is an embedded training loop in this func, trace tools will generate
+# step markers for each iteration.
+def on_device_training_loop(func):
+  # Value for this attribute is from xla.DebugOptions.StepMarkerLocation.
+  setattr(func, "step_marker_location", "STEP_MARK_AT_TOP_LEVEL_WHILE_LOOP")
+  return func

tensorflow/python/tpu/tpu_hardware_feature.py

@@ -0,0 +1,81 @@
+# 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.
+# ==============================================================================
+"""TPU hardware feature info."""
+import enum
+from tensorflow.core.protobuf.tpu import topology_pb2
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("tpu.experimental.HardwareFeature")
+class HardwareFeature(object):
+  """class holds all the feature info about the TPU."""
+
+  def __init__(self, tpu_hardware_feature_proto):
+    """Store TPU hardware feature info.
+
+    Args:
+      tpu_hardware_feature_proto: protobuf which describe the tpu hardware
+        feature.
+    """
+    self.tpu_hardware_feature_proto = tpu_hardware_feature_proto
+
+  class EmbeddingFeature(enum.Enum):
+    """Embedding feature flag strings.
+
+    UNSUPPORTED: No embedding lookup accelerator available on the tpu.
+    V1: Embedding lookup accelerator V1. The embedding lookup operation can only
+        be placed at the beginning of computation. Only one instance of
+        embedding
+        lookup layer is allowed.
+    V2: Embedding lookup accelerator V2. The embedding lookup operation can be
+        placed anywhere of the computation. Multiple instances of embedding
+        lookup layer is allowed.
+    """
+    UNSUPPORTED = "UNSUPPORTED"
+    V1 = "V1"
+    V2 = "V2"
+
+  @classmethod
+  def _embedding_feature_proto_to_string(cls, embedding_feature_proto):
+    """Convert the embedding feature proto to enum string."""
+    embedding_feature_proto_to_string_map = {
+        topology_pb2.TPUHardwareFeature.EmbeddingFeature.UNSUPPORTED:
+            HardwareFeature.EmbeddingFeature.UNSUPPORTED,
+        topology_pb2.TPUHardwareFeature.EmbeddingFeature.V1:
+            HardwareFeature.EmbeddingFeature.V1,
+        topology_pb2.TPUHardwareFeature.EmbeddingFeature.V2:
+            HardwareFeature.EmbeddingFeature.V2
+    }
+    return embedding_feature_proto_to_string_map.get(
+        embedding_feature_proto, HardwareFeature.EmbeddingFeature.UNSUPPORTED)
+
+  @property
+  def embedding_feature(self):
+    """TPU embedding feature.
+
+    Returns:
+      An EmbeddingFeature enum.
+    """
+    return HardwareFeature._embedding_feature_proto_to_string(
+        self.tpu_hardware_feature_proto.embedding_feature)
+
+  @property
+  def num_embedding_devices_per_chip(self):
+    """Number of embedding accelerator devices per chip.
+
+    Returns:
+      Number of embedding devices per chip.
+    """
+    return self.tpu_hardware_feature_proto.num_embedding_devices_per_chip

tensorflow/python/tpu/tpu_infeed_test.py

@@ -0,0 +1,126 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Tests for TPU InfeedQueue methods."""
+
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.platform import test
+from tensorflow.python.tpu import tpu_feed
+
+
+class InfeedTest(test.TestCase):
+
+  def testConstructor(self):
+    """Tests that the constructor can be called with different arguments."""
+    i = tpu_feed.InfeedQueue(number_of_tuple_elements=2)
+    self.assertEqual(i.number_of_tuple_elements, 2)
+    self.assertEqual(i.tuple_types, None)
+    self.assertEqual(i.tuple_shapes, None)
+    self.assertEqual(i.number_of_shards, None)
+    i = tpu_feed.InfeedQueue(
+        tuple_types=[dtypes.float32, dtypes.int32, dtypes.int32])
+    self.assertEqual(i.number_of_tuple_elements, 3)
+    self.assertEqual(i.tuple_types,
+                     [dtypes.float32, dtypes.int32, dtypes.int32])
+    self.assertEqual(i.tuple_shapes, None)
+    self.assertEqual(i.number_of_shards, None)
+    i = tpu_feed.InfeedQueue(tuple_shapes=[[1], [2, 3]])
+    self.assertEqual(i.number_of_tuple_elements, 2)
+    self.assertEqual(i.tuple_types, None)
+    self.assertEqual(i.tuple_shapes, [[1], [2, 3]])
+    self.assertEqual(i.number_of_shards, None)
+    i = tpu_feed.InfeedQueue(shard_dimensions=[1, 0, 7])
+    self.assertEqual(i.number_of_tuple_elements, 3)
+    self.assertEqual(i.tuple_types, None)
+    self.assertEqual(i.tuple_shapes, None)
+    self.assertEqual([p.shard_dimension
+                      for p in i.sharding_policies], [1, 0, 7])
+    with self.assertRaises(ValueError):
+      i = tpu_feed.InfeedQueue()
+    with self.assertRaises(ValueError):
+      i = tpu_feed.InfeedQueue(
+          number_of_tuple_elements=2, tuple_types=[dtypes.float32])
+    with self.assertRaises(ValueError):
+      i = tpu_feed.InfeedQueue(number_of_tuple_elements=2, tuple_shapes=[[1]])
+    with self.assertRaises(ValueError):
+      i = tpu_feed.InfeedQueue(number_of_tuple_elements=2, shard_dimensions=[1])
+    with self.assertRaises(ValueError):
+      i = tpu_feed.InfeedQueue(tuple_shapes=[[1], [2, 3]], shard_dimensions=[1])
+
+  def testModification(self):
+    """Tests modification of the queue post-construction."""
+    i = tpu_feed.InfeedQueue(number_of_tuple_elements=2)
+    i.set_tuple_types([dtypes.float32, dtypes.int32])
+    self.assertEqual(i.tuple_types, [dtypes.float32, dtypes.int32])
+    i.set_tuple_types([dtypes.float32, dtypes.float32])
+    self.assertEqual(i.tuple_types, [dtypes.float32, dtypes.float32])
+    with self.assertRaises(ValueError):
+      i.set_tuple_types([dtypes.float32])
+    i.set_tuple_shapes([[1], [2, 3]])
+    self.assertEqual(i.tuple_shapes, [[1], [2, 3]])
+    i.set_tuple_shapes([[1, 2], [3, 4]])
+    self.assertEqual(i.tuple_shapes, [[1, 2], [3, 4]])
+    with self.assertRaises(ValueError):
+      i.set_tuple_shapes([[1, 2]])
+    i.set_number_of_shards(2)
+    self.assertEqual(i.number_of_shards, 2)
+    i.set_number_of_shards(3)
+    self.assertEqual(i.number_of_shards, 3)
+    t1 = constant_op.constant(1, dtypes.int32, shape=[6])
+    t2 = constant_op.constant(2.0, dtypes.float32, shape=[3, 18])
+    i.set_configuration_from_input_tensors([t1, t2])
+    self.assertEqual(i.tuple_shapes, [[6], [3, 18]])
+    self.assertEqual(i.tuple_types, [dtypes.int32, dtypes.float32])
+    i.set_configuration_from_sharded_input_tensors([[t2, t1], [t2, t1]])
+    self.assertEqual(i.number_of_shards, 2)
+    self.assertEqual(i.tuple_shapes, [[6, 18], [12]])
+    self.assertEqual(i.tuple_types, [dtypes.float32, dtypes.int32])
+    i.set_shard_dimensions([1, 0])
+    i.set_number_of_shards(3)
+    with self.assertRaises(ValueError):
+      i.set_number_of_shards(4)
+
+  def testFreezing(self):
+    """Tests freezing the queue."""
+    i = tpu_feed.InfeedQueue(number_of_tuple_elements=2)
+    t1 = constant_op.constant(1, dtypes.int32, shape=[2])
+    t2 = constant_op.constant(2.0, dtypes.float32, shape=[2, 4])
+    i.set_configuration_from_sharded_input_tensors([[t2, t1], [t2, t1]])
+    self.assertEqual(i.number_of_shards, 2)
+    self.assertEqual(i.tuple_shapes, [[4, 4], [4]])
+    self.assertEqual(i.tuple_types, [dtypes.float32, dtypes.int32])
+    self.assertEqual(i.shard_dimensions, [0, 0])
+    i.freeze()
+    i.set_number_of_shards(2)
+    i.set_tuple_shapes([[4, 4], [4]])
+    i.set_tuple_types([dtypes.float32, dtypes.int32])
+    i.set_shard_dimensions([0, 0])
+    with self.assertRaises(ValueError):
+      i.set_number_of_shards(1)
+    with self.assertRaises(ValueError):
+      i.set_tuple_shapes([[8, 8], [8]])
+    with self.assertRaises(ValueError):
+      i.set_tuple_types([dtypes.int32, dtypes.float32])
+    with self.assertRaises(ValueError):
+      i.set_shard_dimensions([1, 0])
+    self.assertEqual(i.number_of_shards, 2)
+    self.assertEqual(i.tuple_shapes, [[4, 4], [4]])
+    self.assertEqual(i.tuple_types, [dtypes.float32, dtypes.int32])
+    self.assertEqual(i.shard_dimensions, [0, 0])
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/tpu/tpu_name_util.py

@@ -0,0 +1,31 @@
+# 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.
+# ======================================
+"""Helper functions for TPU device names."""
+
+from typing import Text
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["tpu.core"])
+def core(num: int) -> Text:
+  """Returns the device name for a core in a replicated TPU computation.
+
+  Args:
+    num: the virtual core number within each replica to which operators should
+    be assigned.
+  Returns:
+    A device name, suitable for passing to `tf.device()`.
+  """
+  return "device:TPU_REPLICATED_CORE:{}".format(num)

tensorflow/python/tpu/tpu_optimizer.py

@@ -0,0 +1,222 @@
+# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+"""Optimizer that implements cross-shard gradient reduction for TPU."""
+
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops.losses import losses
+from tensorflow.python.platform import tf_logging as logging
+from tensorflow.python.tpu import tpu_function
+from tensorflow.python.tpu.ops import tpu_ops
+from tensorflow.python.training import optimizer
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["tpu.CrossShardOptimizer"])
+class CrossShardOptimizer(optimizer.Optimizer):
+  """An optimizer that averages gradients across TPU shards."""
+
+  def __init__(self,
+               opt,
+               reduction=losses.Reduction.MEAN,
+               name="CrossShardOptimizer",
+               group_assignment=None):
+    """Construct a new cross-shard optimizer.
+
+    Args:
+      opt: An existing `Optimizer` to encapsulate.
+      reduction: The reduction to apply to the shard losses.
+      name: Optional name prefix for the operations created when applying
+        gradients. Defaults to "CrossShardOptimizer".
+      group_assignment: Optional 2d int32 lists with shape
+        [num_groups, num_replicas_per_group] which describles how to apply
+        optimizer to subgroups.
+
+    Raises:
+      ValueError: If reduction is not a valid cross-shard reduction.
+    """
+    accepted_reductions = (losses.Reduction.SUM, losses.Reduction.MEAN)
+    if reduction not in accepted_reductions:
+      raise ValueError(
+          f"Argument `reduction` should be one of {accepted_reductions}. "
+          f"Received: {reduction}")
+    if not isinstance(opt, optimizer.Optimizer):
+      raise TypeError(
+          "CrossShardOptimizer only works with tf.training.Optimizer and not "
+          f"Keras Optimizer. Received: {opt}. "
+          "If you are using TPUStrategy, "
+          "Keras Optimizer will sum gradients across replicas."
+          "If you want to average your gradients, rescale your loss with: "
+          "`loss /= global_batch_size`")
+
+    super(CrossShardOptimizer, self).__init__(False, name)
+    self._opt = opt
+    self._reduction = reduction
+    self._group_assignment = group_assignment
+
+  def _verify_and_get_subgroup_size(self, group_assignment, num_shards):
+    """Verify group_assignment and get the subgroup size".
+
+    Args:
+      group_assignment: list of group ids for applying the optimizer
+        to subgroups.
+      num_shards: The number of TPU shards.
+
+    Returns:
+      The size of one subgroup in group_assignment.
+
+    Raises:
+      ValueError: If group_assignment is invalid.
+    """
+    if not group_assignment:
+      return None
+    if not (isinstance(group_assignment, list) and
+            all(isinstance(i, list) for i in group_assignment)):
+      raise ValueError(
+          f"Argument `group_assignment` must be a list of lists. "
+          f"Received: {group_assignment}")
+
+    replica_ids = set()
+    for g in group_assignment:
+      for i in g:
+        replica_ids.add(i)
+
+    if set(range(num_shards)) != replica_ids:
+      raise ValueError(
+          f"Argument `group_assignment` must be a permutation of "
+          f"range({num_shards}). Received: {group_assignment}")
+
+    subgroup_size_list = [len(group) for group in group_assignment]
+    if all(subgroup_size_list[0] == size for size in subgroup_size_list):
+      return subgroup_size_list[0]
+    else:
+      raise ValueError("The size of each subgroup in `group_assignment` must "
+                       f"be equal. Received: {group_assignment}")
+
+  def compute_gradients(self, loss, var_list=None, **kwargs):
+    """Compute gradients of "loss" for the variables in "var_list".
+
+    This simply wraps `compute_gradients()` from the real optimizer. The
+    gradients will be aggregated in `apply_gradients()` so that user can
+    modify the gradients like clipping with per replica global norm if needed.
+    The global norm with aggregated gradients can be bad as one replica's huge
+    gradients can hurt the gradients from other replicas.
+
+    When the CrossShardOptimizer is constructed with
+    `reduction == losses.Reduction.MEAN` (default), this function scales the
+    loss by `1.0 / num_shards` before computing the gradients. Assuming the
+    optimizer uses the default implementation of `compute_gradients()`, the
+    gradients of the scaled loss are scaled by `1.0 / num_shards` compared to
+    the gradients of the original loss. This scaling factor is important because
+    `apply_gradients()` sums gradients across shards, rather than averaging
+    them. However, the scaling factor must be taken into account when clipping
+    the norm of the gradients or performing other postprocessing.
+
+    Args:
+      loss: A Tensor containing the value to minimize.
+      var_list: Optional list or tuple of `tf.Variable` to update to minimize
+        `loss`.  Defaults to the list of variables collected in the graph
+        under the key `GraphKey.TRAINABLE_VARIABLES`.
+      **kwargs: Keyword arguments for compute_gradients().
+
+    Returns:
+      A list of (gradient, variable) pairs.
+
+    Raises:
+      ValueError: If not within a tpu_shard_context or group_assignment is
+        invalid.
+    """
+    num_shards = tpu_function.get_tpu_context().number_of_shards
+    if num_shards is None:
+      logging.warning(
+          "CrossShardOptimizer should be used within a tpu_shard_context, but "
+          "got unset number_of_shards. Assuming 1.")
+      num_shards = 1
+
+    subgroup_size = self._verify_and_get_subgroup_size(self._group_assignment,
+                                                       num_shards)
+
+    if num_shards > 1 and self._reduction == losses.Reduction.MEAN:
+      if self._group_assignment:
+        scale = 1.0 / subgroup_size
+      else:
+        scale = 1.0 / num_shards
+      loss *= scale
+
+    return self._opt.compute_gradients(loss, var_list=var_list, **kwargs)
+
+  def apply_gradients(self, grads_and_vars, global_step=None, name=None):
+    """Apply gradients to variables.
+
+    Calls tpu_ops.cross_replica_sum() to sum gradient contributions across
+    replicas, and then applies the real optimizer.
+
+    Args:
+      grads_and_vars: List of (gradient, variable) pairs as returned by
+        compute_gradients().
+      global_step: Optional Variable to increment by one after the
+        variables have been updated.
+      name: Optional name for the returned operation.  Default to the
+        name passed to the Optimizer constructor.
+
+    Returns:
+      An `Operation` that applies the gradients. If `global_step` was not None,
+      that operation also increments `global_step`.
+
+    Raises:
+      ValueError: If the grads_and_vars is malformed.
+    """
+    summed_grads_and_vars = []
+    for (grad, var) in grads_and_vars:
+      if grad is None:
+        summed_grads_and_vars.append((grad, var))
+      else:
+        with ops.colocate_with(grad):
+          summed_grads_and_vars.append((tpu_ops.cross_replica_sum(
+              grad, self._group_assignment), var))
+    return self._opt.apply_gradients(summed_grads_and_vars, global_step, name)
+
+  def get_slot(self, *args, **kwargs):
+    """Return a slot named "name" created for "var" by the Optimizer.
+
+    This simply wraps the get_slot() from the actual optimizer.
+
+    Args:
+      *args: Arguments for get_slot().
+      **kwargs: Keyword arguments for get_slot().
+
+    Returns:
+      The `Variable` for the slot if it was created, `None` otherwise.
+    """
+    return self._opt.get_slot(*args, **kwargs)
+
+  def get_slot_names(self, *args, **kwargs):
+    """Return a list of the names of slots created by the `Optimizer`.
+
+    This simply wraps the get_slot_names() from the actual optimizer.
+
+    Args:
+      *args: Arguments for get_slot().
+      **kwargs: Keyword arguments for get_slot().
+
+    Returns:
+      A list of strings.
+    """
+    return self._opt.get_slot_names(*args, **kwargs)
+
+  def variables(self):
+    """Forwarding the variables from the underlying optimizer."""
+    return self._opt.variables()

tensorflow/python/tpu/tpu_outside_compilation_test.py

@@ -0,0 +1,846 @@
+# 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 outside compilation."""
+
+import os
+import tempfile
+import unittest
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.core.util import event_pb2
+from tensorflow.python.distribute import tpu_strategy as tpu_lib
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import remote
+from tensorflow.python.eager import test
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.lib.io import tf_record
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops import summary_ops_v2 as summary
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import flags
+from tensorflow.python.platform import gfile
+from tensorflow.python.tpu import functional as tpu_functional
+from tensorflow.python.tpu import tpu
+from tensorflow.python.tpu import tpu_replication
+from tensorflow.python.tpu.ops import tpu_ops
+try:
+  from tensorboard.plugins.histogram import summary_v2 as histogram_summary_v2  # pylint: disable=g-import-not-at-top
+  from tensorboard.plugins.image import summary_v2 as image_summary_v2  # pylint: disable=g-import-not-at-top
+  from tensorboard.plugins.scalar import summary_v2 as scalar_summary_v2  # pylint: disable=g-import-not-at-top
+  _TENSORBOARD_AVAILABLE = True
+except ImportError:
+  _TENSORBOARD_AVAILABLE = False
+
+FLAGS = flags.FLAGS
+flags.DEFINE_bool(
+    "use_local_tpu",
+    False,
+    "use local TPUs on a TPU VM instead of connecting to a GCP TPU VM or node.",
+)
+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.")
+
+
+def get_tpu_cluster_resolver():
+  if FLAGS.use_local_tpu:
+    return tpu_cluster_resolver.TPUClusterResolver("local")
+  resolver = tpu_cluster_resolver.TPUClusterResolver(
+      tpu=FLAGS.tpu,
+      zone=FLAGS.zone,
+      project=FLAGS.project,
+  )
+  return resolver
+
+
+def get_tpu_strategy():
+  resolver = get_tpu_cluster_resolver()
+  remote.connect_to_cluster(resolver)
+  tpu_cluster_resolver.initialize_tpu_system(resolver)
+  return tpu_lib.TPUStrategyV2(resolver)
+
+
+def computation_with_string_ops(x):
+  output = string_ops.string_format("1{}", x)
+  return string_ops.string_to_number(output)
+
+
+def _events_from_logdir(test_case, logdir):
+  """Reads summary events from log directory."""
+  test_case.assertTrue(gfile.Exists(logdir))
+  files = gfile.ListDirectory(logdir)
+  test_case.assertLen(files, 1)
+  records = list(tf_record.tf_record_iterator(os.path.join(logdir, files[0])))
+  result = []
+  for r in records:
+    event = event_pb2.Event()
+    event.ParseFromString(r)
+    result.append(event)
+  return result
+
+
+def _rewrite_func_wrapper(tf_func):
+
+  def tpu_fn(*args, **kwargs):
+    # tpu.rewrite only accepts list of tensors as input. We need to flatten
+    # keyword arguments to meet this requirement.
+    concrete = tf_func.get_concrete_function(*(list(args) +
+                                               list(kwargs.values())))
+    return tpu.rewrite(concrete.__call__, list(args) + list(kwargs.values()))
+
+  return def_function.function(tpu_fn)
+
+
+def _tpu_partitioned_call_wrapper(tf_func):
+  """Wrap a tensorflow Function with TPUPartitionedCall."""
+
+  def inner_func(*args, **kwargs):
+    concrete = tf_func.get_concrete_function(*args, **kwargs)
+    # TPUPartitionedCall only accepts list of tensors as input args.
+    # Flatten keyword arguments and do some basic ordering:
+    # Positional args + Flattened keyword args + Captured args.
+    op_args = list(args) + list(kwargs.values()) + concrete.captured_inputs
+    return tpu_functional.TPUPartitionedCall(
+        args=op_args,
+        device_ordinal=tpu_ops.tpu_ordinal_selector(),
+        Tout=[o.type for o in concrete.function_def.signature.output_arg],
+        f=concrete)
+
+  return def_function.function(inner_func)
+
+
+class TpuOutsideCompilationTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    super(TpuOutsideCompilationTest, self).setUp()
+    config.set_soft_device_placement(False)
+
+  def testHostNoInput(self):
+    strategy = get_tpu_strategy()
+
+    def outside_fn():
+      logging_ops.print_v2("Outside compiled")
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        tpu_replication.outside_compilation(outside_fn)
+        return x2 + 5.0
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(35., shape=(strategy.num_replicas_in_sync)))
+
+  def testHostInputOnly(self):
+    strategy = get_tpu_strategy()
+
+    def outside_fn(x):
+      logging_ops.print_v2("Outside compiled", x)
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        tpu_replication.outside_compilation(outside_fn, x2)
+        return x2 + 5.0
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(35., shape=(strategy.num_replicas_in_sync)))
+
+  def testJitCompile(self):
+    strategy = get_tpu_strategy()
+
+    def outside_fn(x):
+      logging_ops.print_v2("Outside compiled", x)
+
+    # jit_compile=True should have no effect for TPU.
+    @def_function.function(jit_compile=True)
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        tpu_replication.outside_compilation(outside_fn, x2)
+        return x2 + 5.0
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(35., shape=(strategy.num_replicas_in_sync)))
+
+  def testHostInputOutput(self):
+    strategy = get_tpu_strategy()
+
+    def outside_fn(x):
+      logging_ops.print_v2("Outside compiled", x)
+      return x + 6.0
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        output = tpu_replication.outside_compilation(outside_fn, x2)
+        return output
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(36., shape=(strategy.num_replicas_in_sync)))
+
+  def testHostMultipleInputs(self):
+    strategy = get_tpu_strategy()
+    val0 = np.arange(6).reshape((2, 3)).astype(np.float32)
+    val1 = np.arange(6).reshape((3, 2)).astype(np.float32)
+
+    def outside_fn(arg0, arg1):
+      tmp = array_ops.reshape(arg1, array_ops.shape(arg0))
+      ret0 = arg0 + tmp
+      ret1 = math_ops.matmul(arg0, arg1)
+      ret2 = array_ops.concat([arg0, tmp], 0)
+      return ret0, ret1, ret2
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x, y):
+        a = x + 7.0
+        b = y * 2.0
+        c, d, e = tpu_replication.outside_compilation(outside_fn, a, b)
+        return (math_ops.reduce_max(c) + math_ops.reduce_min(d) +
+                math_ops.reduce_sum(e))
+
+      return strategy.run(tpu_fn, args=(val0, val1))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(213., shape=(strategy.num_replicas_in_sync)))
+
+  def testMultipleClusters(self):
+    strategy = get_tpu_strategy()
+
+    def outside_fn1(x):
+      logging_ops.print_v2("Outside compiled", x)
+      return x + 6.0
+
+    def outside_fn2(x):
+      logging_ops.print_v2("Outside compiled", x)
+      return x - 18.0
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        output1 = tpu_replication.outside_compilation(outside_fn1, x2)
+        x3 = output1 + 3.0
+        output2 = tpu_replication.outside_compilation(outside_fn2, x3)
+        return output2
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(21., shape=(strategy.num_replicas_in_sync)))
+
+  @parameterized.parameters((True), (False))
+  def testOutsideCompilationControlFlowIf(self, take_true_branch):
+    strategy = get_tpu_strategy()
+
+    def outside_fn(x):
+      logging_ops.print_v2("Outside compiled", x)
+      return x + 6.0
+
+    input_value = 51.0 if take_true_branch else 25.0
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        if x < 50.0:
+          return tpu_replication.outside_compilation(outside_fn, x2)
+        else:
+          return x2
+
+      return strategy.run(tpu_fn, args=(input_value,))
+
+    output_value = 36.0
+    if take_true_branch:
+      output_value = 56.0
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(
+            output_value, shape=(strategy.num_replicas_in_sync)))
+
+  def testOutsideCompilationControlFlowWhile(self):
+    strategy = get_tpu_strategy()
+
+    def outside_fn(x):
+      logging_ops.print_v2("Outside compiled", x)
+      return x + 6.0
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        while x2 < 50.0:
+          x2 = tpu_replication.outside_compilation(outside_fn, x2)
+        return x2 + 4.0
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(58., shape=(strategy.num_replicas_in_sync)))
+
+  def testOutsideCompilationHostControlFlow(self):
+    """Tests that control flow on host for outside_compilation works."""
+    strategy = get_tpu_strategy()
+
+    def outside_fn(x):
+      n = 0
+      while n < 4:
+        x = x + 6.0
+        n = n + 1
+      return x
+
+    @def_function.function
+    def train_step():
+
+      def tpu_fn(x):
+        x2 = x + 5.0
+        x2 = tpu_replication.outside_compilation(outside_fn, x2)
+        return x2 + 4.0
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(58., shape=(strategy.num_replicas_in_sync)))
+
+  @unittest.skipIf(not _TENSORBOARD_AVAILABLE, "Tensorboard is not installed.")
+  def testSummary(self):
+    strategy = get_tpu_strategy()
+
+    def host_computation(x):
+      scalar_summary_v2.scalar("x", x, step=0)
+      return x * 2.0
+
+    @def_function.function
+    def step():
+
+      def computation(x):
+        x = x + 1.0
+        y = tpu_replication.outside_compilation(host_computation, x)
+        y = tpu_replication.outside_compilation(host_computation, x)
+        return y + 1.0
+
+      return strategy.run(computation, args=(2.0,))
+
+    summary_writer = summary.create_file_writer(
+        os.path.join(os.getenv("TEST_TMPDIR", "/tmp")), flush_millis=10000)
+    with summary_writer.as_default(), summary.always_record_summaries():
+      self.assertAllEqual(
+          strategy.experimental_local_results(step()),
+          constant_op.constant(7., shape=(strategy.num_replicas_in_sync)))
+
+  @parameterized.parameters((True), (False))
+  @unittest.skipIf(not _TENSORBOARD_AVAILABLE, "Tensorboard is not installed.")
+  def testSummaryInCond(self, take_true_branch):
+    strategy = get_tpu_strategy()
+
+    def host_computation(x):
+      scalar_summary_v2.scalar("x", x, step=0)
+      return x * 2.0
+
+    @def_function.function
+    def step(take_true_branch):
+
+      def computation(x):
+        x = x + 1.0
+        if x < 5.0:
+          y = tpu_replication.outside_compilation(host_computation, x)
+          y = tpu_replication.outside_compilation(host_computation, x)
+          x = y
+        return x + 1.0
+
+      if take_true_branch:
+        return strategy.run(computation, args=(2.0,))
+      else:
+        return strategy.run(computation, args=(10.0,))
+
+    summary_writer = summary.create_file_writer(
+        os.path.join(os.getenv("TEST_TMPDIR", "/tmp")), flush_millis=10000)
+
+    output_value = 12.
+    if take_true_branch:
+      output_value = 7.
+    with summary_writer.as_default(), summary.always_record_summaries():
+      self.assertAllEqual(
+          strategy.experimental_local_results(step(take_true_branch)),
+          constant_op.constant(
+              output_value, shape=(strategy.num_replicas_in_sync)))
+
+  @unittest.skipIf(not _TENSORBOARD_AVAILABLE, "Tensorboard is not installed.")
+  def testSummaryInWhile(self):
+    strategy = get_tpu_strategy()
+
+    def host_computation(x):
+      scalar_summary_v2.scalar("x", x, step=0)
+      return x * 2.0
+
+    @def_function.function
+    def step():
+
+      def computation(x):
+        n = 0
+        while n < 3:
+          x = x + 1.0
+          y = tpu_replication.outside_compilation(host_computation, x)
+          y = tpu_replication.outside_compilation(host_computation, x)
+          x = y
+          n = n + 1
+        return x + 1.0
+
+      return strategy.run(computation, args=(2.0,))
+
+    summary_writer = summary.create_file_writer(
+        os.path.join(os.getenv("TEST_TMPDIR", "/tmp")), flush_millis=10000)
+    with summary_writer.as_default(), summary.always_record_summaries():
+      self.assertAllEqual(
+          strategy.experimental_local_results(step()),
+          constant_op.constant(31., shape=(strategy.num_replicas_in_sync)))
+
+  def testOutsideCompilationAtHeadAndTail(self):
+    """Tests that outside_compilation at head/tail of TPU computation works."""
+    strategy = get_tpu_strategy()
+
+    def host_computation(x):
+      return x * 2.0
+
+    @def_function.function
+    def train_step():
+
+      def computation(x):
+        w = tpu_replication.outside_compilation(host_computation, x)
+        y = w + 1.0
+        z = tpu_replication.outside_compilation(host_computation, y)
+        return z + 5.0
+
+      return strategy.run(computation, args=(2.0,))
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(15., shape=(strategy.num_replicas_in_sync)))
+
+  def testGradientAcrossOutsideCompilation(self):
+    """Tests compiled gradients can contain host computations."""
+    strategy = get_tpu_strategy()
+
+    def host_computation(a):
+      b = a * a
+      c = b * b
+      return c
+
+    @def_function.function
+    def train_step():
+      def computation(x, y):
+        a = x + 7.0
+        b = tpu_replication.outside_compilation(host_computation, a)
+        c = b * y
+        d = gradients_impl.gradients(
+            [c], [x], colocate_gradients_with_ops=True)[0]
+        return d
+
+      return strategy.run(computation, args=(2.0, 3.0))
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(8748., shape=(strategy.num_replicas_in_sync)))
+
+  def testGradientOfGradientAcrossOutsideCompilation(self):
+    """Tests compiled gradients of gradients can contain host computations."""
+    strategy = get_tpu_strategy()
+
+    def host_computation(a):
+      b = a * a
+      c = b * b
+      return c
+
+    @def_function.function
+    def train_step():
+      def computation(x, y):
+        a = x + 7.0
+        b = tpu_replication.outside_compilation(host_computation, a)
+        c = b * y
+        d = gradients_impl.gradients(
+            [c], [x], colocate_gradients_with_ops=True)[0]
+        e = gradients_impl.gradients(
+            [d], [x], colocate_gradients_with_ops=True)[0]
+        return e
+
+      return strategy.run(computation, args=(2.0, 3.0))
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(2916., shape=(strategy.num_replicas_in_sync)))
+
+  def testColocateGradientWithOutsideCompiledOp(self):
+    strategy = get_tpu_strategy()
+
+    @def_function.function
+    def train_step():
+
+      @def_function.function
+      def tpu_fn(x):
+        x1 = tpu_replication.outside_compilation(math_ops.sqrt, x)
+        grad = gradients_impl.gradients([x1], [x],
+                                        colocate_gradients_with_ops=True)[0]
+        sqrt = [
+            op for op in ops.get_default_graph().get_operations()
+            if op.type == "Sqrt"
+        ][0]
+        sqrt_grad = [
+            op for op in ops.get_default_graph().get_operations()
+            if op.type == "SqrtGrad"
+        ][0]
+        assert sqrt.get_attr(
+            tpu_replication._OUTSIDE_COMPILATION_ATTR) == b"0"
+        assert (sqrt_grad.get_attr(
+            tpu_replication._OUTSIDE_COMPILATION_ATTR) == b"0.gradients/uid"
+               )
+        return grad
+
+      return strategy.run(tpu_fn, args=(25.0,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step()),
+        constant_op.constant(.1, shape=(strategy.num_replicas_in_sync)))
+
+
+class OutsideCompilationOnUnsupportedOpTest(test.TestCase,
+                                            parameterized.TestCase):
+
+  def setUp(self):
+    super(OutsideCompilationOnUnsupportedOpTest, self).setUp()
+    config.set_soft_device_placement(True)
+
+  def testStringOpWithManualOutsideCompilation(self):
+    strategy = get_tpu_strategy()
+
+    @def_function.function
+    def train_step(x):
+
+      def computation(x):
+        return tpu_replication.outside_compilation(
+            computation_with_string_ops, x)
+
+      return strategy.run(computation, args=(x,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step(0)),
+        constant_op.constant(10, shape=(strategy.num_replicas_in_sync)))
+
+  def testStringOpWithAutoOutsideCompilation(self):
+    strategy = get_tpu_strategy()
+
+    @def_function.function
+    def train_step(x):
+
+      def computation(x):
+        return computation_with_string_ops(x)
+
+      return strategy.run(computation, args=(x,))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step(0)),
+        constant_op.constant(10, shape=(strategy.num_replicas_in_sync)))
+
+  # Regression test for b/180509859.
+  @unittest.skipIf(not _TENSORBOARD_AVAILABLE, "Tensorboard is not installed.")
+  def testImageSummary(self):
+    strategy = get_tpu_strategy()
+
+    def run():
+
+      @def_function.function
+      def sample_sequence():
+        bsz = 3
+        max_length = 32 * 32
+
+        def f():
+
+          def body(step, tokens):
+            next_token = random_ops.random_uniform([bsz])
+            tokens = tokens.write(step, next_token)
+            return (step + 1, tokens)
+
+          def cond_fn(step, tokens):
+            del tokens
+            return math_ops.less(step, max_length)
+
+          tokens_var = tensor_array_ops.TensorArray(
+              dtype=dtypes.float32,
+              size=max_length,
+              dynamic_size=False,
+              clear_after_read=False,
+              element_shape=(bsz,),
+              name="tokens_accumulator",
+          )
+
+          step = constant_op.constant(0)
+          step, tokens_var = while_loop.while_loop(
+              cond_fn, body, [step, tokens_var]
+          )
+
+          image_flat = array_ops.transpose(tokens_var.stack(), [1, 0])
+          image = array_ops.tile(
+              array_ops.reshape(image_flat, [bsz, 32, 32, 1]), [1, 1, 1, 3]
+          )
+          image_summary_v2.image(
+              "image_sample", image, constant_op.constant(5, dtype=dtypes.int64)
+          )
+
+        return strategy.run(f)
+
+      sample_sequence()
+
+    logdir = tempfile.mkdtemp()
+    summary_writer = summary.create_file_writer(logdir, flush_millis=10000)
+    with summary_writer.as_default(), summary.always_record_summaries():
+      run()
+    events = _events_from_logdir(self, logdir)
+    decoded_image = image_ops.decode_png(
+        events[1].summary.value[0].tensor.string_val[2]).numpy()
+    # Ensure that non-zero values were written to the image summary.
+    self.assertNotAllEqual(
+        array_ops.zeros((3072,), dtype=dtypes.float32),
+        list(decoded_image.flat))
+
+  @unittest.skipIf(not _TENSORBOARD_AVAILABLE, "Tensorboard is not installed.")
+  def testSummaryWithAutoOutsideCompilation(self):
+    strategy = get_tpu_strategy()
+
+    def host_computation(x):
+      scalar_summary_v2.scalar("x", x, step=0)
+      return x * 2.0
+
+    @def_function.function
+    def step():
+
+      def computation(x):
+        x = x + 1.0
+        y = host_computation(x)
+        return y + 1.0
+
+      return strategy.run(computation, args=(2.0,))
+
+    logdir = tempfile.mkdtemp()
+    summary_writer = summary.create_file_writer(logdir, flush_millis=10000)
+    with summary_writer.as_default(), summary.always_record_summaries():
+      self.assertAllEqual(
+          strategy.experimental_local_results(step()),
+          constant_op.constant(7., shape=(strategy.num_replicas_in_sync)))
+    events = _events_from_logdir(self, logdir)
+    # There will be 2 entries: 1 summary file header entry, and 1 entry
+    # written by host.
+    self.assertLen(events, 2)
+    self.assertEqual(events[1].summary.value[0].tag, "x")
+
+  @unittest.skipIf(not _TENSORBOARD_AVAILABLE, "Tensorboard is not installed.")
+  def testNestedFunctionScalarSummary(self):
+    strategy = get_tpu_strategy()
+
+    def host_computation(x):
+      scalar_summary_v2.scalar("x", x, step=0)
+      return x * 2.0
+
+    @def_function.function
+    def step():
+
+      @def_function.function
+      def computation(x):
+        x = x + 1.0
+        y = host_computation(x)
+        return y + 1.0
+
+      return strategy.run(computation, args=(2.0,))
+
+    logdir = tempfile.mkdtemp()
+    summary_writer = summary.create_file_writer(logdir, flush_millis=10000)
+    with summary_writer.as_default(), summary.always_record_summaries():
+      self.assertAllEqual(
+          strategy.experimental_local_results(step()),
+          constant_op.constant(7., shape=(strategy.num_replicas_in_sync)))
+    events = _events_from_logdir(self, logdir)
+    # There will be 2 entries: 1 summary file header entry, and 1 entry
+    # written by host.
+    self.assertLen(events, 2)
+    self.assertEqual(events[1].summary.value[0].tag, "x")
+
+  @unittest.skipIf(not _TENSORBOARD_AVAILABLE, "Tensorboard is not installed.")
+  def testHistogramSummaryWithAutoOutsideCompilation(self):
+    strategy = get_tpu_strategy()
+
+    def host_computation(x):
+      histogram_summary_v2.histogram("x", x, step=0)
+      return x * 2.0
+
+    @def_function.function
+    def step():
+
+      def computation(x):
+        x = x + 1.0
+        y = host_computation(x)
+        return y + 1.0
+
+      return strategy.run(computation, args=(2.0,))
+
+    logdir = tempfile.mkdtemp()
+    summary_writer = summary.create_file_writer(logdir, flush_millis=10000)
+    with summary_writer.as_default(), summary.always_record_summaries():
+      self.assertAllEqual(
+          strategy.experimental_local_results(step()),
+          constant_op.constant(7., shape=(strategy.num_replicas_in_sync)))
+    events = _events_from_logdir(self, logdir)
+    # There will be 2 entries: 1 summary file header entry, and 1 entry
+    # written by host.
+    self.assertLen(events, 2)
+    self.assertEqual(events[1].summary.value[0].tag, "x")
+
+  @parameterized.parameters((True), (False))
+  @unittest.skipIf(not _TENSORBOARD_AVAILABLE, "Tensorboard is not installed.")
+  def testSummaryControlFlowIfWithAutoOutsideCompilation(
+      self, take_true_branch):
+    strategy = get_tpu_strategy()
+
+    @def_function.function
+    def step():
+
+      def computation(x):
+        x = x + 1.0
+        if x < 5:
+          scalar_summary_v2.scalar("x", x, step=0)
+          x = x * 2.0
+        return x + 1.0
+
+      if take_true_branch:
+        return strategy.run(computation, args=(2.0,))
+      else:
+        return strategy.run(computation, args=(10.0,))
+
+    logdir = tempfile.mkdtemp()
+    summary_writer = summary.create_file_writer(logdir, flush_millis=10000)
+    output_value = 12.
+    if take_true_branch:
+      output_value = 7.
+    with summary_writer.as_default(), summary.always_record_summaries():
+      self.assertAllEqual(
+          strategy.experimental_local_results(step()),
+          constant_op.constant(
+              output_value, shape=(strategy.num_replicas_in_sync)))
+    if take_true_branch:
+      events = _events_from_logdir(self, logdir)
+      # There will be 2 entries: 1 summary file header entry, and 1 entry
+      # written by host.
+      #
+      self.assertLen(events, 2)
+      self.assertEqual(events[1].summary.value[0].tag, "cond/x")
+
+  def testAutoOutsideCompilationWithFunctionalNodes(self):
+    strategy = get_tpu_strategy()
+
+    @def_function.function
+    def train_step(a, b):
+
+      def fn(a, b):
+        fn1 = lambda: computation_with_string_ops(a * 100)
+        fn2 = lambda: computation_with_string_ops(a)
+        pred = math_ops.greater_equal(a, b)
+        result = array_ops.identity(
+            cond.cond(pred, fn1, fn2),
+            name="uncompilable_control_flow")
+        return result
+
+      return strategy.run(fn, args=(a, b))
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step(0.0, -1.0)),
+        constant_op.constant(10, shape=(strategy.num_replicas_in_sync)))
+
+  def testRandomOpsWithAutoOutsideCompilation(self):
+    strategy = get_tpu_strategy()
+
+    @def_function.function
+    def train_step():
+
+      def computation():
+        return random_ops.random_normal(shape=[1, 2, 3])
+
+      return strategy.run(computation, args=())
+
+    self.assertAllEqual(
+        strategy.experimental_local_results(train_step())[0].shape, [1, 2, 3])
+
+  def testOutsideCompilationWithTPUPartitionedCallOp(self):
+    """Tests that control flow with TPUPartitionedCall including outside_compilation works."""
+    get_tpu_strategy()
+
+    def host_computation(x):
+      return x + 1
+
+    @def_function.function()
+    def train_step(x):
+      x2 = x + 5.0
+      logging_ops.print_v2(x2)
+      x2 = tpu_replication.outside_compilation(host_computation, x2)
+      return x2 + 4.0
+
+    tpu_fn = _rewrite_func_wrapper(train_step)
+    partitioned_tpu_fn = _tpu_partitioned_call_wrapper(tpu_fn)
+
+    concrete = partitioned_tpu_fn.get_concrete_function(
+        x=tensor.TensorSpec(
+            shape=(1), dtype=dtypes.float32, name="input_tensor"))
+
+    self.assertIsInstance(
+        concrete(array_ops.ones((1), dtype=dtypes.float32))[0], tensor.Tensor)
+
+
+if __name__ == "__main__":
+  test.main()