# 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 embedding layers."""
import abc

from typing import AbstractSet, Dict, Mapping, Sequence, Tuple
import lingvo.compat as tf
from lingvo.core import base_layer
from lingvo.core import py_utils
from lingvo.core import tpu_embedding_layers

# pylint:disable=g-direct-tensorflow-import
from tensorflow.python.tpu import tpu_embedding as tpu_embedding_lib
# pylint:enable=g-direct-tensorflow-import


def _RemovePrivateVar(layer, var_name):
  """Remove a variable by name from `layer`.

  This is usually used to avoid copying the variable to TPU, for example, by the
  tf.cast when accessing layer.theta.

  Args:
    layer: The layer to remove the variable from.
    var_name: The name of the variable to remove.
  """
  # pylint: disable=protected-access
  del layer._private_vars[var_name]
  del layer._private_theta[var_name]
  # pylint: enable=protected-access


class TpuEmbeddingCollection:
  """Manage various TPU embedding related ops and tensors."""

  @classmethod
  def Get(cls):
    """Returns the TpuEmbeddingCollection associated with the current graph."""
    emb_collection = py_utils.GetTpuEmbeddingGraphCollection()
    assert len(emb_collection) <= 1
    if len(emb_collection) == 1:
      tf.logging.info(
          'TpuEmbeddingCollection singleton already exists, reusing'
      )
      return emb_collection[0]
    else:
      singleton = cls()
      emb_collection.append(singleton)
      return singleton

  def __init__(self):
    # Maps table name to a tuple (var_list, is_inference_with_bfloat16), where
    # var_list is the list of variables for the corresponding table, and
    # is_inference_with_bfloat16 is a boolean telling whether this table is
    # using bfloat16 for inference.
    self._table_vars = py_utils.NestedMap()

    # The TPUEmbedding configuration.
    self._tpu_embedding = None

    # Maps table name to the list of ops that loads/retrieves embedding tables
    # to/from TPU.
    self._load_ops_map = py_utils.NestedMap()
    self._retrieve_ops_map = py_utils.NestedMap()

    # Maps task name to the (feature_name -> activation_tensor) dict for the
    # corresponding task.
    self._activations_by_task = {}

    # List of (name, value, weight) tuples for summary.
    self._summary_tensors = []

    # Set of embedding feature names.
    self._feature_names = None

    # Schedule for the value that is used as TPU embedding gradient multiplier.
    self._gradient_multiplier_schedule = None

    # Maps task name to the mode used by that task.
    self._mode_by_task = {}

    # Maps task name to the send gradient op for that task. Mainly used to
    # ensure that send gradient op is created only once for each task.
    self._send_gradient_op_by_task = {}

  def AddTableVariables(self, table_name, var_list, is_inference_with_bfloat16):
    """Add TPU embedding table variable list to the collection."""
    if table_name in self._table_vars:
      raise ValueError(f'Variables for table {table_name} already exist.')
    self._table_vars[table_name] = (var_list, is_inference_with_bfloat16)

  @property
  def table_variables(self) -> py_utils.NestedMap:
    """Returns a nested map of table variables."""
    return self._table_vars.Transform(lambda val: val[0])

  @property
  def inference_with_bfloat16_var_names(self):
    """Returns a list of names of table variables that do bfloat16 inference."""
    result = []
    for var_list, is_inference_with_bfloat16 in self._table_vars.values():
      if is_inference_with_bfloat16:
        result += [v.op.name for v in var_list]
    return result

  @property
  def tpu_embedding(self) -> tpu_embedding_lib.TPUEmbedding:
    return self._tpu_embedding

  @tpu_embedding.setter
  def tpu_embedding(self, tpu_embedding: tpu_embedding_lib.TPUEmbedding):
    if self._tpu_embedding is not None:
      raise ValueError('TPUEmbedding has already been set.')
    self._tpu_embedding = tpu_embedding

  def AddLoadRetrieveOps(self, table_name, load_ops, retrieve_ops):
    if table_name in self._load_ops_map:
      raise ValueError(f'Load ops for table {table_name} already exist.')
    assert table_name not in self._retrieve_ops_map
    self._load_ops_map[table_name] = load_ops
    self._retrieve_ops_map[table_name] = retrieve_ops

  @property
  def load_ops(self):
    return self._load_ops_map

  @property
  def retrieve_ops(self):
    return self._retrieve_ops_map

  def _ValidateTaskScope(self, task_call_scope: str):
    if not task_call_scope:
      raise ValueError(
          'It expects a non-empty task call scope name, but get '
          f'{task_call_scope}. This usually means the current code is not run '
          'under a py_utils.TaskCallScope() context.'
      )

  def AddActivations(self, task_call_scope: str) -> Mapping[str, tf.Tensor]:
    self._ValidateTaskScope(task_call_scope)
    tf.logging.info(
        f'Adding TPU embedding activations for task {task_call_scope}.'
    )
    if task_call_scope not in self._activations_by_task:
      activations = self._tpu_embedding.get_activations()
      self._activations_by_task[task_call_scope] = activations
    return self._activations_by_task[task_call_scope]

  def GetActivations(self, task_call_scope: str) -> Mapping[str, tf.Tensor]:
    tf.logging.info(
        f'Getting TPU embedding activations for task {task_call_scope}.'
    )
    if activations := self._activations_by_task.get(task_call_scope):
      self._ValidateTaskScope(task_call_scope)
    return activations

  def AddSummaryTensor(self, name, value, weight=1.0):
    self._summary_tensors.append((name, value, tf.convert_to_tensor(weight)))

  @property
  def summary_tensors(self) -> Sequence[Tuple[str, tf.Tensor, tf.Tensor]]:
    return self._summary_tensors

  @property
  def feature_names(self) -> AbstractSet[str]:
    return self._feature_names

  @feature_names.setter
  def feature_names(self, feature_names: AbstractSet[str]):
    if self._feature_names and self._feature_names != feature_names:
      raise ValueError(
          'feature_names already exists. '
          f'Existing feature names: {self._feature_names}, '
          f'feature names being added: {feature_names}'
      )
    self._feature_names = feature_names

  def SetGradientMultiplierSchedule(self, multiplier_schedule):
    if self._gradient_multiplier_schedule is not None:
      raise ValueError('gradient_multiplier_schedule was set before.')
    self._gradient_multiplier_schedule = multiplier_schedule

  def SetTaskMode(self, task_call_scope, mode):
    self._ValidateTaskScope(task_call_scope)
    tf.logging.info(
        f'Setting TPU embedding mode for task {task_call_scope} as {mode}.'
    )
    self._mode_by_task[task_call_scope] = mode

  def ShouldStopGradient(self, task_call_scope):
    self._ValidateTaskScope(task_call_scope)
    if task_call_scope not in self._mode_by_task:
      raise ValueError(
          f'TPU embedding mode for task {task_call_scope} not found.'
      )
    should_stop_gradient = self._mode_by_task[task_call_scope] != 'train'
    tf.logging.info(
        ('Disabled' if should_stop_gradient else 'Enabled')
        + f' TPU embedding gradient for task {task_call_scope}.'
    )
    return should_stop_gradient

  def ApplyGradients(self, task_call_scope, feature_to_gradient_dict):
    """Apply tpu embedding gradient updates.

    Args:
      task_call_scope: The current task call scope name.
      feature_to_gradient_dict: A `py_utils.NestedMap` of: tpu embedding feature
        name -> gradient tensor for the embedding feature.

    Returns:
      The gradient update op and a dict of eval metrics.

    Raises:
      ValueError: if gradients have been applied before for the current task.
    """
    self._ValidateTaskScope(task_call_scope)
    if task_call_scope in self._send_gradient_op_by_task:
      raise ValueError(
          f'Send gradient op for task {task_call_scope} already exist.'
      )
    tf.logging.info(
        f'Applying TPU embedding gradients for task {task_call_scope}.'
    )

    # Apply gradient multiplier schedule.
    grad_multiplier = self._gradient_multiplier_schedule.Value()
    feature_to_gradient_dict = feature_to_gradient_dict.Transform(
        lambda g: g * grad_multiplier
    )

    send_gradient_op = self._tpu_embedding.generate_send_gradients_op(
        feature_to_gradient_dict, step=py_utils.GetGlobalStep()
    )
    self._send_gradient_op_by_task[task_call_scope] = send_gradient_op

    activations = self.GetActivations(task_call_scope).values()
    eval_metrics = {
        'tpu_embedding_activation_norm': (
            tf.sqrt(py_utils.SumSquared(activations)),
            tf.constant(1.0),
        ),
        'tpu_embedding_grad_norm': (
            tf.sqrt(py_utils.SumSquared(feature_to_gradient_dict.Flatten())),
            tf.constant(1.0),
        ),
        'tpu_embedding_gradient_multiplier': (
            grad_multiplier,
            tf.constant(1.0),
        ),
    }
    return send_gradient_op, eval_metrics


