# 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 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 _get_strategy(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)
    tpu_cluster_resolver.initialize_tpu_system(self.resolver)
    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):
    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)