code: # Copyright 2017 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ======================================
  """Library of TPU helper functions."""
  
  import enum
  import math
  from typing import List, Optional, Text, Tuple
  
  import numpy as np
  
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.tpu.topology import Topology
  from tensorflow.python.util.tf_export import tf_export
  
  
  SINGLE_CORE_ASSIGNMENT = [[[0, 0, 0, 0]]]
  
  
  def _compute_task_and_cores_to_replicas(core_assignment, topology):
    """Computes a nested dict which maps task and logical core to replicas."""
    task_and_cores_to_replicas = {}
    for replica in range(core_assignment.shape[0]):
      for logical_core in range(core_assignment.shape[1]):
        coordinates = core_assignment[replica, logical_core, :]
        task_id = topology.task_ordinal_at_coordinates(coordinates)
        if task_id not in task_and_cores_to_replicas:
          task_and_cores_to_replicas[task_id] = {}
        if logical_core not in task_and_cores_to_replicas[task_id]:
          task_and_cores_to_replicas[task_id][logical_core] = set()
  
        task_and_cores_to_replicas[task_id][logical_core].add(replica)
  
    task_to_sorted_replica_id = {}
  
    for task, core_to_replicas in task_and_cores_to_replicas.items():
      core_to_sorted_replicas = {}
      for core, replicas in core_to_replicas.items():
        core_to_sorted_replicas[core] = sorted(replicas)
  
      task_to_sorted_replica_id[task] = core_to_sorted_replicas
    return task_to_sorted_replica_id
  
  
  @tf_export("tpu.experimental.DeviceAssignment")
  class DeviceAssignment(object):
    """Mapping from logical cores in a computation to the physical TPU topology.
  
    Prefer to use the `DeviceAssignment.build()` helper to construct a
    `DeviceAssignment`; it is easier if less flexible than constructing a
    `DeviceAssignment` directly.
    """
  
    def __init__(self, topology: Topology, core_assignment: np.ndarray):
      """Constructs a `DeviceAssignment` object.
  
      Args:
        topology: A `Topology` object that describes the physical TPU topology.
        core_assignment: A logical to physical core mapping, represented as a
          rank 3 numpy array. See the description of the `core_assignment`
          property for more details.
  
      Raises:
        ValueError: If `topology` is not `Topology` object.
        ValueError: If `core_assignment` is not a rank 3 numpy array.
      """
      if not isinstance(topology, Topology):
        raise ValueError("topology must be a Topology object, got {}".format(
            type(topology)))
      core_assignment = np.asarray(core_assignment, dtype=np.int32)
  
      self._topology = topology
  
      if core_assignment.ndim != 3:
        raise ValueError("core_assignment must be a rank 3 numpy array, "
                         f"got shape {core_assignment.shape}")
  
      self._num_replicas = core_assignment.shape[0]
      self._num_cores_per_replica = core_assignment.shape[1]
  
      if core_assignment.shape[-1] != topology.mesh_rank:
        raise ValueError(
            "core_assignment.shape[-1] must have size equal to topology "
            f"rank ({topology.mesh_rank}), got "
            f"core_assignment.shape={core_assignment.shape}")
  
      self._core_assignment = core_assignment
      self._task_and_cores_to_replicas = _compute_task_and_cores_to_replicas(
          self._core_assignment, topology)
  
    @property
    def topology(self) -> Topology:
      """A `Topology` that describes the TPU topology."""
      return self._topology
  
    @property
    def num_cores_per_replica(self) -> int:
      """The number of cores per replica."""
      return self._num_cores_per_replica
  
    @property
    def num_replicas(self) -> int:
      """The number of replicas of the computation."""
      return self._num_replicas
  
    @property
    def core_assignment(self) -> np.ndarray:
      """The logical to physical core mapping.
  
      Returns:
        An integer numpy array of rank 3, with shape
        `[num_replicas, num_cores_per_replica, topology_rank]`. Maps
        (replica, logical core) pairs to physical topology coordinates.
      """
      return self._core_assignment
  
    def coordinates(self, replica: int, logical_core: int) -> Tuple:  # pylint:disable=g-bare-generic
      """Returns the physical topology coordinates of a logical core."""
      return tuple(self.core_assignment[replica, logical_core, :])
  
    def lookup_replicas(self, task_id: int, logical_core: int) -> List[int]:
      """Lookup replica ids by task number and logical core.
  
      Args:
        task_id: TensorFlow task number.
        logical_core: An integer, identifying a logical core.
      Returns:
        A sorted list of the replicas that are attached to that task and
        logical_core.
      Raises:
        ValueError: If no replica exists in the task which contains the logical
        core.
      """
      try:
        return self._task_and_cores_to_replicas[task_id][logical_core]
      except KeyError:
        raise ValueError(
            "Can not find any replica in task: {} contains logical_core: {} ".
            format(task_id, logical_core))
  
    def tpu_ordinal(self, replica: int = 0, logical_core: int = 0) -> int:
      """Returns the ordinal of the TPU device assigned to a logical core."""
      coordinates = self.coordinates(replica, logical_core)
      return self._topology.tpu_device_ordinal_at_coordinates(coordinates)
  
    def host_device(self,
                    replica: int = 0,
                    logical_core: int = 0,
                    job: Optional[Text] = None) -> Text:
      """Returns the CPU device attached to a logical core."""
      coordinates = self.coordinates(replica, logical_core)
      return self._topology.cpu_device_name_at_coordinates(coordinates, job=job)
  
    def tpu_device(self,
                   replica: int = 0,
                   logical_core: int = 0,
                   job: Optional[Text] = None) -> Text:
      """Returns the name of the TPU device assigned to a logical core."""
      coordinates = self.coordinates(replica, logical_core)
      return self._topology.tpu_device_name_at_coordinates(coordinates, job=job)
  
    @staticmethod
    def build(topology: Topology,
              computation_shape: Optional[np.ndarray] = None,
              computation_stride: Optional[np.ndarray] = None,
              num_replicas: int = 1) -> "DeviceAssignment":
      return device_assignment(topology, computation_shape, computation_stride,
                               num_replicas)
  
  
  def _open_ring_2d(x_size: int, y_size: int,
                    z_coord: int) -> List[Tuple[int, int, int]]:
    """Ring-order of a X by Y mesh, with a fixed Z coordinate.
  
    For example, in a 4x4 mesh, this returns the following order.
      0 -- 1 -- 2 -- 3
      |    |    |    |
      15-- 6 -- 5 -- 4
      |    |    |    |
      14-- 7 -- 8 -- 9
      |    |    |    |
      13-- 12-- 11-- 10
  
    Note that chip 0 is not included in the output.
  
    Args:
      x_size: An integer represents the mesh size in the x-dimension. Must be
        larger than 1.
      y_size: An integer represents the mesh size in the y-dimension. Must be
        larger than 1.
      z_coord: An integer represents the z-coordinate to use for the chips in the
        ring.
  
    Returns:
      A list of (x,y,z) triples in ring order.
    """
    ret = []
    for i in range(y_size // 2):
      for j in range(1, x_size):
        ret.append((j, 2 * i, z_coord))
      for j in range(x_size - 1, 0, -1):
        ret.append((j, 2 * i + 1, z_coord))
    for i in range(y_size - 1, 0, -1):
      ret.append((0, i, z_coord))
    return ret
  
  
  def _ring_3d(x_size: int, y_size: int,
               z_size: int) -> List[Tuple[int, int, int]]:
    """Ring-order of a X by Y by Z mesh.
  
    Constructs the 3d ring from 2d rings that are stacked in the Z dimension and
    joined in one corner.
  
    z == 0:
      0 -- 1 -- 2 -- 3
      |    |    |    |
      15 - 6 -- 5 -- 4
      |    |    |    |
      14 - 7 -- 8 -- 9
      |    |    |    |
      13 - 12 - 11 - 10
    z == 1:
      63 - 30 - 29 - 28
      |    |    |    |
      16 - 25 - 26 - 27
      |    |    |    |
      17 - 24 - 23 - 22
      |    |    |    |
      18 - 19 - 20 - 21
    z == 2:
      62 - 31 - 32 - 33
      |    |    |    |
      45 - 36 - 35 - 34
      |    |    |    |
      44 - 37 - 38 - 39
      |    |    |    |
      43 - 42 - 41 - 40
    z == 3:
      61 - 60 - 59 - 58
      |    |    |    |
      46 - 55 - 56 - 57
      |    |    |    |
      47 - 54 - 53 - 52
      |    |    |    |
      48 - 49 - 50 - 51
  
    Args:
      x_size: An integer represents the mesh size in the x-dimension. Must be
        larger than 1.
      y_size: An integer represents the mesh size in the y-dimension. Must be
        larger than 1.
      z_size: An integer represents the mesh size in the z-dimension. Must be
        larger than 1.  For example, in a 4x4x4 mesh, this returns the following
        order.
  
    Returns:
      A list of (x,y,z) triples in ring order.
    """
  
    # Handle the case where 2 dimensions are size 1.
    if x_size == 1 and y_size == 1:
      return [(0, 0, i) for i in range(z_size)]
    if x_size == 1 and z_size == 1:
      return [(0, i, 0) for i in range(y_size)]
    if y_size == 1 and z_size == 1:
      return [(i, 0, 0) for i in range(x_size)]
  
    # Handle odd mesh dimensions.  This never happens in practice, so we don't
    # bother to try building something optimal.
    if (x_size > 1 and x_size % 2 != 0) or (y_size > 1 and
                                            y_size % 2 != 0) or (z_size > 1 and
                                                                 z_size % 2 != 0):
      logging.warning("Odd dimension")
      ret = []
      for z in range(z_size):
        for y in range(y_size):
          ret.extend((x, y, z) for x in range(x_size))
      return ret
  
    # Always start with chip 0.
    ret = [(0, 0, 0)]
    # Handle the case where one dimension is size 1.  We just build a flat, 2d
    # ring.
    if z_size == 1:
      ret.extend(_open_ring_2d(x_size, y_size, 0))
      return ret
    if y_size == 1:
      ret = [(0, 0, 0)]
      ret.extend((x, y, z) for (x, z, y) in _open_ring_2d(x_size, z_size, 0))
      return ret
    if x_size == 1:
      ret = [(0, 0, 0)]
      ret.extend((x, y, z) for (y, z, x) in _open_ring_2d(y_size, z_size, 0))
      return ret
  
    # Handle the case where all dimensions have size > 1 and even.
    ret = [(0, 0, 0)]
    for i in range(0, z_size):
      r = _open_ring_2d(x_size, y_size, i)
      if i % 2 == 0:
        ret.extend(r)
      else:
        ret.extend(reversed(r))
    for i in range(z_size - 1, 0, -1):
      ret.append((0, 0, i))
    return ret
  
  
  class DeviceOrderMode(enum.IntEnum):
    """The way of determining device orders when computing device assignment."""
    # By default the mode is set to AUTO, the library will choose to form rings
    # when that is possible.
    AUTO = 0
    # Form rings for replicas and model-parallel cores.
    RING = 1
    # Form meshes for replicas and/or model-parallel cores.
    MESH = 2
  
  
  def device_assignment(
      topology: Topology,
      computation_shape: Optional[np.ndarray] = None,
      computation_stride: Optional[np.ndarray] = None,
      num_replicas: int = 1,
      device_order_mode: DeviceOrderMode = DeviceOrderMode.AUTO
  ) -> DeviceAssignment:
    """Computes a device_assignment of a computation across a TPU topology.
  
    Attempts to choose a compact grid of cores for locality.
  
    Returns a `DeviceAssignment` that describes the cores in the topology assigned
    to each core of each replica.
  
    `computation_shape` and `computation_stride` values should be powers of 2 for
    optimal packing.
  
    Args:
      topology: A `Topology` object that describes the TPU cluster topology. To
        obtain a TPU topology, evaluate the `Tensor` returned by
        `initialize_system` using `Session.run`. Either a serialized
        `TopologyProto` or a `Topology` object may be passed. Note: you must
          evaluate the `Tensor` first; you cannot pass an unevaluated `Tensor`
          here.
      computation_shape: A rank 1 int32 numpy array with size equal to the
        topology rank, describing the shape of the computation's block of cores.
        If None, the `computation_shape` is `[1] * topology_rank`.
      computation_stride: A rank 1 int32 numpy array of size `topology_rank`,
        describing the inter-core spacing of the `computation_shape` cores in the
        TPU topology. If None, the `computation_stride` is `[1] * topology_rank`.
      num_replicas: The number of computation replicas to run. The replicas will
        be packed into the free spaces of the topology.
      device_order_mode: An enum of `DeviceOrderMode` class which indicates
        whether to assign devices to form rings or meshes, or let the library to
        choose.
  
    Returns:
      A DeviceAssignment object, which describes the mapping between the logical
      cores in each computation replica and the physical cores in the TPU
      topology.
  
    Raises:
      ValueError: If `topology` is not a valid `Topology` object.
      ValueError: If `computation_shape` or `computation_stride` are not 1D int32
        numpy arrays with shape [3] where all values are positive.
      ValueError: If computation's replicas cannot fit into the TPU topology.
    """
    # Deserialize the Topology proto, if it is a string.
    if isinstance(topology, bytes):
      topology = Topology(serialized=topology)
  
    if not isinstance(topology, Topology):
      raise ValueError(
          f"`topology` is not a Topology object; got {type(topology)}")
  
    topology_rank = len(topology.mesh_shape)
    mesh_shape = topology.mesh_shape
    if computation_shape is None:
      computation_shape = np.array([1] * topology_rank, dtype=np.int32)
    else:
      computation_shape = np.asarray(computation_shape, dtype=np.int32)
  
    if computation_stride is None:
      computation_stride = np.array([1] * topology_rank, dtype=np.int32)
    else:
      computation_stride = np.asarray(computation_stride, dtype=np.int32)
  
    if computation_shape.shape != (topology_rank,):
      raise ValueError(
          f"computation_shape must have shape [{topology_rank}]; "
          f"got {computation_shape.shape}"
      )
    if computation_stride.shape != (topology_rank,):
      raise ValueError(
          f"computation_stride must have shape [{topology_rank}]; "
          f"got {computation_stride.shape}"
      )
  
    if any(computation_shape < 1):
      raise ValueError(
          "computation_shape must be positive; got computation_shape={}".format(
              computation_shape))
    if any(computation_stride < 1):
      raise ValueError(
          "computation_stride must be positive; got computation_stride={}".format(
              computation_stride))
  
    # Computes the physical size of one computation instance.
    computation_footprint = computation_shape * computation_stride
    if any(computation_footprint > mesh_shape):
      raise ValueError(
          "computation footprint {} does not fit in TPU topology shape {}".format(
              computation_footprint, mesh_shape))
  
    # Computes how many copies of the computation footprint fit in the mesh.
    block_counts = mesh_shape // computation_footprint
  
    replica_counts = block_counts * computation_stride
    max_replicas = np.prod(replica_counts)
    if num_replicas > max_replicas:
      raise ValueError(
          "requested {} replicas but only {} replicas with shape {} and "
          "computation_stride {} fit in a TPU mesh of shape {}".format(
              num_replicas, max_replicas, computation_shape, computation_stride,
              mesh_shape))
  
    def ceil_of_ratio(n, m):
      return (n + m - 1) // m
  
    if topology.missing_devices.size == 0:
      replica_shape = [0] * topology_rank
      if num_replicas > 0:
        remaining_replicas = num_replicas
        remaining_dims = topology_rank
  
        # Choose dimensions as close to an equal cube as possible,
        # in order of increasing dimension size. By visiting dimensions
        # in increasing size, we assign the most constrained dimension
        # first, so we won't make infeasible choices.
        #
        # As a secondary sort order, visit the last dimension (core index) first,
        # then the other dimensions in increasing order. This means we try to use
        # both cores on the same chip in preference to two cores on different
        # chips.  We visit the x dimension first, and the z dimension last, so
        # that we prefer to arrange adjacent replicas on the same machine when
        # possible.
        #
        # For example, if num_replicas == 4, we prefer to use a replica_shape of
        # (2,1,1,2) over (1,1,2,2).
  
        for x, ni in sorted(((x, ((i + 1) % topology_rank))
                             for (i, x) in enumerate(replica_counts))):
          i = (ni + topology_rank - 1) % topology_rank
          target_size = int(math.ceil(remaining_replicas**(1.0 / remaining_dims)))
          replica_shape[i] = min(target_size, x)
          remaining_replicas = ceil_of_ratio(remaining_replicas, replica_shape[i])
          remaining_dims -= 1
  
        assert remaining_replicas == 1 and remaining_dims == 0
  
      # Assigns an offset to each replica such that no two replicas overlap.
      replica_offsets = np.full([num_replicas, topology_rank], -1, dtype=np.int32)
  
      enable_3d_tiling = (
          topology_rank == 4 and
          computation_shape[-1] == mesh_shape[-1]  # Only handle 3D case.
          and np.prod(computation_stride) == 1  # Ensure no stride.
          and num_replicas == max_replicas)  # Full replication.
  
      if device_order_mode != DeviceOrderMode.AUTO:
        if device_order_mode == DeviceOrderMode.RING and not enable_3d_tiling:
          raise ValueError(
              "device_order_mode=DeviceOrderMode.RING is not compatible with the "
              "3D tiling current topology.  Try setting "
              "device_order_mode=DeviceOrderMode.AUTO"
          )
        enable_3d_tiling = device_order_mode == DeviceOrderMode.RING
  
      if enable_3d_tiling:
        assignment = []
        inner_ring = _ring_3d(computation_shape[0], computation_shape[1],
                              computation_shape[2])
        outer_ring = _ring_3d(replica_shape[0], replica_shape[1],
                              replica_shape[2])
  
        for replica in range(num_replicas):
          outer_x, outer_y, outer_z = outer_ring[replica]
          per_replica_assignment = []
          for index in range(np.prod(computation_shape)):
            inner_x, inner_y, inner_z = inner_ring[index // mesh_shape[-1]]
            px = outer_x * computation_shape[0] + inner_x
            py = outer_y * computation_shape[1] + inner_y
            pz = outer_z * computation_shape[2] + inner_z
            pi = index % mesh_shape[-1]
            per_replica_assignment.append([px, py, pz, pi])
          assignment.append(per_replica_assignment)
      else:
        for replica in range(num_replicas):
          # Chooses a replica number in each axis.
          t = replica
          pos = []
          # Visit the core number first.
          for dim in np.concatenate([[replica_shape[-1]], replica_shape[:-1]]):
            pos.append(t % dim)
            t //= dim
          replica_pos = np.concatenate([pos[1:], [pos[0]]])
  
          # Determines where that replica starts in each axis.
          outer = replica_pos // computation_stride
          inner = replica_pos % computation_stride
          replica_offsets[replica, :] = outer * computation_footprint + inner
  
        # Computes a logical core -> physical core mapping for each replica.
        indices = [
            np.arange(0, computation_shape[i] * computation_stride[i],
                      computation_stride[i]) for i in range(topology_rank)
        ]
        indices = np.concatenate(
            [i[..., np.newaxis] for i in np.meshgrid(*indices, indexing="ij")],
            axis=-1)
        indices = indices.reshape((-1, topology_rank))
        assignment = indices + replica_offsets[:, np.newaxis, :]
    else:
      # We have a slice with missing chips. We define a simple assignment by
      # ignoring computation stride. This assignment should enable a consistent
      # and correct device assignment on degraded slices. It is optimal when
      # weights are not sharded. But this device assignment may be sub-optimal for
      # other model parallelism scenarios.
      assert np.prod(computation_stride) == 1
      # Next, we check if we have sufficient devices.
      assert num_replicas * np.prod(
          computation_shape) <= topology.num_tasks * topology.num_tpus_per_task
      # Map replicas to physical devices in task order.
      device_coordinates = topology.device_coordinates
      assignment = []
      devices_per_replica = np.prod(computation_shape)
      for rindex in range(num_replicas):
        replica_assignment = []
        for index in range(devices_per_replica):
          logical_id = rindex * devices_per_replica + index
          # Pick logical cores in task order
          task = logical_id // topology.num_tpus_per_task
          device = logical_id % topology.num_tpus_per_task
          # Append physical cores to the replica assignment
          replica_assignment.append(device_coordinates[task, device, :])
        assignment.append(replica_assignment)
  
    return DeviceAssignment(topology, core_assignment=assignment)
   # Copyright 2020 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Companion classes for mid level API for TPU Embeddings in TF2."""
  
  import abc
  import math
  import typing
  from typing import Any, Dict, Callable, List, Optional, Text, Tuple, TypeVar, Union
  
  from absl import logging
  import six
  
  from tensorflow.core.protobuf.tpu import optimization_parameters_pb2
  from tensorflow.core.protobuf.tpu import tpu_embedding_configuration_pb2
  from tensorflow.python.distribute import sharded_variable
  from tensorflow.python.framework import ops
  from tensorflow.python.framework.tensor_shape import TensorShape
  from tensorflow.python.ops import init_ops_v2
  from tensorflow.python.ops import variables as tf_variables
  from tensorflow.python.tpu.ops import tpu_ops
  from tensorflow.python.types import core
  from tensorflow.python.util.tf_export import tf_export
  
  
  TableVariable = TypeVar("TableVariable", sharded_variable.ShardedVariable,
                          tf_variables.Variable)
  SlotVarCreationFnType = Callable[
      [TableVariable, List[Text], List[init_ops_v2.Initializer]],
      Dict[Text, TableVariable]]
  ClipValueType = Union[Tuple[float, float], float]
  
  
  @six.add_metaclass(abc.ABCMeta)
  class _Optimizer(object):
    """Base class for all optimizers, with common parameters."""
  
    def __init__(
        self,
        learning_rate: Union[float, Callable[[], float]],
        use_gradient_accumulation: bool,
        clip_weight_min: Optional[float],
        clip_weight_max: Optional[float],
        weight_decay_factor: Optional[float],
        multiply_weight_decay_factor_by_learning_rate: bool,
        clipvalue: Optional[ClipValueType] = None,
        slot_variable_creation_fn: Optional[SlotVarCreationFnType] = None):
      self.learning_rate = learning_rate
      self.use_gradient_accumulation = use_gradient_accumulation
      self.clip_weight_min = clip_weight_min
      self.clip_weight_max = clip_weight_max
      if not use_gradient_accumulation and clipvalue is not None:
        raise ValueError(
            f"When `use_gradient_accumulation` is False, gradient clipping "
            f"cannot be used and `clipvalue` should be left as None. "
            f"Received value {clipvalue} for argument `clipvalue`.")
      if clipvalue is None:
        clipvalue = (None, None)
      elif not isinstance(clipvalue, tuple):
        clipvalue = (-1. * clipvalue, clipvalue)
      self.clip_gradient_min, self.clip_gradient_max = clipvalue
  
      self.weight_decay_factor = weight_decay_factor
      self.multiply_weight_decay_factor_by_learning_rate = (
          multiply_weight_decay_factor_by_learning_rate)
  
      if (slot_variable_creation_fn is not None and
          not callable(slot_variable_creation_fn)):
        raise ValueError(
            f"Argument `slot_variable_creation_fn` must be either None or a "
            f"callable. Received: {slot_variable_creation_fn}")
      self.slot_variable_creation_fn = slot_variable_creation_fn
  
    @abc.abstractmethod
    def _slot_names(self) -> List[Text]:
      """Returns the name of all the slot variables.
  
      This does not include the 'parameters' variable and these names must match
      the names of the slots variables as used in the corresponding
      `tpu_ops.load_tpu_embedding_*` ops.
      """
      raise NotImplementedError
  
    @abc.abstractmethod
    def _slot_initializers(self) -> List[init_ops_v2.Initializer]:
      """Returns initializers for slot variables.
  
      This returns a parallel list to self._slot_names().
      """
      raise NotImplementedError
  
    def _set_optimization_parameters(
        self, parameters: optimization_parameters_pb2.OptimizationParameters):
      """Sets the optimizer fields in the OptimizationParameters."""
      if self.use_gradient_accumulation:
        parameters.gradient_accumulation_status = (
            optimization_parameters_pb2.GradientAccumulationStatus.ENABLED)
      else:
        parameters.gradient_accumulation_status = (
            optimization_parameters_pb2.GradientAccumulationStatus.DISABLED)
  
      if self.clip_weight_min is not None:
        parameters.clipping_limits.lower.value = self.clip_weight_min
  
      if self.clip_weight_max is not None:
        parameters.clipping_limits.upper.value = self.clip_weight_max
  
      if self.clip_gradient_min is not None:
        parameters.gradient_clipping_limits.lower.value = self.clip_gradient_min
  
      if self.clip_gradient_max is not None:
        parameters.gradient_clipping_limits.upper.value = self.clip_gradient_max
  
      if self.weight_decay_factor:
        parameters.weight_decay_factor = self.weight_decay_factor
        if self.multiply_weight_decay_factor_by_learning_rate:
          parameters.multiply_weight_decay_factor_by_learning_rate = True
  
    @abc.abstractmethod
    def _load(self) -> Callable[..., ops.Operation]:
      """Returns the load function for the optimizer."""
      raise NotImplementedError
  
    @abc.abstractmethod
    def _retrieve(self) -> Callable[..., core.Tensor]:
      """Returns the retrieve function for the optimizer."""
      raise NotImplementedError
  
    def _create_slots(
        self, table: "TableConfig",
        variable_creator: Callable[[Text, init_ops_v2.Initializer],
                                   tf_variables.Variable]
    ) -> Dict[Text, tf_variables.Variable]:
      """Creates slot variables for table.
  
      Args:
        table: The table variable to create slots for.
        variable_creator: A function which creates variables. Takes parameters
          'name', 'initializer'.
  
      Returns:
        A dict of variables, keyed by self._slot_names().
      """
      if self.slot_variable_creation_fn is not None:
        return self.slot_variable_creation_fn(table, self._slot_names(),
                                              self._slot_initializers())
      else:
        slots = {}
        for slot, initializer in zip(self._slot_names(),
                                     self._slot_initializers()):
          slots[slot] = variable_creator(slot, initializer)
        return slots
  
  
  @tf_export("tpu.experimental.embedding.SGD")
  class SGD(_Optimizer):
    """Optimization parameters for stochastic gradient descent for TPU embeddings.
  
    Pass this to `tf.tpu.experimental.embedding.TPUEmbedding` via the `optimizer`
    argument to set the global optimizer and its parameters:
  
    ```
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        ...
        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
    ```
  
    This can also be used in a `tf.tpu.experimental.embedding.TableConfig` as the
    optimizer parameter to set a table specific optimizer. This will override the
    optimizer and parameters for global embedding optimizer defined above:
  
    ```
    table_one = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...,
        optimizer=tf.tpu.experimental.embedding.SGD(0.2))
    table_two = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
  
    feature_config = (
        tf.tpu.experimental.embedding.FeatureConfig(
            table=table_one),
        tf.tpu.experimental.embedding.FeatureConfig(
            table=table_two))
  
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        batch_size=...
        optimizer=tf.tpu.experimental.embedding.SGD(0.1))
    ```
  
    In the above example, the first feature will be looked up in a table that has
    a learning rate of 0.2 while the second feature will be looked up in a table
    that has a learning rate of 0.1.
  
    See 'tensorflow/core/protobuf/tpu/optimization_parameters.proto' for a
    complete description of these parameters and their impacts on the optimizer
    algorithm.
    """
  
    def __init__(self,
                 learning_rate: Union[float, Callable[[], float]] = 0.01,
                 clip_weight_min: Optional[float] = None,
                 clip_weight_max: Optional[float] = None,
                 weight_decay_factor: Optional[float] = None,
                 multiply_weight_decay_factor_by_learning_rate: bool = None,
                 clipvalue: Optional[ClipValueType] = None):
      """Optimization parameters for stochastic gradient descent.
  
      Args:
        learning_rate: The learning rate. It should be a floating point value or a
          callable taking no arguments for a dynamic learning rate.
        clip_weight_min: the minimum value to clip by; None means -infinity.
        clip_weight_max: the maximum value to clip by; None means +infinity.
        weight_decay_factor: amount of weight decay to apply; None means that the
          weights are not decayed. Weights are decayed by multiplying the weight
          by this factor each step.
        multiply_weight_decay_factor_by_learning_rate: if true,
          `weight_decay_factor` is multiplied by the current learning rate.
        clipvalue: Controls clipping of the gradient. Set to either a single
          positive scalar value to get clipping or a tiple of scalar values (min,
          max) to set a separate maximum or minimum. If one of the two entries is
          None, then there will be no clipping that direction. Note if this is
          set, you may see a decrease in performance as  gradient accumulation
          will be enabled (it is normally off for SGD as it has no affect on
          accuracy). See
          'tensorflow/core/protobuf/tpu/optimization_parameters.proto' for more
          information on gradient accumulation and its impact on tpu embeddings.
      """
      use_gradient_accumulation = clipvalue is not None
  
      super(SGD, self).__init__(
          learning_rate, use_gradient_accumulation, clip_weight_min,
          clip_weight_max, weight_decay_factor,
          multiply_weight_decay_factor_by_learning_rate, clipvalue)
  
    def _slot_names(self) -> List[Text]:
      return []
  
    def _slot_initializers(self) -> List[init_ops_v2.Initializer]:
      return []
  
    def _set_optimization_parameters(
        self, parameters: optimization_parameters_pb2.OptimizationParameters):
      super(SGD, self)._set_optimization_parameters(parameters)
      parameters.stochastic_gradient_descent.SetInParent()
  
    def _load(self) -> Callable[..., ops.Operation]:
      return tpu_ops.load_tpu_embedding_stochastic_gradient_descent_parameters
  
    def _retrieve(self) -> Callable[..., core.Tensor]:
      return tpu_ops.retrieve_tpu_embedding_stochastic_gradient_descent_parameters
  
  
  @tf_export("tpu.experimental.embedding.Adagrad")
  class Adagrad(_Optimizer):
    """Optimization parameters for Adagrad with TPU embeddings.
  
    Pass this to `tf.tpu.experimental.embedding.TPUEmbedding` via the `optimizer`
    argument to set the global optimizer and its parameters:
  
    ```python
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        ...
        optimizer=tf.tpu.experimental.embedding.Adagrad(0.1))
    ```
  
    This can also be used in a `tf.tpu.experimental.embedding.TableConfig` as the
    optimizer parameter to set a table specific optimizer. This will override the
    optimizer and parameters for global embedding optimizer defined above:
  
    ```python
    table_one = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...,
        optimizer=tf.tpu.experimental.embedding.Adagrad(0.2))
    table_two = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
  
    feature_config = (
        tf.tpu.experimental.embedding.FeatureConfig(
            table=table_one),
        tf.tpu.experimental.embedding.FeatureConfig(
            table=table_two))
  
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        batch_size=...
        optimizer=tf.tpu.experimental.embedding.Adagrad(0.1))
    ```
  
    In the above example, the first feature will be looked up in a table that has
    a learning rate of 0.2 while the second feature will be looked up in a table
    that has a learning rate of 0.1.
  
    See 'tensorflow/core/protobuf/tpu/optimization_parameters.proto' for a
    complete description of these parameters and their impacts on the optimizer
    algorithm.
    """
  
    def __init__(
        self,
        learning_rate: Union[float, Callable[[], float]] = 0.001,
        initial_accumulator_value: float = 0.1,
        use_gradient_accumulation: bool = True,
        clip_weight_min: Optional[float] = None,
        clip_weight_max: Optional[float] = None,
        weight_decay_factor: Optional[float] = None,
        multiply_weight_decay_factor_by_learning_rate: bool = None,
        slot_variable_creation_fn: Optional[SlotVarCreationFnType] = None,
        clipvalue: Optional[ClipValueType] = None):
      """Optimization parameters for Adagrad.
  
      Args:
        learning_rate: The learning rate. It should be a floating point value or a
          callable taking no arguments for a dynamic learning rate.
        initial_accumulator_value: initial accumulator for Adagrad.
        use_gradient_accumulation: setting this to `False` makes embedding
          gradients calculation less accurate but faster.
        clip_weight_min: the minimum value to clip by; None means -infinity.
        clip_weight_max: the maximum value to clip by; None means +infinity.
        weight_decay_factor: amount of weight decay to apply; None means that the
          weights are not decayed.
        multiply_weight_decay_factor_by_learning_rate: if true,
          `weight_decay_factor` is multiplied by the current learning rate.
        slot_variable_creation_fn: If you wish do directly control the creation of
          the slot variables, set this to a callable taking three parameters: a
            table variable, a list of slot names to create for it, and a list of
            initializers. This function should return a dict with the slot names
            as keys and the created variables as values with types matching the
            table variable. When set to None (the default), uses the built-in
            variable creation.
        clipvalue: Controls clipping of the gradient. Set to either a single
          positive scalar value to get clipping or a tuple of scalar values (min,
          max) to set a separate maximum or minimum. If one of the two entries is
          None, then there will be no clipping that direction.
      """
      super(Adagrad, self).__init__(
          learning_rate, use_gradient_accumulation, clip_weight_min,
          clip_weight_max, weight_decay_factor,
          multiply_weight_decay_factor_by_learning_rate, clipvalue,
          slot_variable_creation_fn)
      if initial_accumulator_value <= 0:
        raise ValueError(
            f"Argument `initial_accumulator_value` must be a positive float. "
            f"Received: {initial_accumulator_value}")
      self.initial_accumulator_value = initial_accumulator_value
  
    def _slot_names(self) -> List[Text]:
      return ["accumulators"]
  
    def _slot_initializers(self) -> List[init_ops_v2.Initializer]:
      return [init_ops_v2.Constant(self.initial_accumulator_value)]
  
    def _set_optimization_parameters(
        self, parameters: optimization_parameters_pb2.OptimizationParameters):
      super(Adagrad, self)._set_optimization_parameters(parameters)
      parameters.adagrad.SetInParent()
  
    def _load(self) -> Callable[..., ops.Operation]:
      return tpu_ops.load_tpu_embedding_adagrad_parameters
  
    def _retrieve(self) -> Callable[..., core.Tensor]:
      return tpu_ops.retrieve_tpu_embedding_adagrad_parameters
  
  
  @tf_export("tpu.experimental.embedding.FTRL")
  class FTRL(_Optimizer):
    """Optimization parameters for FTRL with TPU embeddings.
  
    See Algorithm 1 of this
    [paper](https://research.google.com/pubs/archive/41159.pdf).
  
    Pass this to `tf.tpu.experimental.embedding.TPUEmbedding` via the `optimizer`
    argument to set the global optimizer and its parameters:
  
    ```python
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        ...
        optimizer=tf.tpu.experimental.embedding.FTRL(0.1))
    ```
  
    This can also be used in a `tf.tpu.experimental.embedding.TableConfig` as the
    optimizer parameter to set a table specific optimizer. This will override the
    optimizer and parameters for global embedding optimizer defined above:
  
    ```python
    table_one = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...,
        optimizer=tf.tpu.experimental.embedding.FTRL(0.2))
    table_two = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
  
    feature_config = (
        tf.tpu.experimental.embedding.FeatureConfig(
            table=table_one),
        tf.tpu.experimental.embedding.FeatureConfig(
            table=table_two))
  
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        batch_size=...
        optimizer=tf.tpu.experimental.embedding.FTRL(0.1))
    ```
  
    In the above example, the first feature will be looked up in a table that has
    a learning rate of 0.2 while the second feature will be looked up in a table
    that has a learning rate of 0.1.
  
    See 'tensorflow/core/protobuf/tpu/optimization_parameters.proto' for a
    complete description of these parameters and their impacts on the optimizer
    algorithm.
    """
  
    def __init__(
        self,
        learning_rate: Union[float, Callable[[], float]] = 0.001,
        learning_rate_power: float = -0.5,
        l1_regularization_strength: float = 0.0,
        l2_regularization_strength: float = 0.0,
        beta: float = 0.0,
        initial_accumulator_value: float = 0.1,
        use_gradient_accumulation: bool = True,
        clip_weight_min: Optional[float] = None,
        clip_weight_max: Optional[float] = None,
        weight_decay_factor: Optional[float] = None,
        multiply_weight_decay_factor_by_learning_rate: bool = None,
        slot_variable_creation_fn: Optional[SlotVarCreationFnType] = None,
        clipvalue: Optional[ClipValueType] = None,
        multiply_linear_by_learning_rate: bool = False,
        allow_zero_accumulator: bool = False):
      """Optimization parameters for Adagrad.
  
      Args:
        learning_rate: The learning rate. It should be a floating point value or a
          callable taking no arguments for a dynamic learning rate.
        learning_rate_power: A float value, must be less or equal to zero.
          Controls how the learning rate decreases during training. Use zero for a
          fixed learning rate.
        l1_regularization_strength: A float value, must be greater than or equal
          to zero.
        l2_regularization_strength: A float value, must be greater than or equal
          to zero.
        beta: A float value, representing the beta value from the paper.
        initial_accumulator_value: The starting value for accumulators. Only zero
          or positive values are allowed.
        use_gradient_accumulation: setting this to `False` makes embedding
          gradients calculation less accurate but faster.
        clip_weight_min: the minimum value to clip by; None means -infinity.
        clip_weight_max: the maximum value to clip by; None means +infinity.
        weight_decay_factor: amount of weight decay to apply; None means that the
          weights are not decayed.
        multiply_weight_decay_factor_by_learning_rate: if true,
          `weight_decay_factor` is multiplied by the current learning rate.
        slot_variable_creation_fn: If you wish do directly control the creation of
          the slot variables, set this to a callable taking three parameters: a
            table variable, a list of slot names to create for it, and a list of
            initializers. This function should return a dict with the slot names
            as keys and the created variables as values with types matching the
            table variable. When set to None (the default), uses the built-in
            variable creation.
        clipvalue: Controls clipping of the gradient. Set to either a single
          positive scalar value to get clipping or a tuple of scalar values (min,
          max) to set a separate maximum or minimum. If one of the two entries is
          None, then there will be no clipping that direction.
        multiply_linear_by_learning_rate: If set to True, a modified formula is
          used for FTRL that treats the "linear" accumulator as being
          pre-multiplied by the learning rate (i.e., the accumulator named
          "linear" actually stores "linear * learning_rate"). Other than
          checkpoint compatibility, this is mathematically equivalent for a static
          learning rate; for a dynamic learning rate, it is nearly the same as
          long as the learning rate does not change quickly. The benefit of this
          is that the modified formula handles zero and near-zero learning rates
          without producing NaNs, improving flexibility for learning rate ramp-up.
        allow_zero_accumulator: If set to True, changes some internal formulas to
          allow zero and near-zero accumulator values at the cost of some
          performance; this only needs to be set if you are using an initial
          accumulator value of zero, which is uncommon.
      """
      super().__init__(learning_rate, use_gradient_accumulation, clip_weight_min,
                       clip_weight_max, weight_decay_factor,
                       multiply_weight_decay_factor_by_learning_rate, clipvalue,
                       slot_variable_creation_fn)
      if initial_accumulator_value <= 0:
        raise ValueError(
            f"Argument `initial_accumulator_value` must be a positive float. "
            f"Received: {initial_accumulator_value}")
      self.initial_accumulator_value = initial_accumulator_value
      self.learning_rate_power = learning_rate_power
      self.l1_regularization_strength = l1_regularization_strength
      self.l2_regularization_strength = l2_regularization_strength
      self.beta = beta
      self.multiply_linear_by_learning_rate = multiply_linear_by_learning_rate
      self.allow_zero_accumulator = allow_zero_accumulator
  
    def _slot_names(self) -> List[Text]:
      return ["accumulators", "linears"]
  
    def _slot_initializers(self) -> List[init_ops_v2.Initializer]:
      return [
          init_ops_v2.Constant(self.initial_accumulator_value),
          init_ops_v2.Constant()
      ]
  
    def _set_optimization_parameters(
        self, parameters: optimization_parameters_pb2.OptimizationParameters):
      super()._set_optimization_parameters(parameters)
      ftrl = parameters.ftrl
      ftrl.l1 = self.l1_regularization_strength
      ftrl.l2 = self.l2_regularization_strength
      ftrl.lr_power = self.learning_rate_power
      ftrl.beta = self.beta
      ftrl.multiply_linear_by_lr = self.multiply_linear_by_learning_rate
      ftrl.allow_zero_accumulator = self.allow_zero_accumulator
  
    def _load(self) -> Callable[..., ops.Operation]:
      return tpu_ops.load_tpu_embedding_ftrl_parameters
  
    def _retrieve(self) -> Callable[..., core.Tensor]:
      return tpu_ops.retrieve_tpu_embedding_ftrl_parameters
  
  
  @tf_export("tpu.experimental.embedding.Adam")
  class Adam(_Optimizer):
    """Optimization parameters for Adam with TPU embeddings.
  
    Pass this to `tf.tpu.experimental.embedding.TPUEmbedding` via the `optimizer`
    argument to set the global optimizer and its parameters:
  
    NOTE: By default this optimizer is lazy, i.e. it will not apply the gradient
    update of zero to rows that were not looked up. You can change this behavior
    by setting `lazy_adam` to `False`.
  
    ```python
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        ...
        optimizer=tf.tpu.experimental.embedding.Adam(0.1))
    ```
  
    This can also be used in a `tf.tpu.experimental.embedding.TableConfig` as the
    optimizer parameter to set a table specific optimizer. This will override the
    optimizer and parameters for global embedding optimizer defined above:
  
    ```python
    table_one = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...,
        optimizer=tf.tpu.experimental.embedding.Adam(0.2))
    table_two = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
  
    feature_config = (
        tf.tpu.experimental.embedding.FeatureConfig(
            table=table_one),
        tf.tpu.experimental.embedding.FeatureConfig(
            table=table_two))
  
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        batch_size=...
        optimizer=tf.tpu.experimental.embedding.Adam(0.1))
    ```
  
    In the above example, the first feature will be looked up in a table that has
    a learning rate of 0.2 while the second feature will be looked up in a table
    that has a learning rate of 0.1.
  
    See 'tensorflow/core/protobuf/tpu/optimization_parameters.proto' for a
    complete description of these parameters and their impacts on the optimizer
    algorithm.
    """
  
    def __init__(
        self,
        learning_rate: Union[float, Callable[[], float]] = 0.001,
        beta_1: float = 0.9,
        beta_2: float = 0.999,
        epsilon: float = 1e-07,
        lazy_adam: bool = True,
        sum_inside_sqrt: bool = True,
        use_gradient_accumulation: bool = True,
        clip_weight_min: Optional[float] = None,
        clip_weight_max: Optional[float] = None,
        weight_decay_factor: Optional[float] = None,
        multiply_weight_decay_factor_by_learning_rate: bool = None,
        slot_variable_creation_fn: Optional[SlotVarCreationFnType] = None,
        clipvalue: Optional[ClipValueType] = None):
      """Optimization parameters for Adam.
  
      See 'tensorflow/core/protobuf/tpu/optimization_parameters.proto' for a
      complete description of these parameters and their impacts on the optimizer
      algorithm.
  
      Args:
        learning_rate: The learning rate. It should be a floating point value or a
          callable taking no arguments for a dynamic learning rate.
        beta_1: A float value. The exponential decay rate for the 1st moment
          estimates.
        beta_2: A float value. The exponential decay rate for the 2nd moment
          estimates.
        epsilon: A small constant for numerical stability.
        lazy_adam: Use lazy Adam instead of Adam. Lazy Adam trains faster.
        sum_inside_sqrt: When this is true, the Adam update formula is changed
          from `m / (sqrt(v) + epsilon)` to `m / sqrt(v + epsilon**2)`. This
          option improves the performance of TPU training and is not expected to
          harm model quality.
        use_gradient_accumulation: Setting this to `False` makes embedding
          gradients calculation less accurate but faster.
        clip_weight_min: the minimum value to clip by; None means -infinity.
        clip_weight_max: the maximum value to clip by; None means +infinity.
        weight_decay_factor: amount of weight decay to apply; None means that the
          weights are not decayed.
        multiply_weight_decay_factor_by_learning_rate: if true,
          `weight_decay_factor` is multiplied by the current learning rate.
        slot_variable_creation_fn: If you wish do directly control the creation of
          the slot variables, set this to a callable taking three parameters: a
            table variable, a list of slot names to create for it, and a list of
            initializers. This function should return a dict with the slot names
            as keys and the created variables as values with types matching the
            table variable. When set to None (the default), uses the built-in
            variable creation.
        clipvalue: Controls clipping of the gradient. Set to either a single
          positive scalar value to get clipping or a tiple of scalar values (min,
          max) to set a separate maximum or minimum. If one of the two entries is
          None, then there will be no clipping that direction.
      """
      super(Adam, self).__init__(
          learning_rate, use_gradient_accumulation, clip_weight_min,
          clip_weight_max, weight_decay_factor,
          multiply_weight_decay_factor_by_learning_rate, clipvalue,
          slot_variable_creation_fn)
      if beta_1 < 0. or beta_1 >= 1.:
        raise ValueError(
            f"Argument `beta_1` must be >= 0 and < 1. Received: {beta_1}.")
      if beta_2 < 0. or beta_2 >= 1.:
        raise ValueError(
            f"Argument `beta_2` must be >= 0 and < 1. Received: {beta_1}.")
      if epsilon <= 0.:
        raise ValueError("epsilon must be positive; got {}.".format(epsilon))
      if not use_gradient_accumulation and not lazy_adam:
        raise ValueError(
            "When disabling lazy Adam (`lazy_adam=False`), "
            "gradient accumulation must be used. "
            "Set `use_gradient_accumulation` to False.")
  
      self.beta_1 = beta_1
      self.beta_2 = beta_2
      self.epsilon = epsilon
      self.lazy_adam = lazy_adam
      self.sum_inside_sqrt = sum_inside_sqrt
  
    def _slot_names(self) -> List[Text]:
      return ["momenta", "velocities"]
  
    def _slot_initializers(self) -> List[init_ops_v2.Initializer]:
      return [init_ops_v2.Constant(), init_ops_v2.Constant()]
  
    def _set_optimization_parameters(
        self, parameters: optimization_parameters_pb2.OptimizationParameters):
      super(Adam, self)._set_optimization_parameters(parameters)
      parameters.adam.beta1 = self.beta_1
      parameters.adam.beta2 = self.beta_2
      parameters.adam.epsilon = self.epsilon
      parameters.adam.use_non_lazy_adam = not self.lazy_adam
      parameters.adam.use_sum_inside_sqrt = self.sum_inside_sqrt
  
    def _load(self) -> Callable[..., ops.Operation]:
      return tpu_ops.load_tpu_embedding_adam_parameters
  
    def _retrieve(self) -> Callable[..., core.Tensor]:
      return tpu_ops.retrieve_tpu_embedding_adam_parameters
  
  
  @tf_export("tpu.experimental.embedding.TableConfig")
  class TableConfig(object):
    """Configuration data for one embedding table.
  
    This class holds the configuration data for a single embedding table. It is
    used as the `table` parameter of a
    `tf.tpu.experimental.embedding.FeatureConfig`. Multiple
    `tf.tpu.experimental.embedding.FeatureConfig` objects can use the same
    `tf.tpu.experimental.embedding.TableConfig` object. In this case a shared
    table will be created for those feature lookups.
  
    ```python
    table_config_one = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
    table_config_two = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
    feature_config = {
        'feature_one': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_one),
        'feature_two': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_one),
        'feature_three': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_two)}
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        batch_size=...
        optimizer=tf.tpu.experimental.embedding.Adam(0.1))
    ```
  
    The above configuration has 2 tables, and three features. The first two
    features will be looked up in the first table and the third feature will be
    looked up in the second table.
  
    """
  
    def __init__(self,
                 vocabulary_size: int,
                 dim: int,
                 initializer: Optional[Callable[[Any], None]] = None,
                 optimizer: Optional[_Optimizer] = None,
                 combiner: Text = "mean",
                 name: Optional[Text] = None):
      """Embedding table configuration.
  
      Args:
        vocabulary_size: Size of the table's vocabulary (number of rows).
        dim: The embedding dimension (width) of the table.
        initializer: A callable initializer taking one parameter, the shape of the
          variable that will be initialized. Will be called once per task, to
          initialize that task's shard of the embedding table. If not specified,
          defaults to `truncated_normal_initializer` with mean `0.0` and standard
          deviation `1/sqrt(dim)`.
        optimizer: An optional instance of an optimizer parameters class, instance
          of one of `tf.tpu.experimental.embedding.SGD`,
          `tf.tpu.experimental.embedding.Adagrad` or
          `tf.tpu.experimental.embedding.Adam`. It set will override the global
          optimizer passed to `tf.tpu.experimental.embedding.TPUEmbedding`.
        combiner: A string specifying how to reduce if there are multiple entries
          in a single row. Currently 'mean', 'sqrtn', 'sum' are supported, with
          'mean' the default. 'sqrtn' often achieves good accuracy, in particular
          with bag-of-words columns. For more information, see
          `tf.nn.embedding_lookup_sparse`.
        name: An optional string used to name the table. Useful for debugging.
  
      Returns:
        `TableConfig`.
  
      Raises:
        ValueError: if `vocabulary_size` is not a positive integer.
        ValueError: if `dim` is not a positive integer.
        ValueError: if `initializer` is specified and is not callable.
        ValueError: if `combiner` is not supported.
      """
      if not isinstance(vocabulary_size, int) or vocabulary_size < 1:
        raise ValueError(
            f"Argument `vocabulary_size` must be an int and must be >= 1. "
            f"Received: {vocabulary_size}")
  
      if not isinstance(dim, int) or dim < 1:
        raise ValueError(
            f"Argument `dim` (embedding dimension) "
            f"must be an int and must be >= 1. Received: {dim}")
  
      if (initializer is not None) and (not callable(initializer)):
        raise ValueError(
            f"Argument `initializer` must be a callable (or None). "
            f"Received: {initializer}")
      if initializer is None:
        initializer = init_ops_v2.TruncatedNormal(mean=0.0,
                                                  stddev=1/math.sqrt(dim))
      accepted_combiners = ("mean", "sum", "sqrtn")
      if combiner not in accepted_combiners:
        raise ValueError(
            f"Argument `combiner` must be one of {accepted_combiners}. "
            f"Received: {combiner}")
  
      self.vocabulary_size = vocabulary_size
      self.dim = dim
      self.initializer = initializer
      self.optimizer = optimizer
      self.combiner = combiner
      self.name = name
  
    def __repr__(self):
      # If using the default initializer, just print "None" for clarity.
      initializer = self.initializer
  
      if isinstance(initializer, init_ops_v2.TruncatedNormal):
        # PY2 type checking can't infer type of initializer even after if.
        initializer = typing.cast(init_ops_v2.TruncatedNormal, initializer)
        if (initializer.mean == 0.0
            and math.isclose(initializer.stddev, 1/math.sqrt(self.dim))):  # pytype: disable=module-attr (math.isclose not in PY2)
          initializer = None
  
      return (
          "TableConfig(vocabulary_size={vocabulary_size!r}, dim={dim!r}, "
          "initializer={initializer!r}, optimizer={optimizer!r}, "
          "combiner={combiner!r}, name={name!r})".format(
              vocabulary_size=self.vocabulary_size,
              dim=self.dim,
              initializer=initializer,
              optimizer=self.optimizer,
              combiner=self.combiner,
              name=self.name,)
      )
  
  
  @tf_export("tpu.experimental.embedding.FeatureConfig")
  class FeatureConfig(object):
    """Configuration data for one embedding feature.
  
    This class holds the configuration data for a single embedding feature. The
    main use is to assign features to `tf.tpu.experimental.embedding.TableConfig`s
    via the table parameter:
  
    ```python
    table_config_one = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
    table_config_two = tf.tpu.experimental.embedding.TableConfig(
        vocabulary_size=...,
        dim=...)
    feature_config = {
        'feature_one': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_one),
        'feature_two': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_one),
        'feature_three': tf.tpu.experimental.embedding.FeatureConfig(
            table=table_config_two)}
    embedding = tf.tpu.experimental.embedding.TPUEmbedding(
        feature_config=feature_config,
        batch_size=...
        optimizer=tf.tpu.experimental.embedding.Adam(0.1))
    ```
  
    The above configuration has 2 tables, and three features. The first two
    features will be looked up in the first table and the third feature will be
    looked up in the second table.
  
    You can also specify the output shape for each feature. The output shape
    should be the expected activation shape excluding the table dimension. For
    dense and sparse tensor, the output shape should be the same as the input
    shape excluding the last dimension. For ragged tensor, the output shape can
    mismatch the input shape.
  
    NOTE: The `max_sequence_length` will be only used when the input tensor has
    rank 2 and the `output_shape` is not set in the feature config.
  
    When feeding features into `embedding.enqueue` they can be `tf.Tensor`s,
    `tf.SparseTensor`s or `tf.RaggedTensor`s. When the argument
    `max_sequence_length` is 0, the default, you should expect a output of
    `embedding.dequeue` for this feature of shape `(batch_size, dim)`. If
    `max_sequence_length` is greater than 0, the feature is embedded as a sequence
    and padded up to the given length. The shape of the output for this feature
    will be `(batch_size, max_sequence_length, dim)`.
    """
  
    def __init__(self,
                 table: TableConfig,
                 max_sequence_length: int = 0,
                 validate_weights_and_indices: bool = True,
                 output_shape: Optional[Union[List[int], TensorShape]] = None,
                 name: Optional[Text] = None):
      """Feature configuration.
  
      Args:
        table: An instance of `tf.tpu.experimental.embedding.TableConfig`,
          describing the table in which this feature should be looked up.
        max_sequence_length: If positive, the feature is a sequence feature with
          the corresponding maximum sequence length. If the sequence is longer
          than this, it will be truncated. If 0, the feature is not a sequence
          feature.
        validate_weights_and_indices: If true, uses safe_embedding_lookup during
          serving which ensures there are no empty rows and all weights and ids
          are positive at the expense of extra compute cost.
        output_shape: Optional argument to config the output shape of the feature
          activation. If provided, the feature feeding to the `embedding.enqueue`
          has to match the shape (for ragged tensor, the input shape and output
          shape can mismatch). If not provided, the shape can be either provided
          to the `embedding.build` or auto detected at the runtime.
        name: An optional name for the feature, useful for debugging.
  
      Returns:
        `FeatureConfig`.
  
      Raises:
        ValueError: if `table` is not an instance of
          `tf.tpu.experimental.embedding.TableConfig`.
        ValueError: if `max_sequence_length` not an integer or is negative.
      """
      if not isinstance(table, TableConfig):
        raise ValueError(f"Argument `table` has invalid type {type(table)}. "
                         "Expected `tf.tpu.experimental.embedding.TableConfig`.")
  
      if not isinstance(max_sequence_length, int) or max_sequence_length < 0:
        raise ValueError(
            f"Argument `max_sequence_length` must be an int and must be >= 0. "
            f"Received: {max_sequence_length}")
  
      self.table = table
      self.max_sequence_length = max_sequence_length
      self.name = name
      self.output_shape = TensorShape(output_shape)
  
      if not isinstance(
          validate_weights_and_indices, bool):
        raise ValueError(
            f"Argument `validate_weights_and_indices` must be a boolean. "
            f"Received: {validate_weights_and_indices}")
  
      self.validate_weights_and_indices = validate_weights_and_indices
  
    def __repr__(self):
      return ("FeatureConfig(table={table!r}, "
              "max_sequence_length={max_sequence_length!r}, "
              "validate_weights_and_indices={"
              "validate_weights_and_indices!r}, name={name!r})".format(
                  table=self.table,
                  max_sequence_length=self.max_sequence_length,
                  validate_weights_and_indices=self.validate_weights_and_indices,
                  name=self.name))
  
  
  def log_tpu_embedding_configuration(
      config: tpu_embedding_configuration_pb2.TPUEmbeddingConfiguration) -> None:
    """Logs a TPUEmbeddingConfiguration proto across multiple statements.
  
    Args:
      config: TPUEmbeddingConfiguration proto to log.  Necessary because
        logging.info has a maximum length to each log statement, which
        particularly large configs can exceed.
    """
    logging.info("Beginning log of TPUEmbeddingConfiguration.")
    for line in str(config).splitlines():
      logging.info(line)
    logging.info("Done with log of TPUEmbeddingConfiguration.")
   # coding=utf-8
  # Copyright 2020 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Utilities for collectives."""
  
  import copy
  import enum
  
  from tensorflow.python.util import deprecation
  from tensorflow.python.util.tf_export import tf_export
  
  
  # TODO(b/170340570): print deprecation warning for CollectiveCommunication.
  @tf_export("distribute.experimental.CommunicationImplementation",
             "distribute.experimental.CollectiveCommunication")
  class CommunicationImplementation(enum.Enum):
    """Cross device communication implementation.
  
    Warning: The alias `tf.distribute.experimental.CollectiveCommunication` is
    deprecated and will be removed in a future version. Use
    `tf.distribute.experimental.CommunicationImplementation` instead.
  
    * `AUTO`: Automatically chosen by Tensorflow.
    * `RING`: TensorFlow's ring algorithms for all-reduce and
      all-gather.
    * `NCCL`: NVIDIA®'s NCCL library. This is now only used for all-reduce on
      GPUs; all-reduce on CPU, all-gather and broadcast fallbacks to RING.
    """
    AUTO = "AUTO"
    RING = "RING"
    NCCL = "NCCL"
    # TODO(ayushd): add ncclAllGather implementation.
  
  
  CollectiveCommunication = CommunicationImplementation
  
  
  @tf_export("distribute.experimental.CommunicationOptions")
  class _OptionsExported(object):
    """Options for cross device communications like All-reduce.
  
    This can be passed to methods like
    `tf.distribute.get_replica_context().all_reduce()` to optimize collective
    operation performance. Note that these are only hints, which may or may not
    change the actual behavior. Some options only apply to certain strategy and
    are ignored by others.
  
    One common optimization is to break gradients all-reduce into multiple packs
    so that weight updates can overlap with gradient all-reduce.
  
    Examples:
  
    ```python
    options = tf.distribute.experimental.CommunicationOptions(
        bytes_per_pack=50 * 1024 * 1024,
        timeout_seconds=120.0,
        implementation=tf.distribute.experimental.CommunicationImplementation.NCCL
    )
    grads = tf.distribute.get_replica_context().all_reduce(
        'sum', grads, options=options)
    optimizer.apply_gradients(zip(grads, vars),
        experimental_aggregate_gradients=False)
    ```
  
    """
  
    def __new__(cls, *args, **kwargs):
      # We expose a dummy class so that we can separate internal and public APIs.
      # Note that __init__ won't be called on the returned object if it's a
      # different class [1].
      # [1] https://docs.python.org/3/reference/datamodel.html#object.__new__
      return Options(*args, **kwargs)
  
    def __init__(self,
                 bytes_per_pack=0,
                 timeout_seconds=None,
                 implementation=CommunicationImplementation.AUTO):
      """Creates a CollectiveHints.
  
      Args:
        bytes_per_pack: a non-negative integer. Breaks collective operations into
          packs of certain size. If it's zero, the value is determined
          automatically. This hint is respected by all multi-replica strategies
          except `TPUStrategy`.
        timeout_seconds: a float or None, timeout in seconds. If not None, the
          collective raises `tf.errors.DeadlineExceededError` if it takes longer
          than this timeout. Zero disables timeout. This can be useful when
          debugging hanging issues.  This should only be used for debugging since
          it creates a new thread for each collective, i.e. an overhead of
          `timeout_seconds * num_collectives_per_second` more threads. This only
          works for `tf.distribute.experimental.MultiWorkerMirroredStrategy`.
        implementation: a
          `tf.distribute.experimental.CommunicationImplementation`. This is a hint
          on the preferred communication implementation. Possible values include
          `AUTO`, `RING`, and `NCCL`. NCCL is generally more performant for GPU,
          but doesn't work for CPU. This only works for
          `tf.distribute.experimental.MultiWorkerMirroredStrategy`.
  
      Raises:
        ValueError: When arguments have invalid value.
      """
      pass
  
  
  class Options(object):
    """Implementation of OptionsInterface."""
  
    def __init__(self,
                 bytes_per_pack=0,
                 timeout_seconds=None,
                 implementation=CommunicationImplementation.AUTO):
      if bytes_per_pack < 0:
        raise ValueError(
            f"Argument `bytes_per_pack` must be >=0, Received {bytes_per_pack}.")
      if isinstance(implementation, str):
        implementation = CommunicationImplementation(implementation.upper())
      if not isinstance(implementation, CommunicationImplementation):
        raise ValueError(
            "Argument `implementation` must be instance of "
            "`tf.distribute.experimental.CommunicationImplementation`.")
      self.bytes_per_pack = bytes_per_pack
      self.timeout_seconds = timeout_seconds
      self.implementation = implementation
  
    __init__.__doc__ = _OptionsExported.__init__.__doc__
  
    def merge(self, options):
      """Merges with another options and returns a new one.
  
      Values specified in the `options` takes precedence if they're not the
      default.
  
      Args:
        options: a `tf.distribute.experimental.CollectiveCommunication`.
  
      Returns:
        A new `tf.distribute.experimental.CollectiveCommunication`.
      """
      merged = copy.deepcopy(self)
      if options is None:
        return merged
      if options.bytes_per_pack != 0:
        merged.bytes_per_pack = options.bytes_per_pack
      if options.timeout_seconds is not None:
        merged.timeout_seconds = options.timeout_seconds
      if options.implementation != CommunicationImplementation.AUTO:
        merged.implementation = options.implementation
      return merged
  
  
  @tf_export("distribute.experimental.CollectiveHints")
  class Hints(object):
    """Hints for collective operations like AllReduce.
  
    This can be passed to methods like
    `tf.distribute.get_replica_context().all_reduce()` to optimize collective
    operation performance. Note that these are only hints, which may or may not
    change the actual behavior. Some options only apply to certain strategy and
    are ignored by others.
  
    One common optimization is to break gradients all-reduce into multiple packs
    so that weight updates can overlap with gradient all-reduce.
  
    Examples:
  
    - bytes_per_pack
  
    ```python
    hints = tf.distribute.experimental.CollectiveHints(
        bytes_per_pack=50 * 1024 * 1024)
    grads = tf.distribute.get_replica_context().all_reduce(
        'sum', grads, experimental_hints=hints)
    optimizer.apply_gradients(zip(grads, vars),
        experimental_aggregate_gradients=False)
    ```
  
    - timeout_seconds
  
    ```python
    strategy = tf.distribute.MirroredStrategy()
    hints = tf.distribute.experimental.CollectiveHints(
        timeout_seconds=120.0)
    try:
      strategy.reduce("sum", v, axis=None, experimental_hints=hints)
    except tf.errors.DeadlineExceededError:
      do_something()
    ```
  
    """
  
    @deprecation.deprecated(
        None, "use distribute.experimental.CommunicationOptions instead")
    def __new__(cls, bytes_per_pack=0, timeout_seconds=None):
      return Options(
          bytes_per_pack=bytes_per_pack, timeout_seconds=timeout_seconds)
  
    def __init__(self, bytes_per_pack=0, timeout_seconds=None):
      """Creates a CollectiveHints.
  
      Args:
        bytes_per_pack: a non-negative integer. Breaks collective operations into
          packs of certain size. If it's zero, the value is determined
          automatically. This only applies to all-reduce with
          `MultiWorkerMirroredStrategy` currently.
        timeout_seconds: a float or None, timeout in seconds. If not None, the
          collective raises `tf.errors.DeadlineExceededError` if it takes longer
          than this timeout. This can be useful when debugging hanging issues.
          This should only be used for debugging since it creates a new thread for
          each collective, i.e. an overhead of `timeout_seconds *
          num_collectives_per_second` more threads.  This only works for
          `tf.distribute.experimental.MultiWorkerMirroredStrategy`.
  
      Raises:
        ValueError: When arguments have invalid value.
      """
      pass
  
apis: numpy.meshgrid tensorflow.python.util.tf_export.tf_export tensorflow.python.tpu.topology.Topology numpy.asarray numpy.arange numpy.prod numpy.array numpy.concatenate numpy.full tensorflow.python.platform.tf_logging.warning
apis: tensorflow.python.framework.tensor_shape.TensorShape tensorflow.python.util.tf_export.tf_export tensorflow.python.ops.init_ops_v2.Constant
apis: tensorflow.python.util.deprecation.deprecated tensorflow.python.util.tf_export.tf_export