class _TPUEmbeddingOptimizerV1Mixin(
    base_layer.BaseLayer, metaclass=base_layer.ABCLayerMeta
):
  """Defines the inferface for optimizers expected by the V1 Layer below."""

  def __init__(self, params):
    super().__init__(params)
    # A dict of slot_variable_name -> slot_variable for checkpointing purposes.
    self._slot_var_dict = {}

  @abc.abstractmethod
  def CreateOptimizerParameters(self, learning_rate: float):
    """Create TPUEmbedding API optimizer parameters."""

  @abc.abstractmethod
  def CreateSlotVariablesAndOps(
      self, table_vars: tf.Tensor, tpu_embedding_table: 'TPUEmbeddingTable'
  ):
    """Create slot variables and load/retrieve ops.

    Args:
      table_vars: A list of all embedding table shard variables.
      tpu_embedding_table: Parent TPUEmbeddingTable layer.

    Returns:
      List of load ops
      List of retrieve ops
    """

  def _GetSelfVariablesDict(self):
    """Returns a dict of variables for checkpointing purposes."""
    return self._slot_var_dict


class TPUEmbeddingSGDOptimizer(
    tpu_embedding_layers.TPUEmbeddingSGDOptimizer,
    _TPUEmbeddingOptimizerV1Mixin,
):
  """SGD optimizer for TPUEmbeddingLayer, TPUEmbeddingTable."""

  def CreateOptimizerParameters(self, learning_rate: float):
    p = self.params
    return tpu_embedding_lib.StochasticGradientDescentParameters(
        learning_rate=learning_rate,
        clip_weight_min=p.clip_weight_min,
        clip_weight_max=p.clip_weight_max,
        weight_decay_factor=p.weight_decay_factor,
        multiply_weight_decay_factor_by_learning_rate=(
            p.multiply_weight_decay_factor_by_learning_rate
        ),
        clip_gradient_min=p.clip_gradient_min,
        clip_gradient_max=p.clip_gradient_max,
    )

  def CreateSlotVariablesAndOps(
      self, table_vars: tf.Tensor, tpu_embedding_table: 'TPUEmbeddingTable'
  ):
    load_op_list = []
    retrieve_op_list = []

    num_tpu_hosts = tpu_embedding_table.params.num_tpu_hosts
    table_name = tpu_embedding_table.table_name

    for host_id, table_var in zip(range(num_tpu_hosts), table_vars):
      # The slot vars should be on the same device as the table var.
      device_name = tpu_embedding_table.GetDeviceName(host_id)
      with tf.device(device_name), py_utils.outside_all_rewrites():
        # Only the Trainer needs these ops.
        if py_utils.use_tpu():
          # TPU Embedding load/retrieve ops need to be in the outer graph scope.
          with tf.init_scope():
            tf.logging.info('creating load and retrieve ops.')
            load_parameters_op = tpu_embedding_lib.tpu_ops.load_tpu_embedding_stochastic_gradient_descent_parameters(
                parameters=table_var,
                table_name=table_name,
                num_shards=num_tpu_hosts,
                shard_id=host_id,
            )
            load_op_list.append(load_parameters_op)

            retrieved_table = tpu_embedding_lib.tpu_ops.retrieve_tpu_embedding_stochastic_gradient_descent_parameters(
                table_name=table_name,
                num_shards=num_tpu_hosts,
                shard_id=host_id,
            )
            retrieve_parameters_op = tpu_embedding_lib.control_flow_ops.group(
                tf.assign(table_var, retrieved_table)
            )
            retrieve_op_list.append(retrieve_parameters_op)

    return load_op_list, retrieve_op_list


