# 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.
# ==============================================================================
"""Mid level API for TPU Embeddings With V2 Embedding Accelerator."""

import collections
import copy
import dataclasses
import functools
import hashlib
import operator
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

from absl import logging

from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.tpu.kernels import sparse_core_layout_pb2
from tensorflow.python.checkpoint import saveable_compat
from tensorflow.python.distribute import device_util
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import tpu_strategy
from tensorflow.python.distribute import tpu_util
from tensorflow.python.distribute import tpu_values
from tensorflow.python.distribute import values
from tensorflow.python.distribute import values_util
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import device as tf_device
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 import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_resource_variable_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import summary_ops_v2
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables as tf_variables
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.tpu import _pywrap_sparse_core_layout
from tensorflow.python.tpu import embedding_context_utils as ecu
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_checkpoint_adapter
from tensorflow.python.tpu import tpu_embedding_v3_utils
from tensorflow.python.tpu import tpu_replication
from tensorflow.python.tpu.ops import gen_xla_ops as xla_ops
from tensorflow.python.trackable import base
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.util import compat
from tensorflow.python.util import nest
from tensorflow.python.util import tf_inspect
from tensorflow.python.util.tf_export import tf_export


_PIPELINE_ATTRIBUTE = "_embedding_pipelining"
_PIPELINE_MODE_FORWARD = "forward"
_PIPELINE_MODE_BACKWARD = "backward"
_PIPELINE_MODEL_SEQUENTIAL = "_sequential"
_PARAMETER_NAME = "parameters"


TableConfig = tpu_embedding_v2_utils.TableConfig
FeatureConfig = tpu_embedding_v2_utils.TableConfig
QuantizationConfig = tpu_embedding_v2_utils.QuantizationConfig


@tf_export("tpu.experimental.embedding.SparseCoreEmbeddingConfig")
@dataclasses.dataclass(frozen=True)
class SparseCoreEmbeddingConfig:
  """Config for sparsecore embedding."""

  disable_table_stacking: bool = False
  max_ids_per_chip_per_sample: int = 64
  max_ids_per_table: Optional[Dict[str, int]] = None
  max_unique_ids_per_table: Optional[Dict[str, int]] = None
  allow_id_dropping: bool = False
  initialize_tables_on_host: bool = True
  enable_fast_table_initialization: bool = False


class EmbeddingPipeliningContext(control_flow_ops.ControlFlowContext):
  """Sets the _embedding_pipelining attribute on all ops created in the scope."""

  def __init__(self, mode: str, enable: bool):
    super().__init__()
    self._name = "EmbeddingPipelinigContext"
    self._mode = attr_value_pb2.AttrValue(s=compat.as_bytes(mode))
    self._enable = enable
    recording_summaries = summary_ops_v2.is_recording_summaries()
    if not isinstance(recording_summaries, bool):
      # We can't handle predicate functions at this point. So, we'll ignore the
      # special casing of summary recording because, presumably, this is not
      # a single step loop so pipelining is still valid.
      recording_summaries = False

    if enable and (
        recording_summaries or not ecu.embedding_pipelining_state.enabled
    ):
      # We'll still flag these ops for the SC forward/backward pass, but we'll
      # run them sequentially. This has to be handled in the MLIR passes
      # embedding_pipelining.cc and embedding_sequencing.cc.
      disable_reason = (
          "Summary recording"
          if recording_summaries
          else "_embedding_pipelining_state.enabled = False"
      )
      logging.info("%s detected, disabling pipelining.", disable_reason)
      self._mode = attr_value_pb2.AttrValue(
          s=compat.as_bytes(mode + _PIPELINE_MODEL_SEQUENTIAL)
      )

  def to_control_flow_context_def(
      self, context_def: Any, export_scope: Any = None
  ):
    # pylint: disable=useless-super-delegation
    # The method is required by `ControlFlowContext`.
    super().to_control_flow_context_def(context_def, export_scope)

  def AddOp(self, op: ops.Operation):
    # pylint: disable=protected-access
    if self._enable:
      op._set_attr(_PIPELINE_ATTRIBUTE, self._mode)
    if self._outer_context:
      self._outer_context.AddOp(op)


class TPUEmbeddingShardedSaveable(saveable_object.SaveableObject):
  """Defines how to save and restore a shard of TPUEmbedding sharded variable."""

  def __init__(
      self,
      variable: tf_variables.Variable,
      shard_id: int,
      num_shards: int,
      shard_dim: int,
      name: str,
  ):
    """Init TPUEmbeddingShardedSaveable."""
    self._shard_id = shard_id
    self._variable = variable

    var_offset = [0] * len(variable.shape)
    # NOTE: always assume even sharding
    var_offset[shard_dim] = shard_id * variable.shape[shard_dim]
    fullshape = variable.shape.as_list()
    fullshape[shard_dim] = num_shards * fullshape[shard_dim]
    save_slice_info = tf_variables.Variable.SaveSliceInfo(
        full_name=name,
        full_shape=fullshape,
        var_offset=var_offset,
        var_shape=variable.shape.as_list(),
    )

    spec = saveable_object.SaveSpec(
        tensor=variable.read_value,
        slice_spec=save_slice_info.spec,
        name=name,
        dtype=variable.dtype,
        device=variable.device,
    )
    super().__init__(variable.read_value, [spec], name)

  def restore(
      self,
      restored_tensors: List[tensor.Tensor],
      restored_shapes: List[tensor_shape.TensorShape],
  ) -> Any:
    del restored_shapes
    restored_tensor = restored_tensors[0]

    return values_util.assign_on_device(
        self._variable.device, self._variable, restored_tensor
    )


def _fielddict():
  return dataclasses.field(default_factory=dict)


@dataclasses.dataclass
class TableStacking:
  """Information about how we stack tables."""

  # Indexed by stacked table name:
  stacked_table_to_tables: Dict[str, TableConfig] = _fielddict()
  quantization_configs: Dict[str, QuantizationConfig] = _fielddict()

  # Indexed by table name:
  table_name_to_table: Dict[str, TableConfig] = _fielddict()
  table_to_padding_rows: Dict[str, int] = _fielddict()
  table_to_padding_columns: Dict[str, int] = _fielddict()
  table_to_sample_count: Dict[str, int] = _fielddict()
  table_to_layout: Dict[str, sparse_core_layout_pb2.SparseCoreTableLayout] = (
      _fielddict()
  )
  # Maps table name to (stacked table, row offset, shard rotation)
  table_to_stacked_table_offset: Dict[str, Tuple[str, int, int]] = _fielddict()

  # Indexed by feature_path the key of flat_features:
  feature_to_sample_offset: Dict[str, int] = _fielddict()


@saveable_compat.legacy_saveable_name("")
class TPUEmbeddingShardedVariable(
    tpu_values.TPUVariableMixin, values.DistributedVariable
):
  """A ShardedVariable class for Embedding tables on TPU."""

  def _is_mirrored(self) -> bool:
    return False

  # Only support sharding on the first dimension.
  @property
  def shard_dim(self) -> int:
    return 0

  @property
  def shape(self) -> tensor_shape.TensorShape:
    """Returns the shape of the embedding variable for the current context."""
    local_shape = self._values[0].shape
    global_shape = local_shape.as_list()
    global_shape[self.shard_dim] = global_shape[self.shard_dim] * len(
        self.values
    )
    return tensor_shape.TensorShape(global_shape)

  def _write_object_proto(self, proto, options):
    super()._write_object_proto(proto, options)
    # TODO(b/305882915): Reset the saved model shape to the local shape
    # for backward compatibility of users that directly access the full
    # variable shape as the shape of values.
    proto.variable.shape.CopyFrom(self._values[0].shape.as_proto())

  def _gather_saveables_for_checkpoint(self) -> Dict[str, Callable[..., Any]]:
    """Overrides Trackable method.

    Returns:
      A dictionary mapping attribute names to `SaveableObject` factories.
    """

    def _saveable_factory(name=self._common_name):
      saveables = []
      num_shards = len(self.values)
      for shard_id in range(num_shards):
        saveables.append(
            TPUEmbeddingShardedSaveable(
                self.values[shard_id],
                shard_id,
                num_shards,
                self.shard_dim,
                name,
            )
        )
      return saveables

    return {base.VARIABLE_VALUE_KEY: _saveable_factory}

  def _dense_var_to_tensor(self, dtype=None, name=None, as_ref=False):
    """Converts a variable to a tensor."""
    # pylint: disable=protected-access
    if tpu_util.enclosing_tpu_context() is None:
      return self._values[0].read_value()
    else:
      return self._read_variable_op()

  def read_value(self) -> Any:
    if tpu_util.enclosing_tpu_context() is None:
      raise NotImplementedError(
          "Reading in cross replica mode is not yet supported"
          "for TPUEmbeddingShardedVariable."
      )
    else:
      return self._read_variable_op()

  def assign(
      self,
      value: Any,
      use_locking: bool = False,
      name: Optional[Any] = None,
      read_value: bool = True,
  ) -> Any:
    if tpu_util.enclosing_tpu_context() is None:
      # Running in a host context
      for device in self.distribute_strategy.extended.worker_devices:
        with ops.device(device):
          self.assign_on_device(device, value)
    return tpu_util.make_raw_assign_fn(
        gen_resource_variable_ops.assign_variable_op
    )(
        self,
        value=value,
        use_locking=use_locking,
        name=name,
        read_value=read_value,
    )

  def assign_on_device(self, device, value):
    if self._packed_var is None:
      raise NotImplementedError("Required packed variable support")
    with ops.device(device):
      gen_resource_variable_ops.assign_variable_op(
          resource=self._packed_var.handle, value=value
      )

  def read_from_device(self, device):
    if self._packed_var is None:
      raise NotImplementedError("Required packed variable support")
    with ops.device(device):
      return gen_resource_variable_ops.read_variable_op(
          resource=self._packed_var.handle, dtype=self.dtype
      )


# TODO(pineapplejuice233): Add debug string representation of the class.
PartitionedCsrFormatTensor = collections.namedtuple(
    "PartitionedCsrFormatTensor",
    [
        "row_pointers",
        "sorted_sample_ids",
        "sorted_token_ids",
        "sorted_gains",
        "sample_count",
        "num_minibatches_per_physical_sparse_core",
    ],
)