class TPUEmbeddingAdagradOptimizer(
    tpu_embedding_layers.TPUEmbeddingAdagradOptimizer,
    _TPUEmbeddingOptimizerV1Mixin,
):
  """Adagrad optimizer for TPUEmbeddingLayer, TPUEmbeddingTable."""

  def CreateOptimizerParameters(self, learning_rate: float):
    p = self.params
    return tpu_embedding_lib.AdagradParameters(
        learning_rate=learning_rate,
        initial_accumulator=p.initial_accumulator,
        use_gradient_accumulation=p.use_gradient_accumulation,
        clip_weight_min=p.clip_weight_min,
        clip_weight_max=p.clip_weight_max,
        weight_decay_factor=p.weight_decay_factor,
        multiply_weight_decay_factor_by_learning_rate=(
            p.multiply_weight_decay_factor_by_learning_rate
        ),
        clip_gradient_min=p.clip_gradient_min,
        clip_gradient_max=p.clip_gradient_max,
    )

  def CreateSlotVariablesAndOps(
      self, table_vars: tf.Tensor, tpu_embedding_table: 'TPUEmbeddingTable'
  ):
    p = self.params

    load_op_list = []
    retrieve_op_list = []

    num_tpu_hosts = tpu_embedding_table.params.num_tpu_hosts
    table_name = tpu_embedding_table.table_name
    slot_var_collections = [tpu_embedding_table.__class__.__name__ + '_vars']

    for host_id, table_var in zip(range(num_tpu_hosts), table_vars):
      # The slot vars should be on the same device as the table var.
      device_name = tpu_embedding_table.GetDeviceName(host_id)
      with tf.device(device_name), py_utils.outside_all_rewrites():
        w_ada = py_utils.WeightParams(
            shape=table_var.shape.as_list(),
            init=py_utils.WeightInit.Constant(p.initial_accumulator),
            dtype=p.dtype,
            collections=slot_var_collections,
        )
        var_name = tpu_embedding_table.GetVariableName(host_id) + '/Adagrad'
        accumulator_var = tpu_embedding_table.CreateOptimizerSlotVariable(
            var_name, w_ada, self._slot_var_dict
        )

        # Only the Trainer needs these ops.
        if py_utils.use_tpu():
          # Remove the slot vars from the variable list to avoid them being
          # copied to TPU.
          _RemovePrivateVar(tpu_embedding_table, var_name)

          # TPU Embedding load/retrieve ops need to be in the outer graph scope.
          with tf.init_scope():
            tf.logging.info('creating load and retrieve ops.')
            load_parameters_op = (
                tpu_embedding_lib.tpu_ops.load_tpu_embedding_adagrad_parameters(
                    parameters=table_var,
                    accumulators=accumulator_var,
                    table_name=table_name,
                    num_shards=num_tpu_hosts,
                    shard_id=host_id,
                )
            )
            load_op_list.append(load_parameters_op)

            retrieved_table, retrieved_accumulator = (
                tpu_embedding_lib.tpu_ops.retrieve_tpu_embedding_adagrad_parameters(
                    table_name=table_name,
                    num_shards=num_tpu_hosts,
                    shard_id=host_id,
                )
            )
            retrieve_parameters_op = tpu_embedding_lib.control_flow_ops.group(
                tf.assign(table_var, retrieved_table),
                tf.assign(accumulator_var, retrieved_accumulator),
            )
            retrieve_op_list.append(retrieve_parameters_op)

    return load_op_list, retrieve_op_list


class TPUEmbeddingAdamOptimizer(
    tpu_embedding_layers.TPUEmbeddingAdamOptimizer,
    _TPUEmbeddingOptimizerV1Mixin,
):
  """Adam optimizer for TPUEmbeddingLayer, TPUEmbeddingTable."""

  def CreateOptimizerParameters(self, learning_rate: float):
    p = self.params
    return tpu_embedding_lib.AdamParameters(
        learning_rate=learning_rate,
        beta1=p.beta1,
        beta2=p.beta2,
        epsilon=p.epsilon,
        lazy_adam=p.lazy_adam,
        sum_inside_sqrt=p.sum_inside_sqrt,
        use_gradient_accumulation=p.use_gradient_accumulation,
        clip_weight_min=p.clip_weight_min,
        clip_weight_max=p.clip_weight_max,
        weight_decay_factor=p.weight_decay_factor,
        multiply_weight_decay_factor_by_learning_rate=(
            p.multiply_weight_decay_factor_by_learning_rate
        ),
        clip_gradient_min=p.clip_gradient_min,
        clip_gradient_max=p.clip_gradient_max,
    )

  def CreateSlotVariablesAndOps(
      self, table_vars: tf.Tensor, tpu_embedding_table: 'TPUEmbeddingTable'
  ):
    p = self.params

    load_op_list = []
    retrieve_op_list = []

    num_tpu_hosts = tpu_embedding_table.params.num_tpu_hosts
    table_name = tpu_embedding_table.table_name
    slot_var_collections = [tpu_embedding_table.__class__.__name__ + '_vars']

    for host_id, table_var in zip(range(num_tpu_hosts), table_vars):
      # The slot vars should be on the same device as the table var.
      device_name = tpu_embedding_table.GetDeviceName(host_id)
      with tf.device(device_name), py_utils.outside_all_rewrites():
        m_adam = py_utils.WeightParams(
            shape=table_var.shape.as_list(),
            init=py_utils.WeightInit.Constant(0.0),
            dtype=p.dtype,
            collections=slot_var_collections,
        )
        var_name_m = tpu_embedding_table.GetVariableName(host_id) + '/Adam/m'
        m_var = tpu_embedding_table.CreateOptimizerSlotVariable(
            var_name_m, m_adam, self._slot_var_dict
        )

        v_adam = py_utils.WeightParams(
            shape=table_var.shape.as_list(),
            init=py_utils.WeightInit.Constant(0.0),
            dtype=p.dtype,
            collections=slot_var_collections,
        )
        var_name_v = tpu_embedding_table.GetVariableName(host_id) + '/Adam/v'
        v_var = tpu_embedding_table.CreateOptimizerSlotVariable(
            var_name_v, v_adam, self._slot_var_dict
        )

        # Only the Trainer needs these ops.
        if py_utils.use_tpu():
          # Remove the slot vars from the variable list to avoid them being
          # copied to TPU.
          _RemovePrivateVar(tpu_embedding_table, var_name_m)
          _RemovePrivateVar(tpu_embedding_table, var_name_v)

          # TPU Embedding load/retrieve ops need to be in the outer graph scope.
          with tf.init_scope():
            tf.logging.info('creating load and retrieve ops.')
            load_parameters_op = (
                tpu_embedding_lib.tpu_ops.load_tpu_embedding_adam_parameters(
                    parameters=table_var,
                    momenta=m_var,
                    velocities=v_var,
                    table_name=table_name,
                    num_shards=num_tpu_hosts,
                    shard_id=host_id,
                )
            )
            load_op_list.append(load_parameters_op)

            retrieved_table, retrieved_m, retrieved_v = (
                tpu_embedding_lib.tpu_ops.retrieve_tpu_embedding_adam_parameters(
                    table_name=table_name,
                    num_shards=num_tpu_hosts,
                    shard_id=host_id,
                )
            )
            retrieve_parameters_op = tpu_embedding_lib.control_flow_ops.group(
                tf.assign(table_var, retrieved_table),
                tf.assign(m_var, retrieved_m),
                tf.assign(v_var, retrieved_v),
            )
            retrieve_op_list.append(retrieve_parameters_op)

    return load_op_list, retrieve_op_list