def _clone_feature_config(feature_config):
  old_to_new_table = {}
  new_features = []

  for old_feature in nest.flatten(feature_config):
    feature = copy.copy(old_feature)
    if feature.table not in old_to_new_table:
      old_to_new_table[feature.table] = copy.copy(feature.table)
    feature.table = old_to_new_table[feature.table]
    new_features.append(feature)

  return nest.pack_sequence_as(feature_config, new_features)


def _stack_tables_with_same_table_dim_and_optimizer(
    table_config: Sequence[TableConfig],
    flat_features: Sequence[Tuple[Any, FeatureConfig]],
    num_partitions: int,
    num_sc_per_partition: int,
    sparse_core_embedding_config: Optional[SparseCoreEmbeddingConfig] = None,
) -> TableStacking:
  """Stack tables with the same table dim and optimizer."""
  logging.info("Number of tables before stacking is %d", len(table_config))
  disable_table_stacking = False
  if sparse_core_embedding_config:
    disable_table_stacking = sparse_core_embedding_config.disable_table_stacking

  if disable_table_stacking:
    logging.warn("Table stacking is disabled.")

  stacker = _pywrap_sparse_core_layout.SparseCoreLayoutStacker(
      num_partitions=num_partitions,
      sparse_cores_per_partition=num_sc_per_partition,
      disable_table_stacking=disable_table_stacking,
  )
  s = TableStacking()
  s.table_name_to_table = {table.name: table for table in table_config}
  table_to_num_samples = {table.name: 0 for table in table_config}
  table_to_num_features = {table.name: 0 for table in table_config}
  for _, feature in flat_features:
    table_to_num_samples[feature.table.name] += functools.reduce(
        operator.mul, feature.output_shape
    )
    table_to_num_features[feature.table.name] += 1

    # First generate stacking for any tables our caller didn't stack for us.
    # Note that we process the tables sorted by name so the ordering is
    # deterministic.
    sorted_tables = sorted(table_config, key=lambda t: t.name)
    for table in sorted_tables:
      if not table.layout:
        # All tables in a stack have to have the same hyperparemeters; this key
        # contains everything we care about. The key is an arbitrary string
        # whose value is not particularly meaningful except that it has to be
        # different if the tables cannot be stacked together.
        #
        # Note that later we rewrite the stack name based on the tables in that
        # stack; this is just a temporary initial name.
        #
        # The key does not need to include the embedding width; that is handled
        # separately.
        key_tuple = (
            # Optimizers don't have a repr but do support hash.
            hash(table.optimizer),
            # Quantization configs don't have a hash but do support repr.
            repr(table.quantization_config),
        )
        key_str = hashlib.sha1(
            repr(key_tuple).encode(),
            usedforsecurity=False,
        ).hexdigest()
        key = "_xxtpuv3internal_" + key_str
        stacker.AddTable(
            table_name=table.name,
            table_height=table.vocabulary_size,
            table_width=table.dim,
            group=key,
            output_samples=table_to_num_samples[table.name],
            num_features=table_to_num_features[table.name],
        )
    # First generate stacking for any tables our caller didn't stack for us.
    # Note that we process the tables sorted by name so the ordering is
    # deterministic.
    # Put the layout information we just computed back into the tables, so we
    # can treat tables whose layouts were given by the caller and tables whose
    # layouts we computed the same.
    for layout in stacker.GetLayouts().tables:
      table = s.table_name_to_table[layout.table_name]
      assert not table.layout  # It's a bug if it was already set.
      table.layout = layout

    # Collect all the layout information from all the tables, whether we just
    # computed it above, or whether the caller passed it as part of the
    # TableConfig:
    tables_by_stack = collections.defaultdict(list)
    for table in sorted_tables:
      layout = table.layout
      assert layout.table_name == table.name
      s.table_to_layout[table.name] = layout
      tables_by_stack[layout.stacked_table_name].append(table)

    for stack_name, tables in tables_by_stack.items():
      s.quantization_configs[stack_name] = tables[0].quantization_config
      s.stacked_table_to_tables[stack_name] = tables

      logging.vlog(1, "Stacked table name: %s", stack_name)
      for table in tables:
        layout = table.layout
        logging.vlog(
            1,
            "  Table %s: offset %d, rotation %d",
            table.name,
            layout.sparse_core_shard_row_offset,
            layout.sparse_core_shard_rotation,
        )
        s.table_to_stacked_table_offset[table.name] = (
            stack_name,
            layout.sparse_core_shard_row_offset
            * num_partitions
            * num_sc_per_partition,
            layout.sparse_core_shard_rotation,
        )
        # Update dimensions in the table to the padded dimensions.
        table.vocabulary_size = layout.unsharded_padded_shape[0]
        table.dim = layout.unsharded_padded_shape[1]
        s.table_to_padding_rows[table.name] = (
            layout.unsharded_padded_shape[0] - layout.unsharded_shape[0]
        )
        s.table_to_padding_columns[table.name] = (
            layout.unsharded_padded_shape[1] - layout.unsharded_shape[1]
        )

    logging.info(
        "Number of tables after stacking is %d.",
        len(s.stacked_table_to_tables),
    )

    s.table_to_sample_count = {
        table_name: 0 for table_name in s.stacked_table_to_tables
    }
    for feature_path, feature in flat_features:
      stacked_table_name = s.table_to_stacked_table_offset[feature.table.name][
          0
      ]
      s.feature_to_sample_offset[feature_path] = s.table_to_sample_count[
          stacked_table_name
      ]
      s.table_to_sample_count[stacked_table_name] += functools.reduce(
          operator.mul, feature.output_shape
      )

    return s