class TPUEmbeddingFTRLOptimizer(
    tpu_embedding_layers.TPUEmbeddingFTRLOptimizer,
    _TPUEmbeddingOptimizerV1Mixin,
):
  """FTRL optimizer for TPUEmbeddingLayer, TPUEmbeddingTable."""

  def CreateOptimizerParameters(self, learning_rate: float):
    p = self.params
    return tpu_embedding_lib.FtrlParameters(
        learning_rate=learning_rate,
        learning_rate_power=p.learning_rate_power,
        initial_accumulator_value=p.initial_accumulator_value,
        l1_regularization_strength=p.l1_regularization_strength,
        l2_regularization_strength=p.l2_regularization_strength,
        use_gradient_accumulation=p.use_gradient_accumulation,
        clip_weight_min=p.clip_weight_min,
        clip_weight_max=p.clip_weight_max,
        weight_decay_factor=p.weight_decay_factor,
        multiply_weight_decay_factor_by_learning_rate=(
            p.multiply_weight_decay_factor_by_learning_rate
        ),
        multiply_linear_by_learning_rate=p.multiply_linear_by_learning_rate,
        beta=p.beta,
        allow_zero_accumulator=p.allow_zero_accumulator,
        clip_gradient_min=p.clip_gradient_min,
        clip_gradient_max=p.clip_gradient_max,
    )

  def CreateSlotVariablesAndOps(
      self, table_vars: tf.Tensor, tpu_embedding_table: 'TPUEmbeddingTable'
  ):
    p = self.params

    load_op_list = []
    retrieve_op_list = []

    num_tpu_hosts = tpu_embedding_table.params.num_tpu_hosts
    table_name = tpu_embedding_table.table_name
    slot_var_collections = [tpu_embedding_table.__class__.__name__ + '_vars']

    for host_id, table_var in zip(range(num_tpu_hosts), table_vars):
      # The slot vars should be on the same device as the table var.
      device_name = tpu_embedding_table.GetDeviceName(host_id)
      with tf.device(device_name), py_utils.outside_all_rewrites():
        accumulator = py_utils.WeightParams(
            shape=table_var.shape.as_list(),
            init=py_utils.WeightInit.Constant(p.initial_accumulator_value),
            dtype=p.dtype,
            collections=slot_var_collections,
        )
        accumulator_name = (
            tpu_embedding_table.GetVariableName(host_id) + '/Ftrl'
        )
        accumulator_var = tpu_embedding_table.CreateOptimizerSlotVariable(
            accumulator_name, accumulator, self._slot_var_dict
        )

        linear = py_utils.WeightParams(
            shape=table_var.shape.as_list(),
            init=py_utils.WeightInit.Constant(p.initial_linear_value),
            dtype=p.dtype,
            collections=slot_var_collections,
        )
        linear_name = tpu_embedding_table.GetVariableName(host_id) + '/Ftrl_1'
        linear_var = tpu_embedding_table.CreateOptimizerSlotVariable(
            linear_name, linear, self._slot_var_dict
        )

        # Only the Trainer needs these ops.
        if py_utils.use_tpu():
          # Remove the slot vars from the variable list to avoid them being
          # copied to TPU.
          _RemovePrivateVar(tpu_embedding_table, accumulator_name)
          _RemovePrivateVar(tpu_embedding_table, linear_name)

          # TPU Embedding load/retrieve ops need to be in the outer graph scope.
          with tf.init_scope():
            tf.logging.info('creating load and retrieve ops.')
            load_parameters_op = (
                tpu_embedding_lib.tpu_ops.load_tpu_embedding_ftrl_parameters(
                    parameters=table_var,
                    accumulators=accumulator_var,
                    linears=linear_var,
                    table_name=table_name,
                    num_shards=num_tpu_hosts,
                    shard_id=host_id,
                )
            )
            load_op_list.append(load_parameters_op)

            retrieved_table, retrieved_accumulator, retrieved_linear = (
                tpu_embedding_lib.tpu_ops.retrieve_tpu_embedding_ftrl_parameters(
                    table_name=table_name,
                    num_shards=num_tpu_hosts,
                    shard_id=host_id,
                )
            )
            retrieve_parameters_op = tpu_embedding_lib.control_flow_ops.group(
                tf.assign(table_var, retrieved_table),
                tf.assign(accumulator_var, retrieved_accumulator),
                tf.assign(linear_var, retrieved_linear),
            )
            retrieve_op_list.append(retrieve_parameters_op)

    return load_op_list, retrieve_op_list


class TPUEmbeddingTable(tpu_embedding_layers.TPUEmbeddingTable):
  """An embedding table controlled by TPUEmbeddingLayer.

  Note that all input_keys needs to be declared upfront.
  """

  def __init__(self, params):
    super().__init__(params)
    p = self.params
    if p.combiner is None:
      assert p.max_sequence_length
    if p.max_sequence_length is not None and p.max_sequence_length > 0:
      assert p.combiner is None

    self._tpu_embedding_collection = TpuEmbeddingCollection.Get()

    def LearningRateFn(step):
      with py_utils.GlobalStepContext(step):
        lr = self.schedule.Value() * p.learning_rate
      self._tpu_embedding_collection.AddSummaryTensor(
          f'tpu_embedding_lr/{p.name}', lr
      )
      return lr

    self._table_config = tpu_embedding_lib.TableConfig(
        self._padded_vocab_size,
        p.embedding_dim,
        combiner=p.combiner,
        learning_rate=None,
        learning_rate_fn=LearningRateFn,
        # All TableConfigs passed to API will have a learning rate function,
        # so the learning_rate in the optimization_parameters is not used.
        optimization_parameters=self.optimizer.CreateOptimizerParameters(
            p.learning_rate
        ),
    )

  @property
  def table_config(self) -> tpu_embedding_lib.TableConfig:
    return self._table_config

  @property
  def auxiliary_variables(self):
    """Returns the auxiliary variables associated with this table."""
    return self._auxiliary_var_dict

  def _CreateLayerVariables(self):
    p = self.params

    # Reuse the singleton table variables if they were created before.
    # TODO(laigd): consider disabling this option. There should be at most one
    # TPU embedding layer in the model.
    all_table_vars = self._tpu_embedding_collection.table_variables
    if self.table_name in all_table_vars:
      embedding_table_vars = all_table_vars[self.table_name]
    else:
      inference_with_merged_var = False
      is_inference_with_bfloat16 = False
      dtype = p.dtype
      init = p.params_init
      if p.is_inference:
        inference_with_merged_var = p.inference_use_merged_variable
        if p.inference_variable_dtype is not None:
          dtype = p.inference_variable_dtype
        is_inference_with_bfloat16 = p.inference_variable_dtype == tf.bfloat16
        if dtype not in [tf.float32, tf.bfloat16]:
          # Note: it doesn't matter what initialization value is used, the
          # actual value of the variable will be loaded from checkpoint.
          # But for backward compatibility we use params_init for floating
          # point dtypes.
          init = py_utils.WeightInit.Constant(dtype.as_numpy_dtype(0))

      w_pc = py_utils.WeightParams(
          shape=[
              p.vocab_size
              if inference_with_merged_var
              else self._ids_per_shard,
              p.embedding_dim,
          ],
          init=init,
          dtype=dtype,
          collections=[self.__class__.__name__ + '_vars'],
      )

      embedding_table_vars = []
      if inference_with_merged_var:
        with py_utils.outside_all_rewrites():
          var_name = 'merged_var'
          self.CreateVariable(var_name, w_pc)
          embedding_var = self.vars[var_name]
          embedding_table_vars.append(embedding_var)
          # Remove from _private_vars / _private_thetas to be added later as wm.
          _RemovePrivateVar(self, var_name)
      else:
        for i in range(p.num_tpu_hosts):
          device_name = self.GetDeviceName(i)
          with tf.device(device_name), py_utils.outside_all_rewrites():
            var_name = self.GetVariableName(i)
            self.CreateVariable(var_name, w_pc)
            embedding_var = self.vars[var_name]
            embedding_table_vars.append(embedding_var)
            # Remove from _private_vars / _private_thetas to be added later as
            # wm.
            _RemovePrivateVar(self, var_name)

      # Track the table variables so they can be excluded from EMA.
      self._tpu_embedding_collection.AddTableVariables(
          self.table_name, embedding_table_vars, is_inference_with_bfloat16
      )

    # Create auxiliary inference-only variables, if any.
    # NOTE: it's not supported to reuse auxiliary variables if they were
    # created before in another TPU embedding layer with the same settings.
    self._auxiliary_var_dict = py_utils.NestedMap()
    if p.is_inference and p.inference_auxiliary_variable_specs:
      assert isinstance(p.inference_auxiliary_variable_specs, dict)
      for name, (dtype, shape) in p.inference_auxiliary_variable_specs.items():
        aux_params = py_utils.WeightParams(
            shape=shape,
            init=py_utils.WeightInit.Constant(dtype.as_numpy_dtype(0)),
            dtype=dtype,
            collections=[self.__class__.__name__ + '_vars'],
        )
        # Set trainable=False to exclude it from EMA.
        self.CreateVariable(name, aux_params, trainable=False)
        self._auxiliary_var_dict[name] = self.vars[name]

    if not py_utils.IsTpuTraining(p):
      # We don't need this for TrainerTpu, as the vars are not directly
      # accessed besides in the TPU embedding load/retrieve ops.
      # However, this is needed for CPU (eval/decode/controller).
      self._private_vars['wm'] = embedding_table_vars
      self._private_theta['wm'] = embedding_table_vars

    # If slot variables and load/retrieve ops were created before, maybe by a
    # different program or task, don't create it again.
    # Note that there should be only one copy of slot variables and
    # load/retrieve ops in the graph and they're shared by different
    # tasks/programs.
    all_load_ops = self._tpu_embedding_collection.load_ops
    if self.table_name not in all_load_ops:
      assert self.table_name not in self._tpu_embedding_collection.retrieve_ops
      # Only trainer and controller (for checkpointing) need slot variables.
      # Only trainer needs load/retrieve ops.
      if not self.do_eval and not p.is_inference:
        load_ops, retrieve_ops = self.optimizer.CreateSlotVariablesAndOps(
            embedding_table_vars, self
        )
        self._tpu_embedding_collection.AddLoadRetrieveOps(
            self.table_name, load_ops, retrieve_ops
        )

  def _GetSelfVariablesDict(self) -> Mapping[str, tf.Tensor]:
    """Returns a dict of variables for checkpointing purposes."""
    all_table_vars = self._tpu_embedding_collection.table_variables
    assert self.table_name in all_table_vars
    return {var.name: var for var in all_table_vars[self.table_name]}

  def CreateOptimizerSlotVariable(self, var_name, var_params, slot_var_dict):
    """Create optimizer slot variable and add it to the given variable dict."""
    self.CreateVariable(var_name, var_params, trainable=False)
    var = self.vars[var_name]
    slot_var_dict[var.name] = var
    return var

  # Return device to place sharded variables on.
  def GetDeviceName(self, host_id):
    if self.params.is_inference:
      # This is to place variables on the same device as other variables.
      return None
    if self.do_eval:
      return '/cpu:0'
    else:
      worker = self.cluster.params.worker.name
      return f'{worker}/replica:0/task:{host_id}/device:CPU:0'

  # Return variable name for embedding table shards.
  def GetVariableName(self, host_id):
    return f'var_{host_id}'