# TODO(b/233952762): Add tests of this version of the mid-level API.
@tf_export("tpu.experimental.embedding.TPUEmbeddingV2")
class TPUEmbeddingV2(tpu_embedding_base.TPUEmbeddingBase):
  """The TPUEmbedding mid level API running on TPU with sparse core accelerator."""

  DEFAULT_MAX_IDS_PER_TABLE = 256
  DEFAULT_MAX_UNIQUE_IDS_PER_TABLE = 256

  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
      pipeline_execution_with_tensor_core: bool = False,
      sparse_core_embedding_config: Optional[SparseCoreEmbeddingConfig] = None,
  ):
    """Creates the TPUEmbeddingV2 mid level API object.

    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.
      pipeline_execution_with_tensor_core: If True, the TPU embedding
        computations will overlap with the TensorCore computations (and hence
        will be one step old). Set to True for improved performance.
      sparse_core_embedding_config: Configs for sparse core embedding including
        settings for table stacking, input feature static buffer size etc.

    Raises:
      ValueError: If optimizer is not one of tf.tpu.experimental.embedding.(SGD,
      Adam or Adagrad) or None when created under a TPUStrategy.
      RuntimeError: If not created under TPUStrategy.
    """
    # We do a clone on the feature_config here as we will alter settings in it
    # and we don't want the user to see these. We can't just use clone here
    # as we need to maintain some object relationships.
    super().__init__(_clone_feature_config(feature_config), optimizer)
    self._strategy = distribute_lib.get_strategy()
    if not isinstance(
        self._strategy, (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV2)
    ):
      raise RuntimeError(
          "TPUEmbeddingV2 should be created under TPUStrategy but found {}."
          .format(self._strategy)
      )

    # TODO(pineapplejuice233): Remove this once weight decay is supported.
    for table in self._table_config:
      if (
          table.optimizer.weight_decay_factor is not None
          or table.optimizer.multiply_weight_decay_factor_by_learning_rate
          is not None
      ):
        raise NotImplementedError(
            "weight_decay_factor and"
            " multiply_weight_decay_factor_by_learning_rate are not supported"
            f" yet. But found in table {table.name} setting."
        )

    self._num_sc_per_chip = (
        self._strategy.extended.tpu_hardware_feature.num_embedding_devices_per_chip
    )
    if self._num_sc_per_chip == 0:
      logging.warning(
          "No embedding devices per chip info is found. Using 4 as the default"
          " value for SparseCore."
      )
      self._num_sc_per_chip = 4

    self._num_sc_shards = (
        self._strategy.num_replicas_in_sync * self._num_sc_per_chip
    )

    # We need this in multiple places, so avoid flattening multiple times.
    # This order will also be used when stacking features.
    self._flat_features = nest.flatten_with_joined_string_paths(
        self._feature_config
    )

    if sparse_core_embedding_config is None:
      self._sparse_core_embedding_config = SparseCoreEmbeddingConfig()
      logging.warning(
          "SparseCoreEmbeddingConfig is not provided. Using default values %s",
          self._sparse_core_embedding_config,
      )
    else:
      self._sparse_core_embedding_config = sparse_core_embedding_config

    self._s = _stack_tables_with_same_table_dim_and_optimizer(
        self._table_config,
        self._flat_features,
        self._strategy.num_replicas_in_sync,
        self._num_sc_per_chip,
        self._sparse_core_embedding_config,
    )

    self._table_name_to_table = self._s.table_name_to_table
    self._stacked_table_to_tables = self._s.stacked_table_to_tables
    self._table_to_padding_columns = self._s.table_to_padding_columns
    self._table_to_padding_rows = self._s.table_to_padding_rows
    self._table_to_stacked_table_offset = self._s.table_to_stacked_table_offset
    self._table_to_sample_count = self._s.table_to_sample_count
    self._feature_to_sample_offset = self._s.feature_to_sample_offset
    self._quantization_configs = self._s.quantization_configs

    # These hyperparameters will be provided by the FDO. Currently hardcode
    # here just for testing.
    self.max_ids_per_chip_per_sample = (
        self._sparse_core_embedding_config.max_ids_per_chip_per_sample
    )
    self.max_minibatches_per_sc = 64

    self._table_to_max_ids_per_sparse_core = {}
    self._table_to_max_unique_ids_per_sparse_core = {}

    self._update_sparse_core_buffer_size_after_table_stacking()

    self._pipelining = pipeline_execution_with_tensor_core
    self._initializers_shard_info_by_default = True

  def _compute_sc_shard_info(
      self,
      table: TableConfig,
      partition_shape: tuple[int, int],
      partition_offset: List[int],
      total_vocab_size: int,
      sc_idx: int,
  ) -> base.ShardInfo:
    # Scale the partition to get sizes for the current table,
    # then select this sc shard.
    sc_shard_size = (
        table.vocabulary_size
        * partition_shape[0]
        // total_vocab_size
        // self._num_sc_per_chip
    )
    sc_shard_offset = (
        table.vocabulary_size
        * partition_offset[0]
        // total_vocab_size
    ) + sc_idx * sc_shard_size

    return base.ShardInfo([sc_shard_size, table.dim], [sc_shard_offset, 0])

  def _compute_sc_shard_idx_and_offset(
      self,
      table_name: str,
      shard_info: base.ShardInfo
  ) -> tuple[int, int]:
    tpu_devices = self._strategy.extended._tpu_devices  # pylint:disable=protected-access
    num_replicas, num_cores_per_replica = tpu_devices.shape
    num_devices = num_replicas * num_cores_per_replica

    shift = self._s.table_to_stacked_table_offset[table_name][2]
    shard_index = shard_info.offset[0] // shard_info.shape[0]
    # Rotate the shards.
    shard_index = (shard_index - shift) % self._num_sc_shards
    num_sc = num_devices * self._num_sc_per_chip

    return shard_index, num_sc

  def _update_sparse_core_buffer_size_after_table_stacking(self):
    """Update the sparse core buffer size after table stacking."""
    for table_name in self._stacked_table_to_tables:
      if (
          self._sparse_core_embedding_config.max_ids_per_table is None
          or table_name
          not in self._sparse_core_embedding_config.max_ids_per_table
      ):
        logging.warning(
            "Table %s is not found in max_ids_per_table provided by"
            " SparseCoreEmbeddingConfig. Using default value 256.",
            table_name,
        )
        self._table_to_max_ids_per_sparse_core[table_name] = (
            self.DEFAULT_MAX_IDS_PER_TABLE
        )
      else:
        self._table_to_max_ids_per_sparse_core[table_name] = (
            self._sparse_core_embedding_config.max_ids_per_table[table_name]
        )
      if (
          self._sparse_core_embedding_config.max_unique_ids_per_table is None
          or table_name
          not in self._sparse_core_embedding_config.max_unique_ids_per_table
      ):
        logging.warning(
            (
                "Table %s is not found in max_unique_ids_per_table provided by"
                " SparseCoreEmbeddingConfig. Using default value 256."
            ),
            table_name,
        )
        self._table_to_max_unique_ids_per_sparse_core[table_name] = (
            self.DEFAULT_MAX_UNIQUE_IDS_PER_TABLE
        )
      else:
        self._table_to_max_unique_ids_per_sparse_core[table_name] = (
            self._sparse_core_embedding_config.max_unique_ids_per_table[
                table_name
            ]
        )

  @property
  def embedding_tables(
      self,
  ) -> Dict[str, tf_variables.Variable]:
    """Returns a dict of embedding tables, keyed by stacked table name."""
    self._maybe_build()
    # Only return the tables and not the slot variables.
    return {
        stacked_table_name: self._variables[stacked_table_name]["parameters"]
        for stacked_table_name in self._stacked_table_to_tables
    }

  @property
  def embedding_table_shards(
      self,
  ) -> Dict[tpu_embedding_v2_utils.TableConfig, List[tf_variables.Variable]]:
    """Returns a dict of embedding tables, keyed by `TableConfig`."""
    self._maybe_build()

    # This reflects the device assignment used by the TPU Strategy.
    ordered_devices = []
    for devices in self._strategy.extended._tpu_devices:  # pylint: disable=protected-access
      ordered_devices.extend(devices)

    table_shards = {
        name: [
            (device, var.read_from_device(device)) for device in ordered_devices
        ]
        for name, var in self.embedding_tables.items()
    }

    return table_shards

  @property
  def embedding_layouts(
      self,
  ) -> Dict[str, sparse_core_layout_pb2.SparseCoreTableLayout]:
    """Returns how the tables are laid out in the variables.

    The SparseCoreTableLayout describes how a table is stored in its internal
    state. You need this only if you need to pull apart the internal state.
    """
    return self._s.table_to_layout

  @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

  def _variable_creator(
      self,
      next_creator: Callable[..., tf_variables.Variable],
      **kwargs,
  ) -> Callable[..., Any]:
    return make_sharded_variable_creator(self._strategy)(next_creator, **kwargs)

  def _wrap_initializer(self, initializer):
    """Wraps the initializer to ensure that it is sharding aware."""
    arg_spec = tf_inspect.getfullargspec(initializer)
    sharding_aware = (
        "shard_info" in arg_spec.args or "shard_info" in arg_spec.kwonlyargs
    )
    if sharding_aware:
      return initializer

    def wrapper(shape, dtype, shard_info: base.ShardInfo):
      del shape
      return initializer(shape=shard_info.shape, dtype=dtype)

    return wrapper

  def _get_shard_info_for_table(
      self,
      table: TableConfig,
      device_idx: int,
      num_devices: int,
  ) -> base.ShardInfo:
    """Returns the shard info for the given table."""
    device_shard_len = table.vocabulary_size // num_devices
    sc_shard_len = device_shard_len // self._num_sc_per_chip
    shift = self._s.table_to_stacked_table_offset[table.name][2]
    shift_rows = shift * sc_shard_len
    row_offset = (
        device_idx * device_shard_len - shift_rows
    ) % table.vocabulary_size
    return base.ShardInfo(
        (device_shard_len, table.dim),
        (row_offset, 0),
    )

  def _initialize_stacked_table_for_device(
      self,
      stacked_tables: List[TableConfig],
      device_idx: int,
      num_devices: int,
  ) -> dict[str, tensor.Tensor]:
    """Initializes the stacked tables and slots shards for a single device."""
    table_dim = stacked_tables[0].dim
    variable_dtype = stacked_tables[0].dtype
    slot_shards = {}
    replicated_slot_shards = {}
    for table in stacked_tables:
      if table.optimizer is not None:
        initializers = {
            slot: initializer
            for slot, initializer in zip(
                table.optimizer._slot_names(),  # pylint:disable=protected-access
                table.optimizer._slot_initializers(),  # pylint:disable=protected-access
            )
        }
      else:
        initializers = {}

      initializers[_PARAMETER_NAME] = table.initializer

      shard_info = self._get_shard_info_for_table(
          table, device_idx, num_devices
      )
      for slot, initializer in initializers.items():
        device_shard = self._wrap_initializer(initializer)(
            shape=(table.vocabulary_size, table.dim),
            dtype=variable_dtype,
            shard_info=shard_info,
        )
        sc_shards = array_ops.reshape(
            device_shard, (self._num_sc_per_chip, -1, table.dim)
        )
        slot_shards.setdefault(slot, []).append(sc_shards)

      # Keep track of the replicated slot variable shards required by the custom
      # combiner.
      if isinstance(table.combiner, tpu_embedding_v2_utils.CustomCombiner):
        combiner_weights_shape = (table.combiner.num_weights,)
        replicated_slot_shards = {
            slot: self._wrap_initializer(initializer)(
                shape=combiner_weights_shape,
                dtype=dtypes.float32,
                shard_info=base.ShardInfo(
                    shape=combiner_weights_shape,
                    offset=device_idx,
                ),
            )
            for slot, initializer in zip(
                table.combiner._slot_names(),  # pylint:disable=protected-access
                table.combiner._slot_initializers(),  # pylint:disable=protected-access
                strict=True,  # number of names and initializers must match.
            )
        }

    slots = {}
    for slot, shards in slot_shards.items():
      concated_shards = array_ops.concat(shards, axis=1)
      concated_shards = array_ops.reshape(concated_shards, (-1, table_dim))
      slots[slot] = concated_shards
    slots.update(replicated_slot_shards)
    return slots

  def _initialize_stacked_table_all_devices(
      self,
      stacked_tables: List[TableConfig],
      stacked_table_name: str,
  ) -> dict[str, tensor.Tensor]:
    """Initializes the stacked tables and slots shards for all devices."""
    table_dim = stacked_tables[0].dim

    tpu_devices = self._strategy.extended._tpu_devices  # pylint:disable=protected-access
    num_devices = tpu_devices.size
    num_scs = num_devices * self._num_sc_per_chip

    for table in stacked_tables:
      if table.vocabulary_size % num_scs != 0:
        raise ValueError(
            "Only evenly sharding across devices is currently supported. "
            f"Got vocab size {table.vocabulary_size} and {num_scs} sparse cores"
        )
      if table.dim != table_dim:
        raise ValueError(
            "Only tables with the same dimension are currently supported. "
            f"Got table {table.name} with dimension {table.dim} in stacked "
            f"table {stacked_table_name} with dimension {table_dim}"
        )

    parameters = {}
    for idx, device in enumerate(tpu_devices.flatten()):
      with ops.device(device):
        device_parameters = self._initialize_stacked_table_for_device(
            stacked_tables, idx, num_devices
        )
        for slot, shard in device_parameters.items():
          parameters.setdefault(slot, []).append(shard)
    return parameters

  @def_function.function
  def _batch_initialize_tables_fn(
      self,
  ) -> Dict[str, Dict[str, List[tensor.Tensor]]]:
    tensors = {}
    for stacked_table_name, tables in self._s.stacked_table_to_tables.items():
      tensors[stacked_table_name] = self._initialize_stacked_table_all_devices(
          tables, stacked_table_name=stacked_table_name
      )
    return tensors

  def _batch_initialize_tables(
      self,
  ) -> Dict[str, Dict[str, List[tensor.Tensor]]]:
    logging.info("Batch initializing embedding tables.")
    tpu_devices = self._strategy.extended._tpu_devices  # pylint:disable=protected-access
    with ops.device(device_util.get_host_for_device(tpu_devices[0][0])):
      return self._batch_initialize_tables_fn()

  def _create_variables(
      self,
      stacked_tables: List[tpu_embedding_v2_utils.TableConfig],
      stacked_table_name: str,
      initialized_tensors: dict[str, dict[str, List[tensor.Tensor]]],
  ) -> Dict[str, tf_variables.Variable]:
    """Create all variables including table variables and slot variables."""
    total_vocab_size = sum([table.vocabulary_size for table in stacked_tables])
    table_dim = stacked_tables[0].dim
    variable_shape = (total_vocab_size, table_dim)
    variable_dtype = dtypes.float32
    optimizer = stacked_tables[0].optimizer

    def table_initialize_fn_non_batch(shape, dtype, shard_info=None):
      # If enable fast table initialization, we will initialize the table
      # directly on the device and use the initializer from the first table.
      if self._sparse_core_embedding_config.enable_fast_table_initialization:
        return stacked_tables[0].initializer(
            shape=(shard_info.shape[0], stacked_tables[0].dim),
            dtype=dtype,
        )

      # Concat all the tables along the first axis.
      concat_tensors = []

      # Temporary patch, we need to initialize tables with the SC level
      # sharding. Note that we need to ensure that the vocab size is divisible
      # by the global number of SC.
      for i in range(self._num_sc_per_chip):
        # Each underlying table has column lookups rotated by 1 to avoid hot
        # spots on core 0 for id=0. We shift the initializer as well to help
        # with comparisons against CPU.
        for table in stacked_tables:
          arg_spec = tf_inspect.getfullargspec(table.initializer)
          sharding_aware = (
              "shard_info" in arg_spec.args
              or "shard_info" in arg_spec.kwonlyargs
          )

          if shard_info:
            sc_shard_info = self._compute_sc_shard_info(
                table,
                shard_info.shape,
                shard_info.offset,
                total_vocab_size,
                i,
            )
            if not sharding_aware:
              shard_index, shard_offset = self._compute_sc_shard_idx_and_offset(
                  table.name, sc_shard_info
              )
              sc_shard = table.initializer(
                  shape=(table.vocabulary_size, table.dim), dtype=dtype
              )[shard_index::shard_offset, :]
            else:
              sc_shard = table.initializer(
                  shape=(table.vocabulary_size, table.dim),
                  dtype=dtype,
                  shard_info=sc_shard_info,
              )
          else:
            sc_shard = table.initializer(
                shape=(
                    (table.vocabulary_size * shape[0])
                    // total_vocab_size
                    // self._num_sc_per_chip,
                    table.dim,
                ),
                dtype=dtype,
            )
          concat_tensors.append(sc_shard)
      return array_ops.concat(concat_tensors, axis=0)

    def batch_initialize_fn(name, shard_info, replicated=False):
      if not initialized_tensors:
        initialized_tensors.update(self._batch_initialize_tables())
      if not replicated:
        # This is the path for the table and optimizer slot variables (which
        # currently have the same shape as the table).
        shard_id = shard_info.offset[0] // shard_info.shape[0]
      else:
        # This is the path for the combiner slot variable (1D and is of size
        # combiner.num_weights). The device ID on which this shard is
        # initialized is passed via the first element of shard_info.offset.
        shard_id = shard_info.offset[0]
      res = initialized_tensors[stacked_table_name][name][shard_id]
      assert shard_info.shape == res.shape
      return res

    def table_initialize_fn(shape, dtype, shard_info=None):
      if not shard_info or not self._initializers_shard_info_by_default:
        # This is the legacy loading path.
        return table_initialize_fn_non_batch(shape, dtype, shard_info)
      return batch_initialize_fn(_PARAMETER_NAME, shard_info)

    def slot_initialize_fn(
        name, shape, dtype, initializer, shard_info=None, replicated=False
    ):
      if not shard_info or not self._initializers_shard_info_by_default:
        # This is the legacy loading path.
        return initializer(shape, dtype)
      return batch_initialize_fn(name, shard_info, replicated)

    def getter(name, shape, dtype, initializer, trainable):
      initial_value = functools.partial(initializer, shape=shape, dtype=dtype)
      # _add_variable_with_custom_getter clears the shape sometimes, so we
      # take the global shape from outside the getter.
      return tf_variables.Variable(
          name=name,
          initial_value=initial_value,
          shape=shape,
          dtype=dtype,
          trainable=trainable,
      )

    def variable_creator(name, initializer, shape, dtype):
      # 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=shape,
          dtype=dtype,
          getter=getter,
          trainable=False,
      )

    with variable_scope.variable_creator_scope(self._variable_creator):
      parameters = variable_creator(
          stacked_table_name,
          table_initialize_fn,
          variable_shape,
          variable_dtype,
      )

    def slot_creator(name, initializer, shape, dtype):
      return variable_creator(
          stacked_table_name + "/" + name, initializer, shape, dtype
      )

    slot_vars = {}
    if optimizer is not None:
      with variable_scope.variable_creator_scope(self._variable_creator):
        slot_vars = optimizer._create_slots(  # pylint: disable=protected-access
            parameters,
            variable_creator=functools.partial(
                slot_creator, shape=variable_shape, dtype=variable_dtype
            ),
            initializer_wrapper=lambda slot, initializer: functools.partial(
                slot_initialize_fn,
                name=slot,
                initializer=initializer,
                replicated=False,
            ),
        )

    slot_vars[_PARAMETER_NAME] = parameters
    return slot_vars

  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.
    """
    variables = {}
    # Batch initialized table and slot variables. It will be lazily initialized
    # when it is first used inside the variable initializers.
    table_tensors = {}
    for stacked_table_name, tables in self._s.stacked_table_to_tables.items():
      variables[stacked_table_name] = self._create_variables(
          tables,
          stacked_table_name=stacked_table_name,
          initialized_tensors=table_tensors,
      )
    return variables

  def _track_restore_info_for_cpu(self) -> None:
    layouts = sparse_core_layout_pb2.SparseCoreTableLayouts()
    layouts.tables.extend(self.embedding_layouts.values())
    logging.info(
        "Adding sparse core layouts for %s tables", len(layouts.tables)
    )
    with ops.device("/cpu:0"):
      self._track_trackable(
          tpu_embedding_v3_utils.SparseCoreLayoutsTrackable(
              constant_op.constant(
                  layouts.SerializeToString(), dtype=dtypes.string
              )
          ),
          tpu_embedding_v3_utils.SPARSECORE_LAYOUTS_CHECKPOINT_KEY,
      )

  def _checkpoint_adapter(self, path):
    # The TPUEmbedding may need to reshard checkpoint values during restore.
    return tpu_embedding_v3_checkpoint_adapter.TpuEmbeddingV3CheckpointAdapter.create_from_checkpoint(
        path
    )

  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 build(self):
    """Create variables and slots variables for TPU embeddings."""
    if self._built:
      return
    self._variables = self._create_variables_and_slots()
    self._track_restore_info_for_cpu()
    self._built = True
    logging.info("TPUEmbedding built.")

  def apply_gradients(
      self,
      gradients: Any,
      preserved_outputs: Dict[str, PartitionedCsrFormatTensor],
  ):
    """Applies the gradient update to the embedding tables.

    If a gradient of `None` is passed in any position of the nested structure,
    then a gradient update with a zero gradient is applied for that feature.
    For optimizers like SGD or Adagrad, this is the same as applying no update
    at all. For lazy Adam and other sparsely applied optimizers with decay,
    ensure you understand the effect of applying a zero gradient.

    Args:
      gradients: A nested structure of gradients, with structure matching the
        `feature_config` passed to this object.
      preserved_outputs: A dicts of PartitionedCsrFormatTensor, coming from the
        second output of the embedding lookup call.

    Raises:
      RuntimeError: if not built.
      ValueError: If a non-`tf.Tensor` non-`None` gradient is passed in, or a
        `tf.Tensor` of the incorrect shape is passed in. Also if
        the size of any sequence in `gradients` does not match corresponding
        sequence in `feature_config`.
      TypeError: If the type of any sequence in `gradients` does not match
        corresponding sequence in `feature_config`.
    """
    if not self._built:
      raise RuntimeError(
          "apply_gradients called on unbuilt TPUEmbeddingV2 object. Please"
          " either call the embedding lookup method first or manually call the"
          " build method."
      )
    nest.assert_same_structure(self._feature_config, gradients)

    # Note that stacking gradients is placed on the core of the trianing step
    # to reduce the number of input/output arguments of the training loop during
    # pipelining.
    gradients = self._stack_gradients(gradients)

    context = EmbeddingPipeliningContext(
        _PIPELINE_MODE_BACKWARD, self._pipelining
    )
    context.Enter()

    def _wrap_param(param, dtype=dtypes.float32):
      if callable(param):
        param = math_ops.cast(param(), dtype=dtype)
      return ops.convert_to_tensor(param, dtype=dtype)

    # Take num_minibatches_per_physical_sparse_core from any table as
    # they are the same across tables.
    num_minibatches_per_physical_sparse_core = list(preserved_outputs.values())[
        0
    ].num_minibatches_per_physical_sparse_core

    for table_name in self._stacked_table_to_tables:
      gradient = gradients[table_name]
      partitioned_tensor = preserved_outputs[table_name]

      table = self.variables[table_name]["parameters"]
      optimizer = self._stacked_table_to_tables[table_name][0].optimizer
      if isinstance(optimizer, tpu_embedding_v2_utils.SGD):
        updated_embedding_table = xla_ops.xla_sparse_dense_matmul_grad_with_sgd_and_static_buffer_size(
            row_pointers=partitioned_tensor.row_pointers,
            sorted_sample_ids=partitioned_tensor.sorted_sample_ids,
            sorted_token_ids=partitioned_tensor.sorted_token_ids,
            sorted_gains=partitioned_tensor.sorted_gains,
            activation_gradients=gradient,
            learning_rate=_wrap_param(optimizer.learning_rate),
            embedding_table=table.read_value(),
            num_minibatches_per_physical_sparse_core=num_minibatches_per_physical_sparse_core,
            max_ids_per_sparse_core=self._table_to_max_ids_per_sparse_core[
                table_name
            ],
            max_unique_ids_per_sparse_core=self._table_to_max_unique_ids_per_sparse_core[
                table_name
            ],
            table_name=table_name,
        )
        table.assign(updated_embedding_table)
      elif isinstance(optimizer, tpu_embedding_v2_utils.Adagrad):
        accumulators = self.variables[table_name]["accumulators"]
        updated_embedding_table, updated_accumulator = (
            xla_ops.xla_sparse_dense_matmul_grad_with_adagrad_and_static_buffer_size(
                row_pointers=partitioned_tensor.row_pointers,
                sorted_sample_ids=partitioned_tensor.sorted_sample_ids,
                sorted_token_ids=partitioned_tensor.sorted_token_ids,
                sorted_gains=partitioned_tensor.sorted_gains,
                activation_gradients=gradient,
                learning_rate=_wrap_param(optimizer.learning_rate),
                embedding_table=table.read_value(),
                accumulator=accumulators.read_value(),
                num_minibatches_per_physical_sparse_core=num_minibatches_per_physical_sparse_core,
                max_ids_per_sparse_core=self._table_to_max_ids_per_sparse_core[
                    table_name
                ],
                max_unique_ids_per_sparse_core=self._table_to_max_unique_ids_per_sparse_core[
                    table_name
                ],
                table_name=table_name,
            )
        )
        accumulators.assign(updated_accumulator)
        table.assign(updated_embedding_table)
      elif isinstance(optimizer, tpu_embedding_v2_utils.AdagradMomentum):
        accumulators = self.variables[table_name]["accumulators"]
        momenta = self.variables[table_name]["momenta"]
        updated_embedding_table, updated_accumulator, updated_momenta = (
            xla_ops.xla_sparse_dense_matmul_grad_with_adagrad_momentum_and_static_buffer_size(
                row_pointers=partitioned_tensor.row_pointers,
                sorted_sample_ids=partitioned_tensor.sorted_sample_ids,
                sorted_token_ids=partitioned_tensor.sorted_token_ids,
                sorted_gains=partitioned_tensor.sorted_gains,
                activation_gradients=gradient,
                learning_rate=_wrap_param(optimizer.learning_rate),
                embedding_table=table.read_value(),
                accumulator=accumulators.read_value(),
                momenta=momenta.read_value(),
                num_minibatches_per_physical_sparse_core=num_minibatches_per_physical_sparse_core,
                use_nesterov=optimizer.use_nesterov,
                exponent=optimizer.exponent,
                beta1=optimizer.momentum,
                beta2=optimizer.beta2,
                epsilon=optimizer.epsilon,
                max_ids_per_sparse_core=self._table_to_max_ids_per_sparse_core[
                    table_name
                ],
                max_unique_ids_per_sparse_core=self._table_to_max_unique_ids_per_sparse_core[
                    table_name
                ],
                table_name=table_name,
            )
        )
        momenta.assign(updated_momenta)
        accumulators.assign(updated_accumulator)
        table.assign(updated_embedding_table)
      elif isinstance(optimizer, tpu_embedding_v2_utils.Adam):
        momenta = self.variables[table_name]["momenta"]
        velocity = self.variables[table_name]["velocities"]
        updated_embedding_table, updated_momenta, updated_velocity = (
            xla_ops.xla_sparse_dense_matmul_grad_with_adam_and_static_buffer_size(
                row_pointers=partitioned_tensor.row_pointers,
                sorted_sample_ids=partitioned_tensor.sorted_sample_ids,
                sorted_token_ids=partitioned_tensor.sorted_token_ids,
                sorted_gains=partitioned_tensor.sorted_gains,
                activation_gradients=gradient,
                learning_rate=_wrap_param(optimizer.learning_rate),
                embedding_table=table.read_value(),
                momenta=momenta.read_value(),
                velocity=velocity.read_value(),
                num_minibatches_per_physical_sparse_core=num_minibatches_per_physical_sparse_core,
                use_sum_inside_sqrt=optimizer.sum_inside_sqrt,
                beta1=optimizer.beta_1,
                beta2=optimizer.beta_2,
                epsilon=optimizer.epsilon,
                max_ids_per_sparse_core=self._table_to_max_ids_per_sparse_core[
                    table_name
                ],
                max_unique_ids_per_sparse_core=self._table_to_max_unique_ids_per_sparse_core[
                    table_name
                ],
                table_name=table_name,
            )
        )
        velocity.assign(updated_velocity)
        momenta.assign(updated_momenta)
        table.assign(updated_embedding_table)
      elif isinstance(optimizer, tpu_embedding_v2_utils.FTRL):
        accumulators = self.variables[table_name]["accumulators"]
        linears = self.variables[table_name]["linears"]
        (updated_table_tensor, updated_accum_tensor, updated_linear_tensor) = (
            xla_ops.xla_sparse_dense_matmul_grad_with_ftrl_and_static_buffer_size(
                row_pointers=partitioned_tensor.row_pointers,
                sorted_sample_ids=partitioned_tensor.sorted_sample_ids,
                sorted_token_ids=partitioned_tensor.sorted_token_ids,
                sorted_gains=partitioned_tensor.sorted_gains,
                activation_gradients=gradient,
                learning_rate=_wrap_param(optimizer.learning_rate),
                embedding_table=table.read_value(),
                accumulator=accumulators.read_value(),
                linear=linears.read_value(),
                num_minibatches_per_physical_sparse_core=num_minibatches_per_physical_sparse_core,
                multiply_linear_by_learning_rate=optimizer.multiply_linear_by_learning_rate,
                beta=optimizer.beta,
                learning_rate_power=optimizer.learning_rate_power,
                l1_regularization_strength=optimizer.l1_regularization_strength,
                l2_regularization_strength=optimizer.l2_regularization_strength,
                max_ids_per_sparse_core=self._table_to_max_ids_per_sparse_core[
                    table_name
                ],
                max_unique_ids_per_sparse_core=self._table_to_max_unique_ids_per_sparse_core[
                    table_name
                ],
                table_name=table_name,
            )
        )
        linears.assign(updated_linear_tensor)
        accumulators.assign(updated_accum_tensor)
        table.assign(updated_table_tensor)
      else:
        raise ValueError("Unsupported optimizer in minibatching mode.")

    context.Exit()

  def __call__(
      self, features: Any, weights: Optional[Any] = None
  ) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]:
    """Call the mid level api to do embedding lookup."""
    return self.embedding_lookup(features, weights)

  # TODO(pineapplejuice233): Duplicated helper function from tpu_embedding_v2.py. Remove
  # this once this file is open souced.
  def _raise_error_for_incorrect_control_flow_context(self):
    """Raises an error if we are not in the TPUReplicateContext."""
    # Do not allow any XLA control flow (i.e. control flow in between a
    # TPUStrategy's run call and the call to this function), as we can't
    # extract the enqueue from the head when in XLA control flow.
    graph = ops.get_default_graph()
    in_tpu_ctx = False
    while graph is not None:
      ctx = graph._get_control_flow_context()  # pylint: disable=protected-access
      while ctx is not None:
        if isinstance(ctx, tpu_replication.TPUReplicateContext):
          in_tpu_ctx = True
          break
        ctx = ctx.outer_context
      if in_tpu_ctx:
        break
      graph = getattr(graph, "outer_graph", None)
    if graph != ops.get_default_graph() and in_tpu_ctx:
      raise RuntimeError(
          "Current graph {} does not match graph which contains "
          "TPUReplicateContext {}. This is most likely due to the fact that "
          "enqueueing embedding data is called inside control flow or a "
          "tf.function inside `strategy.run`. This is not supported because "
          "outside compilation fails to extract the enqueue ops as the head of "
          "a computation.".format(ops.get_default_graph(), graph)
      )
    return in_tpu_ctx

  @classmethod
  def compute_sparse_core_stats(
      cls,
      features: Any,
      feature_config: Union[FeatureConfig, Iterable],  # pylint:disable=g-bare-generic
      num_tpu_chips: int,
      num_sc_per_chip: int = 4,
      optimizer: Optional[tpu_embedding_v2_utils._Optimizer] = None,  # pylint:disable=protected-access
      sparse_core_embedding_config: Optional[SparseCoreEmbeddingConfig] = None,
  ) -> Tuple[Any, Any]:
    """Computes the max_ids/unique ids settings from the input features."""
    copy_feature_config = _clone_feature_config(feature_config)
    table_config_list = list(
        {feature.table for feature in nest.flatten(copy_feature_config)}
    )

    for table in table_config_list:
      if table.optimizer is None:
        table.optimizer = optimizer

    flat_features = nest.flatten_with_joined_string_paths(copy_feature_config)

    s = _stack_tables_with_same_table_dim_and_optimizer(
        table_config_list,
        flat_features,
        num_tpu_chips,
        num_sc_per_chip,
        sparse_core_embedding_config,
    )

    flat_inputs = nest.flatten(features)

    # First process them to be COO tensors.
    table_to_list_of_coos = (
        TPUEmbeddingV2._preprocess_inputs_and_weights_to_list_of_coo_tensors(
            flat_inputs=flat_inputs,
            flat_weights=[None] * len(flat_inputs),
            flat_features=flat_features,
            stacked_table_to_tables=s.stacked_table_to_tables,
            table_to_stacked_table_offset=s.table_to_stacked_table_offset,
            feature_to_sample_offset=s.feature_to_sample_offset,
            num_sc_per_chip=num_sc_per_chip,
            stacked_table_to_sample_count=s.table_to_sample_count,
            num_sc_shards=num_sc_per_chip * num_tpu_chips,
        )
    )

    table_to_max_ids_per_sparse_core = {
        table_name: 0 for table_name in s.stacked_table_to_tables
    }
    table_to_max_unique_ids_per_sparse_core = {
        table_name: 0 for table_name in s.stacked_table_to_tables
    }

    for table_name in s.stacked_table_to_tables:
      feature_width = s.stacked_table_to_tables[table_name][0].dim

      total_vocab_size = sum([
          table.vocabulary_size
          for table in s.stacked_table_to_tables[table_name]
      ])
      for i in range(num_sc_per_chip):
        row_ids_list = table_to_list_of_coos[table_name][0][i]
        col_ids_list = table_to_list_of_coos[table_name][1][i]
        gains_list = table_to_list_of_coos[table_name][2][i]
        sample_count_list = table_to_list_of_coos[table_name][3]
        col_offset_list = table_to_list_of_coos[table_name][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,
            col_ids_list=col_ids_list,
            gains_list=gains_list,
            sample_count_list=sample_count_list,
            col_offset_list=col_offset_list,
            num_replica=num_tpu_chips,
            table_vocab_size=total_vocab_size,
            feature_width=feature_width,
            num_sc_per_chip=num_sc_per_chip,
            table_name=table_name,
        )

        table_to_max_ids_per_sparse_core[table_name] = math_ops.maximum(
            table_to_max_ids_per_sparse_core[table_name],
            max_ids_per_sparse_core,
        )
        table_to_max_unique_ids_per_sparse_core[table_name] = math_ops.maximum(
            table_to_max_unique_ids_per_sparse_core[table_name],
            max_unique_ids_per_sparse_core,
        )
    return (
        table_to_max_ids_per_sparse_core,
        table_to_max_unique_ids_per_sparse_core,
    )

  def enqueue(
      self,
      features: Any,
      weights: Optional[Any] = None,
      device: Optional[str] = None,
  ) -> Any:
    """Preprocessing the features on host."""
    nest.assert_same_structure(self._feature_config, features)

    flat_inputs = nest.flatten(features)
    flat_weights = [None] * len(flat_inputs)
    if weights is not None:
      nest.assert_same_structure(self._feature_config, weights)
      flat_weights = nest.flatten(weights)

    in_tpu_context = self._raise_error_for_incorrect_control_flow_context()

    if in_tpu_context:
      # Automatically apply outside compilation if we are in tpu context.
      return tpu_replication.outside_compilation(
          self._preprocess_features,
          num_replicas_in_sync=self._strategy.num_replicas_in_sync,
          max_ids_per_chip_per_sample=self.max_ids_per_chip_per_sample,
          max_minibatches_per_sc=self.max_minibatches_per_sc,
          num_sc_per_chip=self._num_sc_per_chip,
          num_sc_shards=self._num_sc_shards,
          stacked_table_to_tables=self._stacked_table_to_tables,
          table_to_stacked_table_offset=self._table_to_stacked_table_offset,
          table_to_sample_count=self._table_to_sample_count,
          feature_to_sample_offset=self._feature_to_sample_offset,
          flat_features=self._flat_features,
          flat_inputs=flat_inputs,
          flat_weights=flat_weights,
      )
    elif device is None:
      # This is used by keras function tracing. Use any of the TPU devices
      # and trace once for a single device.
      tpu_devices = self._strategy.extended._tpu_devices  # pylint:disable=protected-access

      with ops.device(device_util.get_host_for_device(tpu_devices[0][0])):
        return self._preprocess_features(
            num_replicas_in_sync=self._strategy.num_replicas_in_sync,
            max_ids_per_chip_per_sample=self.max_ids_per_chip_per_sample,
            max_minibatches_per_sc=self.max_minibatches_per_sc,
            num_sc_per_chip=self._num_sc_per_chip,
            num_sc_shards=self._num_sc_shards,
            stacked_table_to_tables=self._stacked_table_to_tables,
            table_to_stacked_table_offset=self._table_to_stacked_table_offset,
            table_to_sample_count=self._table_to_sample_count,
            feature_to_sample_offset=self._feature_to_sample_offset,
            flat_features=self._flat_features,
            flat_inputs=flat_inputs,
            flat_weights=flat_weights,
        )
    else:
      device_spec = tf_device.DeviceSpec.from_string(device)
      if device_spec.device_type != "TPU":
        raise ValueError("Non-TPU device {} passed to enqueue.".format(device))

      with ops.device(device_util.get_host_for_device(device)):
        return self._preprocess_features(
            num_replicas_in_sync=self._strategy.num_replicas_in_sync,
            max_ids_per_chip_per_sample=self.max_ids_per_chip_per_sample,
            max_minibatches_per_sc=self.max_minibatches_per_sc,
            num_sc_per_chip=self._num_sc_per_chip,
            num_sc_shards=self._num_sc_shards,
            stacked_table_to_tables=self._stacked_table_to_tables,
            table_to_stacked_table_offset=self._table_to_stacked_table_offset,
            table_to_sample_count=self._table_to_sample_count,
            feature_to_sample_offset=self._feature_to_sample_offset,
            flat_features=self._flat_features,
            flat_inputs=flat_inputs,
            flat_weights=flat_weights,
        )

  def _copy_tensors_to_device(
      self,
      partitioned_tensors: Dict[str, Any],
  ) -> Any:
    """Copy tensors to device."""
    partitioned_device_tensors = {}
    for table_name in partitioned_tensors:
      partitioned_tensor = partitioned_tensors[table_name][0]
      row_pointers_unpadded_size = partitioned_tensors[table_name][1]
      ids_unpadded_size = partitioned_tensors[table_name][2]

      row_pointers, sorted_sample_ids, sorted_token_ids, sorted_gains = (
          xla_ops.tpu_copy_with_dynamic_shape(
              [
                  partitioned_tensor.row_pointers,
                  partitioned_tensor.sorted_sample_ids,
                  partitioned_tensor.sorted_token_ids,
                  partitioned_tensor.sorted_gains,
              ],
              [
                  row_pointers_unpadded_size,
                  ids_unpadded_size,
                  ids_unpadded_size,
                  ids_unpadded_size,
              ],
          )
      )

      # Placeholder Op for pipelining.
      row_pointers, sorted_sample_ids, sorted_token_ids, sorted_gains = (
          xla_ops.tpu_annotate_tensors_with_dynamic_shape([
              row_pointers,
              sorted_sample_ids,
              sorted_token_ids,
              sorted_gains,
          ])
      )

      partitioned_device_tensors[table_name] = PartitionedCsrFormatTensor(
          row_pointers=row_pointers,
          sorted_sample_ids=sorted_sample_ids,
          sorted_token_ids=sorted_token_ids,
          sorted_gains=sorted_gains,
          sample_count=partitioned_tensor.sample_count,
          num_minibatches_per_physical_sparse_core=(
              partitioned_tensor.num_minibatches_per_physical_sparse_core
          ),
      )
    return partitioned_device_tensors

  def dequeue(
      self,
      partitioned_tensors: Tuple[
          Dict[str, PartitionedCsrFormatTensor], int, int
      ],
  ) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]:
    """Perform embedding lookup."""
    # We expect this dequeue function will always run inside tpu context.
    context = EmbeddingPipeliningContext(
        _PIPELINE_MODE_FORWARD, self._pipelining
    )
    context.Enter()
    partitioned_tensors = tpu_replication.outside_compilation(
        self._copy_tensors_to_device,
        partitioned_tensors=partitioned_tensors,
    )

    activations = {}
    # Take num_minibatches_per_physical_sparse_core from any table as
    # they are the same across tables.
    num_minibatches_per_physical_sparse_core = list(
        partitioned_tensors.values()
    )[0].num_minibatches_per_physical_sparse_core

    for table_name in self._stacked_table_to_tables:
      partitioned_tensor = partitioned_tensors[table_name]

      table = self.variables[table_name]["parameters"]
      quantization_config = self._quantization_configs[table_name]
      if not isinstance(partitioned_tensor, PartitionedCsrFormatTensor):
        raise ValueError(
            "Expect PartitionedCsrFormatTensor but get"
            f" {type(partitioned_tensor)}."
        )
      activation = xla_ops.xla_sparse_dense_matmul_with_static_buffer_size(
          row_pointers=partitioned_tensor.row_pointers,
          sorted_sample_ids=partitioned_tensor.sorted_sample_ids,
          sorted_token_ids=partitioned_tensor.sorted_token_ids,
          sorted_gains=partitioned_tensor.sorted_gains,
          input_size=self._table_to_sample_count[table_name],
          embedding_table=table,
          num_minibatches_per_physical_sparse_core=num_minibatches_per_physical_sparse_core,
          quantization_config_low=(
              quantization_config.lower if quantization_config else 0
          ),
          quantization_config_high=(
              quantization_config.upper if quantization_config else 0
          ),
          quantization_config_num_buckets=(
              quantization_config.num_buckets if quantization_config else 0
          ),
          max_ids_per_sparse_core=self._table_to_max_ids_per_sparse_core[
              table_name
          ],
          max_unique_ids_per_sparse_core=self._table_to_max_unique_ids_per_sparse_core[
              table_name
          ],
          table_name=table_name,
      )

      activations[table_name] = activation

    context.Exit()
    # Note that unstacking gradients is placed on the core of the trianing step
    # to reduce the number of input/output arguments of the training loop during
    # pipelining.
    activations = self._unstack_activations(activations)

    return (activations, partitioned_tensors)

  def embedding_lookup(
      self, features: Any, weights: Optional[Any] = None
  ) -> Tuple[Any, Dict[str, PartitionedCsrFormatTensor]]:
    """Perform embedding lookup on the input feature.

    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.

    Raises:
      ValueError: If the input feature is not one of the Tensor, SparseTensor or
        RaggedTensor type.
      TypeError: If the type of any sequence in `features` does not match
        corresponding sequence in `feature_config`. Similarly for `weights`, if
        not `None`.

    Returns:
      packed_activations: Embedding lookup results packed as the same sequence
        of the input feature.
      packed_output: A dict of PartitionedCsrFormatTensors.
    """
    if not self._built:
      self._maybe_build()

    context = EmbeddingPipeliningContext(
        _PIPELINE_MODE_FORWARD, self._pipelining
    )
    context.Enter()

    partitioned_tensors = self.enqueue(features, weights)

    context.Exit()

    result = self.dequeue(partitioned_tensors)

    return result

  def _preprocess_features(
      self,
      num_replicas_in_sync: int,
      max_ids_per_chip_per_sample: int,
      max_minibatches_per_sc: int,
      num_sc_per_chip: int,
      num_sc_shards: int,
      stacked_table_to_tables: Dict[str, Any],
      table_to_stacked_table_offset: Dict[str, Tuple[str, int, int]],
      table_to_sample_count: Dict[str, int],
      feature_to_sample_offset: Dict[str, int],
      flat_features: Any,
      flat_inputs: Any,
      flat_weights: Optional[Any] = None,
  ) -> Any:
    """Function to preprocess features."""
    # Preprocess the inputs into list of COO tensor.
    table_to_list_of_coos = (
        TPUEmbeddingV2._preprocess_inputs_and_weights_to_list_of_coo_tensors(
            flat_inputs,
            flat_weights,
            flat_features,
            stacked_table_to_tables,
            table_to_stacked_table_offset,
            feature_to_sample_offset,
            num_sc_per_chip,
            table_to_sample_count,
            num_sc_shards,
        )
    )

    # Sort the COO tensors.
    table_to_sorted_coo_tensor = self._sort_list_of_coo_tensors(
        num_replicas_in_sync,
        table_to_list_of_coos,
        stacked_table_to_tables,
        num_sc_per_chip,
    )

    table_to_csr_format_tensor = (
        self._get_csr_wrapped_coo_from_sorted_coo_tensor(
            num_replicas_in_sync,
            max_ids_per_chip_per_sample,
            max_minibatches_per_sc,
            table_to_sorted_coo_tensor,
            stacked_table_to_tables,
            table_to_sample_count,
            num_sc_per_chip,
        )
    )

    return table_to_csr_format_tensor

  @classmethod
  def _convert_input_feature_to_list_of_coo_tensors(
      cls,
      input_feature: Union[
          tensor.Tensor, sparse_tensor.SparseTensor, ragged_tensor.RaggedTensor
      ],
      weight: Optional[tensor.Tensor],
      feature_config: tpu_embedding_v2_utils.FeatureConfig,
      row_offset: int,
      col_offset: int,
      col_shift: int,
      unused_vocab_size: int,
      num_sc_per_chip: int,
      num_sc_shards: int,
      stacked_table_sample_count: int,
  ) -> Any:
    """Convert any of the expected input types to a COO format."""
    sample_count = functools.reduce(operator.mul, feature_config.output_shape)
    if isinstance(input_feature, tensor.Tensor):
      input_feature = array_ops.reshape(input_feature, [-1])
      if weight is None:
        weight = array_ops.ones_like(input_feature, dtype=dtypes.float32)
      elif isinstance(weight, tensor.Tensor):
        weight = array_ops.reshape(weight, [-1])
      else:
        raise ValueError(
            f"Expect weight to be Tensor type but got {type(weight)}"
        )
      row_ids_list, col_ids_list, gains_list = (
          xla_ops.convert_to_list_of_sparse_core_coo_tensors(
              indices_or_row_splits=array_ops.zeros((0,), dtype=dtypes.int32),
              values=math_ops.cast(input_feature, dtype=dtypes.int32),
              weights=math_ops.cast(weight, dtypes.float32),
              sample_count=sample_count,
              combiner=feature_config.table.combiner,
              num_sc_per_chip=num_sc_per_chip,
              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,
          )
      )
    elif isinstance(input_feature, sparse_tensor.SparseTensor):
      if weight is None:
        weight = array_ops.ones_like(input_feature.values, dtype=dtypes.float32)
      elif isinstance(weight, sparse_tensor.SparseTensor):
        weight = weight.values
      else:
        raise ValueError(
            f"Expect weight to be SparseTensor type but got {type(weight)}"
        )
      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(
                  input_feature.indices, dtype=dtypes.int32
              ),
              values=math_ops.cast(input_feature.values, dtype=dtypes.int32),
              weights=math_ops.cast(weight, dtypes.float32),
              sample_count=sample_count,
              combiner=feature_config.table.combiner,
              num_sc_per_chip=num_sc_per_chip,
              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,
          )
      )
    elif isinstance(input_feature, ragged_tensor.RaggedTensor):
      if weight is None:
        weight = array_ops.ones_like(input_feature.values, dtype=dtypes.float32)
      elif isinstance(weight, ragged_tensor.RaggedTensor):
        weight = weight.values
      else:
        raise ValueError(
            f"Expect weight to be RaggedTensor type but got {type(weight)}"
        )
      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(
                  input_feature.row_splits, dtype=dtypes.int32
              ),
              values=math_ops.cast(input_feature.values, dtype=dtypes.int32),
              weights=math_ops.cast(weight, dtypes.float32),
              sample_count=sample_count,
              combiner=feature_config.table.combiner,
              num_sc_per_chip=num_sc_per_chip,
              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,
          )
      )
    else:
      raise ValueError(
          f"Input of unknown type {type(input_feature)}. Please only pass "
          "Tensor, SparseTensor or RaggedTensor as input to embedding "
          "lookup."
      )
    return row_ids_list, col_ids_list, gains_list, sample_count

  @classmethod
  def _preprocess_inputs_and_weights_to_list_of_coo_tensors(
      cls,
      flat_inputs: Any,
      flat_weights: Any,
      flat_features: Any,
      stacked_table_to_tables: Dict[str, Any],
      table_to_stacked_table_offset: Dict[str, Tuple[str, int, int]],
      feature_to_sample_offset: Dict[str, int],
      num_sc_per_chip: int,
      stacked_table_to_sample_count: Dict[str, int],
      num_sc_shards: int,
  ) -> Dict[str, Any]:
    """Convert the raw inputs into list of coo tensors."""
    table_to_list_of_coos = {  # pylint: disable=g-complex-comprehension
        table_name: (
            [[], [], [], []],
            [[], [], [], []],
            [[], [], [], []],
            [],
            [],
        )
        for table_name in stacked_table_to_tables
    }
    for inp, weight, (feature_path, feature) in zip(
        flat_inputs, flat_weights, flat_features
    ):
      table_name, col_offset, col_shift = table_to_stacked_table_offset[
          feature.table.name
      ]
      stacked_table_sample_count = stacked_table_to_sample_count[table_name]
      row_offset = feature_to_sample_offset[feature_path]
      # Consider making this into one op per table rather than per feature?
      row_ids_list, col_ids_list, gains_list, sample_count = (
          TPUEmbeddingV2._convert_input_feature_to_list_of_coo_tensors(
              inp,
              weight,
              feature,
              row_offset,
              col_offset,
              col_shift,
              feature.table.vocabulary_size,
              num_sc_per_chip,
              num_sc_shards,
              stacked_table_sample_count,
          )
      )
      for i in range(num_sc_per_chip):
        table_to_list_of_coos[table_name][0][i].append(row_ids_list[i])
        table_to_list_of_coos[table_name][1][i].append(col_ids_list[i])
        table_to_list_of_coos[table_name][2][i].append(gains_list[i])
      table_to_list_of_coos[table_name][3].append(
          sample_count // num_sc_per_chip
      )
      table_to_list_of_coos[table_name][4].append(col_offset)
    return table_to_list_of_coos

  def _sort_list_of_coo_tensors(
      self,
      num_replicas_in_sync: int,
      table_to_list_of_coos: Dict[str, Any],
      stacked_table_to_tables: Dict[str, Any],
      num_sc_per_chip: int,
  ) -> Tuple[Dict[str, Any], List[tensor.Tensor]]:
    """Sort the coo tensors by replica."""
    table_to_sorted_coo_tensor = {
        table_name: ([], [], [], []) for table_name in stacked_table_to_tables
    }
    for table_name in stacked_table_to_tables:
      # Feature width are the same across stacked tables.
      feature_width = stacked_table_to_tables[table_name][0].dim

      total_vocab_size = sum([
          table.vocabulary_size for table in stacked_table_to_tables[table_name]
      ])
      for i in range(num_sc_per_chip):
        row_ids_list = table_to_list_of_coos[table_name][0][i]
        col_ids_list = table_to_list_of_coos[table_name][1][i]
        gains_list = table_to_list_of_coos[table_name][2][i]
        sample_count_list = table_to_list_of_coos[table_name][3]
        col_offset_list = table_to_list_of_coos[table_name][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,
            col_ids_list=col_ids_list,
            gains_list=gains_list,
            sample_count_list=sample_count_list,
            col_offset_list=col_offset_list,
            num_replica=num_replicas_in_sync,
            table_vocab_size=total_vocab_size,
            feature_width=feature_width,
            num_sc_per_chip=num_sc_per_chip,
            max_ids_per_sparse_core=self._table_to_max_ids_per_sparse_core[
                table_name
            ],
            max_unique_ids_per_sparse_core=self._table_to_max_unique_ids_per_sparse_core[
                table_name
            ],
            table_name=table_name,
        )

        table_to_sorted_coo_tensor[table_name][0].append(sorted_row_ids)
        table_to_sorted_coo_tensor[table_name][1].append(sorted_col_ids)
        table_to_sorted_coo_tensor[table_name][2].append(sorted_gains)
        table_to_sorted_coo_tensor[table_name][3].append(id_counts)

    return table_to_sorted_coo_tensor

  def _get_csr_wrapped_coo_from_sorted_coo_tensor(
      self,
      num_replicas_in_sync: int,
      max_ids_per_chip_per_sample: int,
      max_minibatches_per_sc: int,
      table_to_sorted_coo_tensor: Dict[str, Any],
      stacked_table_to_tables: Dict[str, Any],
      table_to_sample_count: Dict[str, int],
      num_sc_per_chip: int,
  ) -> Any:
    """Get csr wrapped coo tensor from the sorted coo tensor."""
    table_to_csr_format_tensor = {}
    for table_name in stacked_table_to_tables:
      (
          sorted_row_ids_list,
          sorted_col_ids_list,
          sorted_gains_list,
          id_counts_list,
      ) = table_to_sorted_coo_tensor[table_name]

      # Feature width are the same across stacked tables.
      feature_width = stacked_table_to_tables[table_name][0].dim

      total_vocab_size = sum([
          table.vocabulary_size for table in stacked_table_to_tables[table_name]
      ])
      (
          row_pointers,
          sorted_sample_ids,
          sorted_token_ids,
          sorted_gains,
          row_pointers_unpadded_size,
          ids_unpadded_size,
          num_minibatches_per_physical_sparse_core,
      ) = 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
          ),  # no splits are needed.
          sample_count_per_sc=table_to_sample_count[table_name]
          // num_sc_per_chip,
          num_replica=num_replicas_in_sync,
          max_minibatches_per_sc=max_minibatches_per_sc,
          max_ids_per_chip_per_sample=max_ids_per_chip_per_sample,
          table_vocab_size=total_vocab_size,
          feature_width=feature_width,
          table_name=table_name,
          allow_id_dropping=self._sparse_core_embedding_config.allow_id_dropping,
      )
      table_to_csr_format_tensor[table_name] = (
          PartitionedCsrFormatTensor(
              row_pointers=row_pointers,
              sorted_sample_ids=sorted_sample_ids,
              sorted_token_ids=sorted_token_ids,
              sorted_gains=sorted_gains,
              sample_count=table_to_sample_count[table_name],
              num_minibatches_per_physical_sparse_core=num_minibatches_per_physical_sparse_core,
          ),
          row_pointers_unpadded_size,
          ids_unpadded_size,
      )
    return table_to_csr_format_tensor

  def _unstack_activations(self, activations: Dict[str, tensor.Tensor]):
    """Untack the incoming per table activations into per feature."""

    # Activations are stacked in a particular order. That order is the order
    # features appear in the self._flat_features.
    flattened_activations = []
    table_to_current_offset = {
        table_name: 0 for table_name in self._stacked_table_to_tables
    }
    for table_name in self._stacked_table_to_tables:
      activation_shape = activations[table_name].shape
      activations[table_name] = array_ops.reshape(
          activations[table_name],
          [self._num_sc_per_chip, -1, activation_shape[-1]],
      )
    for _, feature in self._flat_features:
      sample_count = functools.reduce(operator.mul, feature.output_shape)
      table_name = self._table_to_stacked_table_offset[feature.table.name][0]
      extra_cols = self._table_to_padding_columns[feature.table.name]
      activation = array_ops.slice(
          activations[table_name],
          [0, table_to_current_offset[table_name], 0],
          [
              self._num_sc_per_chip,
              sample_count // self._num_sc_per_chip,
              feature.table.dim - extra_cols,
          ],
      )
      # Reshape to follow the user's requested output shape.
      activation = array_ops.reshape(
          activation,
          list(feature.output_shape) + [feature.table.dim - extra_cols],
      )
      flattened_activations.append(activation)
      table_to_current_offset[table_name] += (
          sample_count // self._num_sc_per_chip
      )

    return nest.pack_sequence_as(self._feature_config, flattened_activations)

  def _stack_gradients(self, gradients):
    """Stack the incoming gradients to per table gradients."""

    # Gradients are stacked in a particular order. That order is the order
    # features appear in the self._flat_features.
    table_to_gradient_list = {
        table_name: [[], [], [], []]
        for table_name in self._stacked_table_to_tables
    }
    flattend_gradients = nest.flatten(gradients)
    for gradient, (path, feature) in zip(
        flattend_gradients, self._flat_features
    ):
      sample_count = functools.reduce(operator.mul, feature.output_shape)
      if gradient is not None and not isinstance(gradient, tensor.Tensor):
        raise ValueError(
            f"found non-tensor type: {type(gradient)} at path {path}."
        )
      if gradient is None:
        # TODO(bfontain): In the case that an entire table's gradient is gone
        # then maybe we can just omit the update all together?
        logging.warning(
            (
                "No gradient passed for feature %s, sending zero "
                "gradient. This may not be correct behavior for certain "
                "optimizers like Adam."
            ),
            path,
        )
        gradient = array_ops.zeros(
            (sample_count, feature.table.dim), dtype=dtypes.float32
        )
      table_name = self._table_to_stacked_table_offset[feature.table.name][0]
      extra_cols = self._table_to_padding_columns[feature.table.name]
      gradient = array_ops.reshape(
          gradient, [-1, feature.table.dim - extra_cols]
      )
      if extra_cols != 0:
        gradient = array_ops.pad(gradient, [[0, 0], [0, extra_cols]])
        # Ensure static shape after padding.
        gradient.set_shape([sample_count, feature.table.dim])
      per_sc_sample_count = sample_count // self._num_sc_per_chip
      for i in range(self._num_sc_per_chip):
        table_to_gradient_list[table_name][i].append(
            array_ops.slice(
                gradient,
                [i * per_sc_sample_count, 0],
                [
                    per_sc_sample_count,
                    feature.table.dim,
                ],
            )
        )
    for table_name in table_to_gradient_list:
      table_to_gradient_list[table_name] = array_ops.concat(
          [
              array_ops.concat(table_to_gradient_list[table_name][i], axis=0)
              for i in range(self._num_sc_per_chip)
          ],
          axis=0,
      )
    return table_to_gradient_list


# TODO(pineapplejuice233): Merge this function with the one in tpu_embeding_v2.py once
# this file is OSSed.
def extract_variable_info(
    kwargs: Any,
) -> Tuple[
    str, Tuple[int, ...], dtypes.DType, Callable[[], Any], Optional[int]
]:
  """Extracts the variable creation attributes from the kwargs.

  Args:
    kwargs: a dict of keyword arguments that were passed to a variable creator
      scope.

  Returns:
    A tuple of variable name, shape, dtype, initialization function,
    restore_uid.
  """

  def get_restore_uid(initial_value: Callable[..., Any]) -> int | None:
    return getattr(initial_value, "restore_uid", None)

  if isinstance(kwargs["initial_value"], functools.partial) and (
      "shape" in kwargs["initial_value"].keywords
      or kwargs["initial_value"].args
  ):
    # Sometimes shape is passed positionally, sometimes it's passed as a kwarg.
    if "shape" in kwargs["initial_value"].keywords:
      shape = kwargs["initial_value"].keywords["shape"]
    else:
      shape = kwargs["initial_value"].args[0]
    return (
        kwargs["name"],
        shape,
        kwargs["initial_value"].keywords.get("dtype", kwargs["dtype"]),
        kwargs["initial_value"].func,
        get_restore_uid(kwargs["initial_value"].func),
    )
  elif (
      "shape" not in kwargs
      or kwargs["shape"] is None
      or not callable(kwargs["initial_value"])
  ):
    raise ValueError(
        "Unable to extract initializer function and shape from {}. Please "
        "either pass a function that expects a shape and dtype as the "
        "initial value for your variable or functools.partial object with "
        "the shape and dtype kwargs set. This is needed so that we can "
        "initialize the shards of the ShardedVariable locally.".format(
            kwargs["initial_value"]
        )
    )
  else:
    return (
        kwargs["name"],
        kwargs["shape"],
        kwargs["dtype"],
        kwargs["initial_value"],
        get_restore_uid(kwargs["initial_value"]),
    )


def is_checkpoint_initial_value(initial_value: Any) -> bool:
  """Whether the initial value is from checkpoint."""
  return (
      isinstance(initial_value, base.CheckpointInitialValue)
      or isinstance(initial_value, base.CheckpointInitialValueCallable)
      or (
          isinstance(initial_value, functools.partial)
          and isinstance(
              initial_value.func, base.CheckpointInitialValueCallable
          )
      )
  )


def make_sharded_variable_creator(
    strategy: distribute_lib.Strategy,
) -> Callable[..., Any]:
  """Create a variable creator which shards across all the tpu device.

  Args:
    strategy: a TPUStrategy object.

  Returns:
    The sharded variable creator.
  """
  tpu_devices = strategy.extended._tpu_devices  # pylint:disable=protected-access

  def _create_sharded_variable(next_creator, *args, **kwargs):
    """Create a TPUEmbeddingShardedVariable."""
    # Avoid the default mirror variable creator.
    kwargs["skip_mirrored_creator"] = True

    # Only support sharding on the first dimension.
    shard_dim = 0

    num_replicas, num_cores_per_replica = tpu_devices.shape

    is_ckpt_init_value = is_checkpoint_initial_value(kwargs["initial_value"])

    arg_spec = tf_inspect.getfullargspec(kwargs["initial_value"])
    if (
        is_ckpt_init_value
        and "shard_info" not in arg_spec.args
        and "shard_info" not in arg_spec.kwonlyargs
    ):
      raise ValueError(
          "When a sharded variable is initialized from a checkpoint, "
          "shard_info must be in arguments of the init function."
      )

    name, shape, dtype, unwrapped_initial_value, restore_uid = (
        extract_variable_info(kwargs)
    )

    shape = ops.tensor_shape.TensorShape(shape)
    num_devices = num_replicas * num_cores_per_replica
    # NOTE: only support sharding variables evenly across devices.
    if shape[shard_dim] % num_devices != 0:
      raise ValueError(
          "Only evenly sharding across devices is currently supported. "
          "Got shape {} and {} devices".format(shape, num_devices)
      )

    partition_shape = shape.as_list()
    partition_shape[shard_dim] = partition_shape[shard_dim] // num_devices

    unwrapped_arg_spec = tf_inspect.getargspec(unwrapped_initial_value)
    sharding_aware = "shard_info" in unwrapped_arg_spec.args

    variables = []
    # Keep track of offset for sharding aware initializers.
    partition_offset = [0] * len(shape)
    for replica_id in range(num_replicas):
      for logic_core_id in range(num_cores_per_replica):
        with ops.device(tpu_devices[replica_id][logic_core_id]):
          kwargs["name"] = f"{name}/{replica_id}"
          kwargs["shape"] = partition_shape
          if sharding_aware:
            # TODO(pineapplejuice233): Change this to use MOD sharding logic.
            shard_info = base.ShardInfo(
                tensor_shape.as_shape(partition_shape),
                copy.deepcopy(partition_offset),
            )
            kwargs["initial_value"] = functools.partial(
                kwargs["initial_value"], shard_info=shard_info
            )
            partition_offset[shard_dim] += partition_shape[shard_dim]
          else:
            kwargs["initial_value"] = functools.partial(
                unwrapped_initial_value, shape=partition_shape, dtype=dtype
            )
          variables.append(next_creator(*args, **kwargs))

    result = TPUEmbeddingShardedVariable(
        strategy, variables, tf_variables.VariableAggregation.NONE, None
    )
    if restore_uid is not None:
      result._maybe_initialize_trackable()  # pylint: disable=protected-access
      result._update_uid = restore_uid  # pylint: disable=protected-access
    return result

  return _create_sharded_variable