class TPUEmbeddingLayer(tpu_embedding_layers.TPUEmbeddingLayer):
  """Monolithic interface to TPU embedding.

  This layer has some important caveats, due to the interface of the
  TPU embedding hardware. Its behavior most closely mimics that of
  tf.nn.embedding_lookup_sparse.

  Supports multiple tables and multiple input_keys per table.
  Requires its own optimizer parameters.
  """

  @classmethod
  def Params(cls):
    p = super().Params()
    # We override this parameter so that it has a valid default.
    p.optimizer = TPUEmbeddingAdagradOptimizer.Params()
    return p

  def __init__(self, params):
    super().__init__(params)
    self._tpu_embedding_collection = TpuEmbeddingCollection.Get()

    # Save embedding feature names in the collection.
    feature_names = set()
    for table in self.tables:
      for feature in table.input_keys:
        if feature in feature_names:
          raise ValueError(f'Input key {feature} was used by multiple tables.')
        feature_names.add(feature)
    self._tpu_embedding_collection.feature_names = feature_names

  def _child_variable_scope_override(self):
    return {**super()._child_variable_scope_override(), 'tables': []}

  def _CheckTPUEmbeddingConfig(
      self,
      tpu_embedding: tpu_embedding_lib.TPUEmbedding,
      table_to_config_dict: Mapping[str, tpu_embedding_lib.TableConfig],
      feature_to_config_dict: Mapping[str, tpu_embedding_lib.FeatureConfig],
      global_batch_size: int,
  ):
    """Check that the existing tpu_embedding config matches the given ones."""

    def _Match(d1, d2, namedtuple_attrs_to_check):
      if len(d1) != len(d2):
        return False
      for k, v1 in d1.items():
        if k not in d2:
          return False
        v2 = d2[k]
        for attr in namedtuple_attrs_to_check:
          if getattr(v1, attr) != getattr(v2, attr):
            return False
      return True

    # We just check numeric/string settings for simplicity, this excludes things
    # like learning_rate_fn, optimization_parameters, etc since it's hard to
    # compare them.
    if not _Match(
        tpu_embedding.table_to_config_dict,
        table_to_config_dict,
        ['vocabulary_size', 'dimension', 'combiner'],
    ):
      raise ValueError(
          'table_to_config_dict mismatch. '
          f'Expecting {tpu_embedding.table_to_config_dict}, '
          f'got {table_to_config_dict}'
      )
    if not _Match(
        tpu_embedding.feature_to_config_dict,
        feature_to_config_dict,
        ['table_id', 'max_sequence_length'],
    ):
      raise ValueError(
          'feature_to_config_dict mismatch. '
          f'Expecting {tpu_embedding.feature_to_config_dict}, '
          f'got {feature_to_config_dict}'
      )
    if (
        tpu_embedding.batch_size_per_core * tpu_embedding.num_cores
        != global_batch_size
    ):
      raise ValueError(
          'global_batch_size mismatch. '
          f'batch_size_per_core: {tpu_embedding.batch_size_per_core}, '
          f'num_cores: {tpu_embedding.num_cores}, '
          f'global_batch_size: {global_batch_size}'
      )

  def _CreateLayerVariables(self):
    super()._CreateLayerVariables()
    p = self.params

    # Track which features are intended to be "sequence embeddings".
    self._sequence_features = {}

    if py_utils.IsTpuTraining(p):
      num_cores = self.cluster.params.worker.tpus_per_replica
      global_batch_size = (
          self.params.batch_size * self.cluster.num_splits_per_client
      )

      table_to_config_dict: Dict[str, tpu_embedding_lib.TableConfig] = {}
      feature_to_config_dict: Dict[str, tpu_embedding_lib.FeatureConfig] = {}
      for table in self.tables:
        table_to_config_dict[table.table_name] = table.table_config
        for feature in table.input_keys:
          if table.max_sequence_length > 0:
            self._sequence_features[feature] = True
          feature_to_config_dict[feature] = tpu_embedding_lib.FeatureConfig(
              table.table_name, max_sequence_length=table.max_sequence_length
          )

      tpu_embedding = self._tpu_embedding_collection.tpu_embedding
      if tpu_embedding:
        self._CheckTPUEmbeddingConfig(
            tpu_embedding,
            table_to_config_dict,
            feature_to_config_dict,
            global_batch_size,
        )
        tf.logging.info('TPUEmbedding API singleton already exists, reusing')
        self._tpu_embedding = tpu_embedding
      else:
        device_config = tpu_embedding_lib.DeviceConfig(
            num_cores=num_cores,
            num_hosts=self.params.tables[0].num_tpu_hosts,
            job_name=self.cluster.params.worker.name,
        )
        self._tpu_embedding = tpu_embedding_lib.TPUEmbedding(
            table_to_config_dict,
            feature_to_config_dict,
            global_batch_size,
            mode=tpu_embedding_lib.TRAINING,
            master=None,
            pipeline_execution_with_tensor_core=(
                self.params.pipeline_execution_with_tensor_core
            ),
            partition_strategy=p.partition_strategy,
            device_config=device_config,
        )
        self._tpu_embedding_collection.tpu_embedding = self._tpu_embedding
        self._tpu_embedding_collection.SetGradientMultiplierSchedule(
            self.gradient_multiplier_schedule
        )

  def _TpuEmbLookup(self, ids_map: py_utils.NestedMap) -> py_utils.NestedMap:
    """See base class."""
    task_call_scope = py_utils.GetTaskCallScope()
    activations = self._tpu_embedding_collection.AddActivations(task_call_scope)

    ret = py_utils.NestedMap()
    for k, v in activations.items():
      if ids_map.Get(k) is not None:
        if k in self._sequence_features:
          ret.Set(k, v)
        else:
          # Non-sequence embeddings, we add a "time" dimension of size 1.
          with tf.name_scope(k):
            ret.Set(k, tf.expand_dims(v, axis=1))
    return ret

  def GetIdsMap(self, input_batch: py_utils.NestedMap) -> py_utils.NestedMap:
    """Returns a NestedMap containing the embedding ids to lookup.

    Args:
      input_batch: The input batch to get the embedding id tensors from.

    Returns:
      A py_utils.NestedMap containing the embedding id tensors.
    """
    keys = set().union(*[table.input_keys for table in self.tables])
    return input_batch.GetSlice(keys)

  def GetAuxiliaryVariables(self, input_key: str) -> py_utils.NestedMap:
    """Returns the auxiliary variables associated with the given table.

    Args:
      input_key: A valid input_key as specified in one of the tables.

    Returns:
      The auxiliary variables associated with the embedding table identified by
      input_key.

    Raises:
      ValueError: if input_key is not a valid input key.
    """
    for table in self.tables:
      if input_key in table.input_keys:
        return table.auxiliary_variables
    raise ValueError(f'{input_key} is not a valid input key.')
