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

import functools
import sys
import time

import six

from tensorflow.python.autograph.operators import py_builtins
from tensorflow.python.data.experimental.ops import batching
from tensorflow.python.data.experimental.ops import cardinality as cardinality_lib
from tensorflow.python.data.experimental.ops import distribute
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.data.ops import multi_device_iterator_ops
from tensorflow.python.data.ops import optional_ops
from tensorflow.python.distribute import device_util
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import distribute_utils
from tensorflow.python.distribute import input_ops
from tensorflow.python.distribute import reduce_util
from tensorflow.python.distribute import values
from tensorflow.python.distribute.distribute_lib import InputReplicationMode
from tensorflow.python.eager import context
from tensorflow.python.eager import monitoring
from tensorflow.python.framework import composite_tensor
from tensorflow.python.framework import device as tf_device
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.framework import sparse_tensor
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_util
from tensorflow.python.framework import type_spec
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import cond as tf_cond
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.types import distribute as distribute_types
from tensorflow.python.util import nest
from tensorflow.python.util.compat import collections_abc


_distributed_dataset_initialization_time_milliseconds = monitoring.Sampler(
    "/tensorflow/api/distribution_strategy/"
    "distributed_dataset_initialization_time_milliseconds",
    monitoring.ExponentialBuckets(scale=1, growth_factor=2, bucket_count=26),
    "Track the time (in milliseconds) to initialize distributed datasets.",
    "strategy", "workers")

_distributed_dataset_from_function_initialization_time_milliseconds = (
    monitoring.Sampler(
        "/tensorflow/api/distribution_strategy/"
        "distributed_dataset_from_function_initialization_time_milliseconds",
        monitoring.ExponentialBuckets(
            scale=1, growth_factor=2, bucket_count=26),
        "Track the time (in milliseconds) to initialize distributed datasets "
        "from function.",
        "strategy", "workers"))


def get_iterator_spec_from_dataset(strategy, dataset):
  """Returns an iterator spec from dataset function.

  This function constructs type spec for iterator obtained from
  iter(dataset).

  Args:
    strategy: a `tf.distribute.Strategy` object, used to run all-reduce to
        handle last partial batch.
    dataset: A tf.data.Dataset instance. If using a function that returns a
      tf.data.Dataset instance, pass dataset_fn.structured_outputs.

  Returns:
    A type_spec for iterator for dataset instance.

  """
  # pylint: disable=protected-access
  output_element_spec = dataset.element_spec
  if isinstance(dataset._type_spec,
                (DistributedDatasetSpec,
                 DistributedDatasetsFromFunctionSpec)):
    iterator_type_spec = DistributedIteratorSpec(
        strategy.extended._input_workers_with_options(),
        output_element_spec,
        strategy.extended._container_strategy(),
        options=None,
        cardinality=dataset.cardinality,
        enable_get_next_as_optional=True)
  else:
    if strategy.extended._num_gpus_per_worker:
      logging.warning(
          f"{strategy.extended._num_gpus_per_worker} GPUs "
          "are allocated per worker. Please use DistributedDataset by "
          "calling strategy.experimental_distribute_dataset or strategy."
          "distribute_datasets_from_function to make best use of GPU "
          "resources"
      )
    iterator_type_spec = iterator_ops.IteratorSpec(output_element_spec)
  return iterator_type_spec
  # pylint: enable=protected-access


class InputWorkers(object):
  """A 1-to-many mapping from input worker devices to compute devices."""

  # TODO(ishark): Remove option canonicalize_devices and make all the callers
  # pass canonicalized or raw device strings as relevant from strategy.
  def __init__(self,
               worker_device_pairs,
               canonicalize_devices=True):
    """Initialize an `InputWorkers` object.

    Args:
      worker_device_pairs: A sequence of pairs: `(input device, a tuple of
        compute devices fed by that input device)`.
      canonicalize_devices: Whether to canonicalize devices for workers fully or
        partially. If False, it will partially canonicalize devices by removing
        job and task.
    """
    self._worker_device_pairs = worker_device_pairs
    self._input_worker_devices = tuple(d for d, _ in self._worker_device_pairs)
    self._canonicalize_devices = canonicalize_devices
    if canonicalize_devices:
      self._fed_devices = tuple(
          tuple(device_util.canonicalize(d)
                for d in f)
          for _, f in self._worker_device_pairs)
    else:
      self._fed_devices = tuple(
          tuple(device_util.canonicalize_without_job_and_task(d)
                for d in f)
          for _, f in self._worker_device_pairs)

  @property
  def num_workers(self):
    return len(self._input_worker_devices)

  @property
  def worker_devices(self):
    return self._input_worker_devices

  def compute_devices_for_worker(self, worker_index):
    return self._fed_devices[worker_index]

  def __repr__(self):
    devices = self.worker_devices
    debug_repr = ",\n".join("  %d %s: %s" %
                            (i, devices[i], self._fed_devices[i])
                            for i in range(len(devices)))
    return "%s:{\n%s}" % (self.__class__.__name__, debug_repr)

  def serialize(self):
    return (self._worker_device_pairs, self._canonicalize_devices)

  def deserialize(self, serialized):
    return InputWorkers(serialized)


def _calculate_replicas_with_values(strategy, input_workers, optional_list):
  """Computes the number of replicas that have values.

  Args:
    strategy: the `tf.distribute.Strategy`.
    input_workers: the `InputWorkers`.
    optional_list: a list of lists `tf.experimental.Optional`. The values from
      each compute device grouped by the input device.

  Returns:
    A scalar Tensor.
  """
  worker_has_values = []
  for worker, optionals in zip(input_workers.worker_devices, optional_list):
    with ops.device(worker):
      device_has_values = [
          math_ops.cast(v.has_value(), dtypes.int64) for v in optionals
      ]
      worker_has_values.append(
          math_ops.reduce_sum(device_has_values, keepdims=True))
  client_has_values = math_ops.reduce_sum(worker_has_values, keepdims=True)
  if strategy.extended._in_multi_worker_mode():  # pylint: disable=protected-access
    global_has_values = strategy.reduce(
        reduce_util.ReduceOp.SUM, client_has_values, axis=None)
    return array_ops.reshape(global_has_values, [])
  else:
    return array_ops.reshape(client_has_values, [])


def _is_statically_shaped(element_spec):
  """Test if an iterator output is statically shaped.

  For sparse and ragged tensors this only tests the batch dimension.

  Args:
    element_spec: a nest structure of `tf.TypeSpec`. The element spec of the
      dataset of the iterator.

  Returns:
    True if the shape is static, false otherwise.
  """

  for spec in nest.flatten(element_spec):
    if isinstance(
        spec, (sparse_tensor.SparseTensorSpec, ragged_tensor.RaggedTensorSpec)):
      # For sparse or ragged tensor, we should only check the first
      # dimension in order to get_next_as_optional. This is because
      # when these tensors get batched by dataset only the batch dimension
      # is set.
      if spec.shape.rank > 0 and spec.shape.as_list()[0] is None:
        return False
    else:
      for component in spec._flat_tensor_specs:  # pylint: disable=protected-access
        if not component.shape.is_fully_defined():
          return False
  return True


class DistributedIteratorBase(collections_abc.Iterator,
                              distribute_types.DistributedIteratorInterface):
  """Common implementation for all input iterators."""

  # pylint: disable=super-init-not-called
  def __init__(
      self,
      input_workers,
      iterators,
      strategy,
      cardinality,
      enable_get_next_as_optional,
      replica_order=None,
  ):
    assert isinstance(input_workers, InputWorkers)
    if not input_workers.worker_devices:
      raise ValueError("Should have at least one worker for input iterator.")

    self._iterators = iterators
    self._input_workers = input_workers
    self._strategy = strategy
    self._cardinality = cardinality
    self._enable_get_next_as_optional = enable_get_next_as_optional
    self._replica_order = replica_order

  def next(self):
    return self.__next__()

  def __next__(self):
    try:
      return self.get_next()
    except errors.OutOfRangeError:
      raise StopIteration

  def __iter__(self):
    return self

  def get_next_as_optional(self):
    # Ideally get_next_as_optional() should be consistent with get_next(), but
    # we used to always do partial batch handling in get_next_as_optional(). We
    # are keeping this behavior for now until we understantd the impact.

    # Skip partial batch handling when the dataset is infinite or empty, as
    # there won't be any partial batches in those cases. This gives the user
    # more static shapes as it avoids the tf.cond. Note that for empty datasets,
    # we can only skip in single client mode, as the dataset can be non-empty on
    # other workers.
    if self._cardinality == cardinality_lib.INFINITE:
      return optional_ops.Optional.from_value(
          self._get_next_no_partial_batch_handling())
    if (self._cardinality == 0 and
        not self._strategy.extended._in_multi_worker_mode()):  # pylint: disable=protected-access
      return optional_ops.Optional.empty(self._element_spec)

    optional_list = []
    for i, worker in enumerate(self._input_workers.worker_devices):
      with ops.device(worker):
        optional_list.append(self._iterators[i].get_next_as_optional_list())

    def _create_optional_with_dummy():
      value_list = _get_value_or_dummy(
          self._input_workers, optional_list, produce_dummy=True)

      if self._replica_order is not None:
        value_list = self._reorder_replicas(value_list)

      per_replica = _create_per_replica(value_list, self._strategy)
      return optional_ops.Optional.from_value(per_replica)

    def _create_empty_optional():
      return optional_ops.Optional.empty(self._element_spec)

    num_replicas_with_values = _calculate_replicas_with_values(
        self._strategy, self._input_workers, optional_list)

    return tf_cond.cond(
        num_replicas_with_values > 0,
        _create_optional_with_dummy,
        _create_empty_optional,
        strict=True)

  def get_next(self, name=None):
    """Returns the next input from the iterator for all replicas."""
    with distribute_lib.enter_or_assert_strategy(
        self._strategy):
      if distribute_lib.get_replica_context() is not None:
        raise ValueError("next(iterator) should be called from outside of "
                         "replica_fn. e.g. strategy.run(replica_fn, "
                         "args=(next(iterator),))")

    if not self._enable_get_next_as_optional:
      return self._get_next_no_partial_batch_handling(name)

    optional_list = []
    for i, worker in enumerate(self._input_workers.worker_devices):
      with ops.device(worker):
        optional_list.append(self._iterators[i].get_next_as_optional_list())
    num_replicas_with_values = _calculate_replicas_with_values(
        self._strategy, self._input_workers, optional_list)

    def _value_or_dummy():
      value_list = _get_value_or_dummy(
          self._input_workers, optional_list, produce_dummy=True)

      if self._replica_order is not None:
        value_list = self._reorder_replicas(value_list)

      return _create_per_replica(value_list, self._strategy)

    def _eof():
      # Optional.get_value raises InvalidArgumentError when there's no value,
      # so we need to call GetNext to raise EOFError.
      return self._get_next_no_partial_batch_handling()

    return tf_cond.cond(
        num_replicas_with_values > 0, _value_or_dummy, _eof, strict=True)

  def _get_next_no_partial_batch_handling(self, name=None):
    replicas = []
    for i, worker in enumerate(self._input_workers.worker_devices):
      if name is not None:
        d = tf_device.DeviceSpec.from_string(worker)
        new_name = "%s_%s_%d" % (name, d.job, d.task)
      else:
        new_name = None
      with ops.device(worker):
        # Make `replicas` a flat list of values across all replicas.
        replicas.extend(self._iterators[i].get_next_as_list(new_name))

    if self._replica_order is not None:
      replicas = self._reorder_replicas(replicas)

    return _create_per_replica(replicas, self._strategy)

  def _reorder_replicas(self, replicas):
    assert len(self._replica_order) == len(
        replicas
    ), "replica order size ({}) != replicas size ({})!".format(
        len(self._replica_order), len(replicas)
    )
    return [replicas[i] for i in self._replica_order]


class DistributedDatasetAndIteratorSpec(type_spec.TypeSpec):
  """Common Type specification for `DistributedDataset and DistributedDatasetsFromFunction."""

  __slots__ = [
      "_input_workers", "_element_spec", "_strategy", "_cardinality",
      "_enable_get_next_as_optional", "_options", "_canonicalize_devices"
  ]

  def __init__(
      self,
      input_workers,
      element_spec,
      strategy,
      options,
      cardinality=cardinality_lib.UNKNOWN,
      enable_get_next_as_optional=None,
      replica_order=None,
  ):
    # We don't want to allow deserialization of this class because we don't
    # serialize the strategy object. Currently the only places where
    # _deserialize is called is when we save/restore using SavedModels.
    if isinstance(input_workers, tuple):
      raise NotImplementedError("DistributedIteratorSpec does not have support "
                                "for deserialization.")
    else:
      self._input_workers = input_workers
      self._element_spec = element_spec
      self._strategy = strategy
      self._cardinality = cardinality
      self._enable_get_next_as_optional = enable_get_next_as_optional
      self._options = options
      if self._strategy:
        self._canonicalize_devices = getattr(self._strategy,
                                             "_canonicalize_devices", True)
      else:
        self._canonicalize_devices = True
      self._replica_order = replica_order

  def _serialize(self):
    # We cannot serialize the strategy object so we convert it to an id that we
    # can use for comparison.
    return (self._input_workers.serialize(), self._element_spec,
            id(self._strategy), id(self._options))

  def _deserialize(self):
    raise ValueError(
        f"Deserialization is currently unsupported for {type(self)}.")

  def sanity_check_type(self, other):
    """Returns the most specific TypeSpec compatible with `self` and `other`.

    Args:
      other: A `TypeSpec`.

    Raises:
      ValueError: If there is no TypeSpec that is compatible with both `self`
        and `other`.
    """
    # pylint: disable=protected-access
    if type(self) is not type(other):
      raise ValueError("No TypeSpec is compatible with both %s and %s" %
                       (self, other))
    if self._input_workers.serialize() != other._input_workers.serialize():
      raise ValueError("_input_workers is not compatible with both %s "
                       "and %s" % (self, other))
    if self._strategy is not other._strategy:
      raise ValueError("tf.distribute strategy is not compatible with both %s "
                       "and %s" % (self, other))

  def is_subtype_of(self, other):
    """Returns True if `self` is subtype of `other`.

    Args:
      other: A `TypeSpec`.
    """
    try:
      self.sanity_check_type(other)
      nest.assert_same_structure(self._element_spec, other._element_spec)  # pylint: disable=protected-access
    except (TypeError, ValueError):
      return False

    self_elements = nest.flatten(self._element_spec)
    other_elements = nest.flatten(other._element_spec)  # pylint: disable=protected-access

    return all(
        self_element.is_subtype_of(other_element)
        for (self_element, other_element) in zip(self_elements, other_elements))

  def most_specific_common_supertype(self, others):
    """Returns the most specific supertype of `self` and `others`.

    Args:
      others: A Sequence of `TypeSpec`.

    Returns `None` if a supertype does not exist.
    """
    try:
      for other in others:
        self.sanity_check_type(other)
        nest.assert_same_structure(self._element_spec, other._element_spec)  # pylint: disable=protected-access
    except (TypeError, ValueError):
      return None

    self_elements = nest.flatten(self._element_spec)
    others_elements = [nest.flatten(other._element_spec) for other in others]  # pylint: disable=protected-access
    common_elements = [None] * len(self_elements)

    for i, self_element in enumerate(self_elements):
      common_elements[i] = self_element.most_specific_common_supertype(
          [other_elements[i] for other_elements in others_elements])
      if common_elements[i] is None:
        return None
    common_element_spec = nest.pack_sequence_as(self._element_spec,
                                                common_elements)
    return type(self)(
        self._input_workers,
        common_element_spec,
        self._strategy,
        self._options,
        cardinality=self._cardinality,
        enable_get_next_as_optional=self._enable_get_next_as_optional)

  def _with_tensor_ranks_only(self):
    element_spec = nest.map_structure(
        lambda s: s._with_tensor_ranks_only(),  # pylint: disable=protected-access
        self._element_spec)
    return type(self)(
        self._input_workers,
        element_spec,
        self._strategy,
        self._options,
        cardinality=self._cardinality,
        enable_get_next_as_optional=self._enable_get_next_as_optional)

  # TODO(b/206014848): Remove once names are not used.
  def _without_tensor_names(self):
    element_spec = nest.map_structure(
        lambda s: s._without_tensor_names(),  # pylint: disable=protected-access
        self._element_spec)
    return type(self)(
        self._input_workers,
        element_spec,
        self._strategy,
        self._options,
        cardinality=self._cardinality,
        enable_get_next_as_optional=self._enable_get_next_as_optional)


class DistributedIteratorSpec(DistributedDatasetAndIteratorSpec):
  """Type specification for `DistributedIterator`."""

  @property
  def value_type(self):
    return DistributedIterator

  @property
  def _component_specs(self):
    specs = []
    worker_device_pairs = self._input_workers._worker_device_pairs  # pylint: disable=protected-access

    for i, (input_device, compute_devices) in enumerate(worker_device_pairs):
      element_spec = nest.map_structure(
          functools.partial(_replace_per_replica_spec, i=i), self._element_spec)
      specs.append(
          _SingleWorkerDatasetIteratorSpec(input_device, compute_devices,
                                           element_spec, self._options,
                                           self._canonicalize_devices))
    return specs

  def _to_components(self, value):
    return value._iterators  # pylint: disable=protected-access

  def _from_components(self, components):
    return DistributedIterator(
        input_workers=self._input_workers,
        iterators=None,
        components=components,
        element_spec=self._element_spec,
        strategy=self._strategy,
        cardinality=self._cardinality,
        enable_get_next_as_optional=self._enable_get_next_as_optional,
        options=self._options,
        replica_order=self._replica_order,
    )

  @staticmethod
  def from_value(value):
    # pylint: disable=protected-access
    return DistributedIteratorSpec(
        value._input_workers,
        value._element_spec,
        value._strategy,
        value._options,
        cardinality=value._cardinality,
        enable_get_next_as_optional=value._enable_get_next_as_optional)


class DistributedIterator(DistributedIteratorBase,
                          composite_tensor.CompositeTensor):
  """Input Iterator for a distributed dataset."""

  def __init__(
      self,
      input_workers=None,
      iterators=None,
      strategy=None,
      components=None,
      element_spec=None,
      cardinality=cardinality_lib.UNKNOWN,
      enable_get_next_as_optional=False,
      options=None,
      replica_order=None,
  ):
    if input_workers is None:
      raise ValueError("`input_workers` should be "
                       "provided.")

    error_message = ("Either `input_workers` or "
                     "both `components` and `element_spec` need to be "
                     "provided.")
    self._options = options

    if iterators is None:
      if (components is None or element_spec is None):
        raise ValueError(error_message)
      self._element_spec = element_spec
      self._input_workers = input_workers
      self._iterators = components
      self._strategy = strategy
      self._cardinality = cardinality
      self._enable_get_next_as_optional = enable_get_next_as_optional
      self._replica_order = replica_order
    else:
      if (components is not None and element_spec is not None):
        raise ValueError(error_message)

      super(DistributedIterator, self).__init__(
          input_workers,
          iterators,
          strategy,
          cardinality,
          enable_get_next_as_optional,
          replica_order,
      )

  @property
  def element_spec(self):
    # When partial batch handling is enabled, always set the batch dimension to
    # None, otherwise we just follow element_spec of the underlying dataset
    # (whose batch dimension may also be None). This is because with partial
    # batching handling we could always produce empty batches.
    if (self._enable_get_next_as_optional and
        self._strategy.extended._in_multi_worker_mode()):  # pylint: disable=protected-access
      return nest.map_structure(
          _rebatch_as_dynamic, self._element_spec, expand_composites=False)
    return self._element_spec

  @property
  def _type_spec(self):
    # Note that we use actual element_spec instead of the rebatched-as-dynamic
    # one to create DistributedIteratorSpec, to be consistent with the
    # underlying iterators' specs.
    return DistributedIteratorSpec(
        self._input_workers,
        self._element_spec,
        self._strategy,
        self._options,
        self._cardinality,
        self._enable_get_next_as_optional,
        self._replica_order,
    )


class _IterableInput(collections_abc.Iterable,
                     distribute_types.DistributedDatasetInterface):
  """Base class for iterable inputs for distribution strategies."""

  # pylint: disable=super-init-not-called
  def __init__(self, input_workers):
    assert isinstance(input_workers, InputWorkers)
    self._input_workers = input_workers

  def __iter__(self):
    raise NotImplementedError("must be implemented in descendants")

  def reduce(self, initial_state, reduce_fn):
    """Execute a `reduce_fn` over all the elements of the input."""
    iterator = iter(self)
    optional_data = iterator.get_next_as_optional()

    def cond(optional_data, state):
      del state  # Unused.
      return optional_data.has_value()

    def loop_body(optional_data, state):
      """Executes `reduce_fn` in a loop till the dataset is empty."""
      state = reduce_fn(state, optional_data.get_value())
      optional_data = iterator.get_next_as_optional()
      return optional_data, state

    optional_data, final_state = while_loop.while_loop(
        cond,
        loop_body, [optional_data, initial_state],
        parallel_iterations=1,
        return_same_structure=True)
    return final_state


class DistributedDatasetSpec(DistributedDatasetAndIteratorSpec):
  """Type specification for `DistributedDataset."""

  @property
  def value_type(self):
    return DistributedDataset

  @property
  def _component_specs(self):
    specs = []
    worker_device_pairs = self._input_workers._worker_device_pairs  # pylint: disable=protected-access

    for i, _ in enumerate(worker_device_pairs):
      element_spec = nest.map_structure(
          functools.partial(_replace_per_replica_spec, i=i), self._element_spec)
      specs.append(dataset_ops.DatasetSpec(element_spec))
    return specs

  def _to_components(self, value):
    return value._cloned_datasets  # pylint: disable=protected-access

  def _from_components(self, components):
    return DistributedDataset(
        input_workers=self._input_workers,
        strategy=self._strategy,
        components=components,
        element_spec=self._element_spec,
        enable_get_next_as_optional=self._enable_get_next_as_optional,
        options=self._options,
        replica_order=self._replica_order,
    )

  @staticmethod
  def from_value(value):
    # pylint: disable=protected-access
    return DistributedDatasetSpec(
        value._input_workers,
        value._element_spec,
        value._strategy,
        value._options,
        enable_get_next_as_optional=value._enable_get_next_as_optional)
    # pylint: enable=protected-access


class DistributedDataset(_IterableInput, composite_tensor.CompositeTensor):
  """Distributed dataset that supports prefetching to multiple devices."""

  def __init__(
      self,
      input_workers,
      strategy,
      dataset=None,
      num_replicas_in_sync=None,
      input_context=None,
      components=None,
      element_spec=None,
      enable_get_next_as_optional=None,
      build=True,
      options=None,
      replica_order=None,
  ):
    """Distribute the dataset on all workers.

    If `num_replicas_in_sync` is not None, we split each batch of the dataset
    into `num_replicas_in_sync` smaller batches, to be distributed among that
    worker's replicas, so that the batch size for a global step (across all
    workers and replicas) is as expected.

    Args:
      input_workers: an `InputWorkers` object.
      strategy: a `tf.distribute.Strategy` object, used to run all-reduce to
        handle last partial batch.
      dataset: `tf.data.Dataset` that will be used as the input source. Either
        dataset or components field should be passed when constructing
        DistributedDataset. Use this when constructing DistributedDataset from a
        new `tf.data.Dataset`. Use components when constructing using
        DistributedDatasetSpec.
      num_replicas_in_sync: Optional integer. If this is not None, the value is
        used to decide how to rebatch datasets into smaller batches so that the
        total batch size for each step (across all workers and replicas) adds up
        to `dataset`'s batch size.
      input_context: `InputContext` for sharding. Only pass this in for between
        graph multi-worker cases where there is only one `input_worker`. In
        these cases, we will shard based on the `input_pipeline_id` and
        `num_input_pipelines` in the `InputContext`.
      components: datasets when DistributedDataset is constructed from
        DistributedDatasetSpec. Either field dataset or components should be
        passed.
      element_spec: element spec for DistributedDataset when constructing from
        DistributedDatasetSpec. This will be used to set the element_spec for
        DistributedDataset and verified against element_spec from components.
      enable_get_next_as_optional: this is required when components is passed
        instead of dataset.
      build: whether to build underlying datasets when this object is created.
        This is only useful for `ParameterServerStrategy` now.
      options: `tf.distribute.InputOptions` used to control options on how this
        dataset is distributed.
      replica_order: the order of the replicas, which will be used to reorder
        the iterators to match the device order.
    """
    super(DistributedDataset, self).__init__(input_workers=input_workers)
    if input_workers is None or strategy is None:
      raise ValueError("input_workers and strategy are required arguments")
    if dataset is not None and components is not None:
      raise ValueError("Only one of dataset or components should be present")
    if dataset is None and components is None:
      raise ValueError("At least one of dataset or components should be passed")

    self._input_workers = input_workers
    self._strategy = strategy
    self._options = options
    self._input_context = input_context
    self._num_replicas_in_sync = num_replicas_in_sync
    self._replica_order = replica_order

    if dataset is not None:
      self._original_dataset = dataset
      self._built = False
      if build:
        self.build()
    else:
      if not build:
        raise ValueError(
            "When constructing DistributedDataset with components, build "
            "should not be False. This is an internal error. Please file a "
            "bug.")
      if enable_get_next_as_optional is None:
        raise ValueError(
            "When constructing DistributedDataset with components, " +
            "enable_get_next_as_optional should also be passed")
      self._cloned_datasets = components
      self._cardinality = _cardinality(self._cloned_datasets[0])
      self._enable_get_next_as_optional = enable_get_next_as_optional

      assert element_spec is not None
      if element_spec != _create_distributed_tensor_spec(
          self._strategy, self._cloned_datasets[0].element_spec):
        raise ValueError("Mismatched element_spec from the passed components")
      self._element_spec = element_spec

      self._built = True

  def build(self, dataset_to_replace=None):
    assert not self._built
    dataset = dataset_to_replace or self._original_dataset
    self._cardinality = _cardinality(dataset)
    self._enable_get_next_as_optional = _enable_get_next_as_optional(
        self._strategy, dataset, self._cardinality)
    distribute_start_time_ns = time.time_ns()
    self._create_cloned_datasets_from_dataset(dataset, self._input_context,
                                              self._input_workers,
                                              self._strategy,
                                              self._num_replicas_in_sync)
    if context.executing_eagerly():
      # Records the time to initialize the distributed dataset.
      context.async_wait()
      distribute_duration_ms = (time.time_ns() -
                                distribute_start_time_ns) // 1_000_000
      _distributed_dataset_initialization_time_milliseconds.get_cell(
          self._strategy.__class__.__name__,
          str(self._input_workers.num_workers)).add(distribute_duration_ms)
    self._element_spec = _create_distributed_tensor_spec(
        self._strategy, self._cloned_datasets[0].element_spec)
    self._built = True

  def auto_shard(self, num_shards, shard_ix):
    assert (
        len(self._cloned_datasets) == len(self._input_workers.worker_devices)
    ), (
        f"datasets: {len(self._cloned_datasets)}, "
        f"input workers: {len(self._input_workers.worker_devices)}"
    )
    sharded_datasets = []
    for i in range(len(self._input_workers.worker_devices)):
      with ops.colocate_with(self._cloned_datasets[i]._variant_tensor):  # pylint:disable=protected-access
        sharded_datasets.append(
            input_ops.auto_shard_dataset(
                self._cloned_datasets[i], num_shards, shard_ix,
                self._num_replicas_in_sync
                ))
    return DistributedDataset(
        self._input_workers,
        self._strategy,
        components=sharded_datasets,
        element_spec=self._element_spec,
        options=self._options,
        enable_get_next_as_optional=self._enable_get_next_as_optional)

  @property
  def cardinality(self):
    if not self._built:
      raise ValueError(
          "Cannot get the cardinality of a dataset that is not built")
    return self._cardinality

  def _create_cloned_datasets_from_dataset(self, dataset, input_context,
                                           input_workers, strategy,
                                           num_replicas_in_sync):
    # We clone and shard the dataset on each worker. The current setup tries to
    # shard the dataset by files if possible so that each worker sees a
    # different subset of files. If that is not possible, will attempt to shard
    # the final input such that each worker will run the entire preprocessing
    # pipeline and only receive its own shard of the dataset.

    # Additionally, we rebatch the dataset on each worker into
    # `num_replicas_in_sync` smaller batches to be distributed among that
    # worker's replicas, so that the batch size for a global step (across all
    # workers and replicas) adds up to the original dataset's batch size.
    if num_replicas_in_sync is not None and num_replicas_in_sync > 1:
      num_workers = input_context.num_input_pipelines if input_context else len(
          input_workers.worker_devices)
      rebatch_fn = self._make_rebatch_fn(dataset, num_workers,
                                         num_replicas_in_sync)
    else:
      rebatch_fn = None
    self._cloned_datasets = []
    if input_context:
      # Between-graph where we rely on the input_context for sharding
      assert input_workers.num_workers == 1
      if rebatch_fn is not None:
        dataset = rebatch_fn(dataset, input_context.input_pipeline_id)
      dataset = input_ops.auto_shard_dataset(dataset,
                                             input_context.num_input_pipelines,
                                             input_context.input_pipeline_id,
                                             num_replicas_in_sync)
      self._cloned_datasets.append(dataset)
    else:
      replicated_ds = distribute.replicate(dataset,
                                           input_workers.worker_devices)
      for i, worker in enumerate(input_workers.worker_devices):
        with ops.device(worker):
          cloned_dataset = replicated_ds[worker]
          if rebatch_fn is not None:
            cloned_dataset = rebatch_fn(cloned_dataset, i)
          cloned_dataset = input_ops.auto_shard_dataset(
              cloned_dataset, len(input_workers.worker_devices), i,
              num_replicas_in_sync)
          self._cloned_datasets.append(cloned_dataset)

  def _make_rebatch_fn(self, dataset, num_workers, num_replicas_in_sync):
    """Returns a callable that rebatches the input dataset.

    Args:
      dataset: A `tf.data.Dataset` representing the dataset to be distributed.
      num_workers: An integer representing the number of workers to distribute
        `dataset` among.
      num_replicas_in_sync: An integer representing the number of replicas in
        sync across all workers.
    """
    if num_replicas_in_sync % num_workers:
      raise ValueError(
          "tf.distribute expects every worker to have the same number of "
          "replicas. However, encountered `num_replicas_in_sync` ({}) that "
          "cannot be divided by `num_workers` ({})".format(
              num_replicas_in_sync, num_workers))

    num_replicas_per_worker = num_replicas_in_sync // num_workers
    with ops.colocate_with(dataset._variant_tensor):  # pylint: disable=protected-access
      batch_size = distribute.compute_batch_size(dataset)

    def rebatch_fn(dataset, worker_index):
      try:

        def apply_rebatch():
          batch_sizes = distribute.batch_sizes_for_worker(
              batch_size, num_workers, num_replicas_per_worker, worker_index)
          return dataset.rebatch(batch_sizes).prefetch(num_replicas_per_worker)

        # pylint: disable=protected-access
        def apply_legacy_rebatch():
          return distribute._LegacyRebatchDataset(
              dataset, num_replicas_in_sync).prefetch(num_replicas_per_worker)

        with ops.colocate_with(dataset._variant_tensor):
          return tf_cond.cond(
              math_ops.not_equal(batch_size, -1),
              true_fn=apply_rebatch,
              false_fn=apply_legacy_rebatch)
      except errors.InvalidArgumentError as e:
        if "without encountering a batch" in str(e):
          six.reraise(
              ValueError,
              ValueError(
                  "Call the `batch` method on the input Dataset in order to be "
                  "able to split your input across {} replicas.\n Please see "
                  "the tf.distribute.Strategy guide. {}".format(
                      num_replicas_in_sync, e)),
              sys.exc_info()[2])
        else:
          raise

    return rebatch_fn

  def __iter__(self):
    if not (context.executing_eagerly() or
            ops.get_default_graph().building_function):
      raise RuntimeError("__iter__() is only supported inside of tf.function "
                         "or when eager execution is enabled.")
    if not self._built:
      raise ValueError("To use this dataset, you need to pass this dataset to "
                       "ClusterCoordinator.create_per_worker_dataset.")

    canonicalize_devices = getattr(self._strategy, "_canonicalize_devices",
                                   True)

    worker_iterators = _create_iterators_per_worker(
        self._cloned_datasets,
        self._input_workers,
        options=self._options,
        canonicalize_devices=canonicalize_devices)
    iterator = DistributedIterator(
        self._input_workers,
        worker_iterators,
        self._strategy,
        cardinality=self._cardinality,
        enable_get_next_as_optional=self._enable_get_next_as_optional,
        options=self._options,
        replica_order=self._replica_order,
    )
    iterator._element_spec = self._element_spec  # pylint: disable=protected-access

    # When async eager is enabled, sometimes the iterator may not finish
    # initialization before passing to a multi device function, add a sync point
    # here to make sure all underlying iterators are initialized.
    if context.executing_eagerly():
      context.async_wait()

    return iterator

  @property
  def element_spec(self):
    """The type specification of an element of this dataset."""
    # When partial batch handling is enabled, always set the batch dimension to
    # None, otherwise we just follow element_spec of the underlying dataset
    # (whose batch dimension may also be None). This is because with partial
    # batching handling we could always produce empty batches.
    if (self._enable_get_next_as_optional and
        self._strategy.extended._in_multi_worker_mode()):  # pylint: disable=protected-access
      return nest.map_structure(
          _rebatch_as_dynamic, self._element_spec, expand_composites=False)
    return self._element_spec

  @property
  def _type_spec(self):
    return DistributedDatasetSpec(
        self._input_workers,
        self._element_spec,
        self._strategy,
        self._options,
        enable_get_next_as_optional=self._enable_get_next_as_optional)


class DistributedDatasetsFromFunctionSpec(DistributedDatasetAndIteratorSpec):
  """Type specification for `DistributedDatasetsFromFunction."""

  @property
  def value_type(self):
    return DistributedDatasetsFromFunction

  @property
  def _component_specs(self):
    specs = []
    worker_device_pairs = self._input_workers._worker_device_pairs  # pylint: disable=protected-access

    for i, _ in enumerate(worker_device_pairs):
      element_spec = nest.map_structure(
          functools.partial(_replace_per_replica_spec, i=i), self._element_spec)
      specs.append(dataset_ops.DatasetSpec(element_spec))
    return specs

  def _to_components(self, value):
    return value._datasets  # pylint: disable=protected-access

  def _from_components(self, components):
    return DistributedDatasetsFromFunction(
        input_workers=self._input_workers,
        strategy=self._strategy,
        components=components,
        element_spec=self._element_spec,
        options=self._options)

  @staticmethod
  def from_value(value):
    # pylint: disable=protected-access
    return DistributedDatasetsFromFunctionSpec(
        input_workers=value._input_workers,
        element_spec=value._element_spec,
        strategy=value._strategy,
        options=value._options)


# TODO(priyag): Add other replication modes.
class DistributedDatasetsFromFunction(_IterableInput,
                                      composite_tensor.CompositeTensor):
  """Inputs created from dataset function."""

  def __init__(
      self,
      input_workers,
      strategy,
      input_contexts=None,
      dataset_fn=None,
      options=None,
      components=None,
      element_spec=None,
      build=True,
      replica_order=None,
  ):
    """Makes an iterable from datasets created by the given function.

    Args:
      input_workers: an `InputWorkers` object.
      strategy: a `tf.distribute.Strategy` object, used to run all-reduce to
        handle last partial batch.
      input_contexts: A list of `InputContext` instances to be passed to call(s)
        to `dataset_fn`. Length and order should match worker order in
        `worker_device_pairs`.
      dataset_fn: A function that returns a `Dataset` given an `InputContext`.
        Either dataset_fn or components should be passed to construct
        DistributedDatasetsFromFunction. Use this when constructing
        DistributedDataset using a function. Use components when constructing
        using DistributedDatasetsFromFunctionSpec.
      options: `tf.distribute.InputOptions` used to control options on how this
        dataset is distributed.
      components: datasets when DistributedDatasetsFromFunction is constructed
        from DistributedDatasetsFromFunctionSpec. Only one of dataset or
        components should be passed.
      element_spec: element spec for DistributedDataset when constructing from
        DistributedDatasetSpec. This will be used to set the element_spec for
        DistributedDatasetsFromFunctionSpec and verified against element_spec
        from components.
      build: whether to build underlying datasets when this object is created.
        This is only useful for `ParameterServerStrategy` now.
      replica_order: the order of the replicas, which will be used to reorder
        the iterators to match the device order.
    """
    super(DistributedDatasetsFromFunction, self).__init__(
        input_workers=input_workers)
    self._input_workers = input_workers
    self._strategy = strategy
    self._options = options
    self._replica_order = replica_order
    if dataset_fn is not None and components is not None:
      raise ValueError("Only one of dataset_fn or components should be set")
    if dataset_fn is None and components is None:
      raise ValueError("At least one of dataset_fn or components should be set")

    if dataset_fn is not None:
      if input_workers.num_workers != len(input_contexts):
        raise ValueError(
            "Number of input workers (%d) is not same as number of "
            "input_contexts (%d)" %
            (input_workers.num_workers, len(input_contexts)))
      self._input_contexts = input_contexts
      self._num_replicas_in_sync = self._input_contexts[0].num_replicas_in_sync
      self._dataset_fn = dataset_fn
      self._built = False
      if build:
        self.build()
    else:
      if element_spec is None:
        raise ValueError(
            "element_spec should also be passed when passing components")
      if not build:
        raise ValueError(
            "When constructing DistributedDatasetFromFunction with components, "
            "build should not be False. This is an internal error. Please file "
            "a bug.")
      self._element_spec = element_spec
      self._datasets = components
      self._num_replicas_in_sync = None
      self._built = True
      self._cardinality = _cardinality(self._datasets[0])
      self._enable_get_next_as_optional = _enable_get_next_as_optional(
          self._strategy, self._datasets[0], self._cardinality)

  def build(self):
    assert not self._built
    distribute_start_time_ns = time.time_ns()
    self._datasets, element_spec = (
        _create_datasets_from_function_with_input_context(
            self._input_contexts, self._input_workers, self._dataset_fn))
    if context.executing_eagerly():
      # Records the time to initialize the distributed dataset.
      context.async_wait()
      distribute_duration_ms = (time.time_ns() -
                                distribute_start_time_ns) // 1_000_000
      _distributed_dataset_from_function_initialization_time_milliseconds.get_cell(
          self._strategy.__class__.__name__,
          str(self._input_workers.num_workers)).add(distribute_duration_ms)

    self._element_spec = _create_distributed_tensor_spec(
        self._strategy, element_spec)
    self._cardinality = _cardinality(self._datasets[0])
    self._enable_get_next_as_optional = _enable_get_next_as_optional(
        self._strategy, self._datasets[0], self._cardinality)
    self._built = True

  def auto_shard(self, num_shards, shard_ix):
    assert (
        len(self._datasets) == len(self._input_workers.worker_devices)
    ), (
        f"datasets: {len(self._datasets)}, "
        f"input workers: {len(self._input_workers.worker_devices)}"
    )
    sharded_datasets = []
    for i in range(len(self._input_workers.worker_devices)):
      with ops.colocate_with(self._datasets[i]._variant_tensor):  # pylint: disable=protected-access
        sharded_datasets.append(
            input_ops.auto_shard_dataset(
                self._datasets[i], num_shards, shard_ix,
                self._num_replicas_in_sync
            )
        )
    return DistributedDatasetsFromFunction(self._input_workers, self._strategy,
                                           components=sharded_datasets,
                                           element_spec=self._element_spec,
                                           options=self._options)

  @property
  def cardinality(self):
    if not self._built:
      raise ValueError(
          "Cannot get the cardinality of a dataset that is not built")
    return self._cardinality

  def __iter__(self):
    if not (ops.executing_eagerly_outside_functions() or
            ops.get_default_graph().building_function):
      raise RuntimeError("__iter__() is only supported inside of tf.function "
                         "or when eager execution is enabled.")

    if not self._built:
      raise ValueError("You need to use this dataset in "
                       "ClusterCoordinator.create_per_worker_dataset.")

    canonicalize_devices = getattr(self._strategy, "_canonicalize_devices",
                                   True)

    iterators = _create_iterators_per_worker(
        self._datasets,
        self._input_workers,
        options=self._options,
        canonicalize_devices=canonicalize_devices)
    iterator = DistributedIterator(
        input_workers=self._input_workers,
        iterators=iterators,
        strategy=self._strategy,
        cardinality=self._cardinality,
        enable_get_next_as_optional=self._enable_get_next_as_optional,
        options=self._options,
        replica_order=self._replica_order,
    )
    iterator._element_spec = self._element_spec  # pylint: disable=protected-access

    # When async eager is enabled, sometimes the iterator may not finish
    # initialization before passing to a multi device function, add a sync
    # point here to make sure all underlying iterators are initialized.
    if context.executing_eagerly():
      context.async_wait()

    return iterator

  @property
  def element_spec(self):
    """The type specification of an element of this dataset."""
    # When partial batch handling is enabled, always set the batch dimension to
    # None, otherwise we just follow element_spec of the underlying dataset
    # (whose batch dimension may also be None). This is because with partial
    # batching handling we could always produce empty batches.
    if (self._enable_get_next_as_optional and
        self._strategy.extended._in_multi_worker_mode()):  # pylint: disable=protected-access
      return nest.map_structure(
          _rebatch_as_dynamic, self._element_spec, expand_composites=False)
    return self._element_spec

  @property
  def _type_spec(self):
    return DistributedDatasetsFromFunctionSpec(self._input_workers,
                                               self._element_spec,
                                               self._strategy, self._options)


def _dummy_tensor_fn(value_structure):
  """A function to create dummy tensors from `value_structure`."""

  def create_dummy_tensor(spec):
    """Create a dummy tensor with possible batch dimensions set to 0."""
    if hasattr(spec, "_create_empty_value"):
      # Type spec may overwrite default dummy values behavior by declaring the
      # `_create_empty_value(self)` method. This method must return a value
      # compatible with the type spec with batch dimensions set to 0 or fail if
      # such a value does not exist. This allows a composite tensor to customize
      # dummy values creation as, in general, its dummy value is not composed
      # from dummy components (e.g. `row_splits` tensor of a RaggedTensor is
      # never allowed to be empty). See b/183969859 for more discussions.
      # TODO(b/186079336): reconsider CompositeTensor support.
      return spec._create_empty_value()  # pylint: disable=protected-access

    if isinstance(spec, ragged_tensor.RaggedTensorSpec):
      # Splice out the ragged dimensions.
      # pylint: disable=protected-access
      feature_shape = spec._shape[:1].concatenate(
          spec._shape[(1 + spec._ragged_rank):])
      feature_type = spec._dtype
      # pylint: enable=protected-access
    else:
      feature_shape = spec.shape
      feature_type = spec.dtype
    # Ideally we should set the batch dimension to 0, however as in
    # DistributionStrategy we don't know the batch dimension, we try to
    # guess it as much as possible. If the feature has unknown dimensions, we
    # will set them to 0. If the feature shape is already static, we guess the
    # first dimension as batch dimension and set it to 0.
    dims = ([dim if dim is not None else 0 for dim in feature_shape.as_list()]
            if feature_shape else [])
    if dims and (isinstance(spec, ragged_tensor.RaggedTensorSpec) or
                 feature_shape.is_fully_defined()):
      dims[0] = tensor_shape.Dimension(0)

    if isinstance(spec, sparse_tensor.SparseTensorSpec):
      return sparse_tensor.SparseTensor(
          values=array_ops.zeros(0, feature_type),
          indices=array_ops.zeros((0, len(dims)), dtypes.int64),
          dense_shape=dims)

    # Create the dummy tensor.
    dummy_tensor = array_ops.zeros(tensor_shape.TensorShape(dims), feature_type)
    if isinstance(spec, ragged_tensor.RaggedTensorSpec):
      # Reinsert the ragged dimensions with size 0.
      # pylint: disable=protected-access
      row_splits = array_ops.zeros(1, spec._row_splits_dtype)
      dummy_tensor = ragged_tensor.RaggedTensor.from_nested_row_splits(
          dummy_tensor, (row_splits,) * spec._ragged_rank, validate=False)
      # pylint: enable=protected-access
    return dummy_tensor

  return nest.map_structure(create_dummy_tensor, value_structure)


def _get_value_or_dummy(input_workers, optional_list, produce_dummy):
  """Returns the value of the optionals or dummy values.

  Args:
    input_workers: the `InputWorkers`.
    optional_list: a list of lists `tf.experimental.Optional`. The values from
      each compute device grouped by the input device.
    produce_dummy: a bool. Whether to produce dummy tensors when the optional
      doesn't have a value.

  Returns:
    A flatten list of Tensors.

  """
  value_list = []
  for i, worker in enumerate(input_workers.worker_devices):
    with ops.device(worker):
      devices = input_workers.compute_devices_for_worker(i)
      for j, device in enumerate(devices):
        with ops.device(device):
          if produce_dummy:
            # pylint: disable=cell-var-from-loop
            value_list.append(
                tf_cond.cond(
                    optional_list[i][j].has_value(),
                    lambda: optional_list[i][j].get_value(),  # pylint: disable=unnecessary-lambda
                    lambda: _dummy_tensor_fn(optional_list[i][j].element_spec),
                    strict=True,
                ))
            # pylint: enable=cell-var-from-loop
          else:
            value_list.append(optional_list[i][j].get_value())
  return value_list


class _SingleWorkerDatasetIteratorBase(object):
  """Iterator for a single `tf.data.Dataset`."""

  def __init__(self, dataset, worker, devices, options=None):
    """Create iterator for the `dataset` to fetch data to worker's `devices` .

    A `MultiDeviceIterator`  or `OwnedMultiDeviceIterator` is used to prefetch
    input to the devices on the given worker.

    Args:
      dataset: A `tf.data.Dataset` instance.
      worker: Worker on which ops should be created.
      devices: Distribute data from `dataset` to these devices.
      options: options.
    """
    self._dataset = dataset
    self._worker = worker
    self._devices = devices
    self._element_spec = dataset.element_spec
    self._options = options
    self._make_iterator()

  def _make_iterator(self):
    raise NotImplementedError("must be implemented in descendants")

  def _format_data_list_with_options(self, data_list):
    """Change the data in to a list type if required.

    The OwnedMultiDeviceIterator returns the list data type,
    while the PER_REPLICA iterator (when used with prefetch disabled)
    returns without the enclosed list. This is to fix the inconsistency.
    Args:
      data_list: data_list
    Returns:
      list
    """
    if (self._options and self._options.experimental_replication_mode ==
        InputReplicationMode.PER_REPLICA and
        not self._options.experimental_fetch_to_device):
      return [data_list]
    else:
      return data_list

  def get_next(self, device, name=None):
    """Get next element for the given device."""
    del name
    with ops.device(self._worker):
      if _should_use_multi_device_iterator(self._options):
        return self._iterator.get_next(device)
      else:
        return self._iterator.get_next()

  def get_next_as_list(self, name=None):
    """Get next element from the underlying iterator.

    Runs the iterator get_next() within a device scope. Since this doesn't use
    get_next_as_optional(), it is considerably faster than get_next_as_list(),
    but it raises EOFError if any of the device doesn't get any data.

    Args:
      name: not used.

    Returns:
      A list consisting of the next data from each device.
    """
    del name
    with ops.device(self._worker):
      return self._format_data_list_with_options(self._iterator.get_next())

  def get_next_as_optional_list(self):
    with ops.device(self._worker):
      return self._format_data_list_with_options(
          self._iterator.get_next_as_optional())


class _SingleWorkerDatasetIteratorSpec(type_spec.TypeSpec):
  """Type specification for `_SingleWorkerOwnedDatasetIterator`."""

  __slots__ = [
      "_worker", "_devices", "_element_spec", "_options",
      "_canonicalize_devices"
  ]

  def __init__(self, worker, devices, element_spec, options,
               canonicalize_devices=True):
    self._worker = worker
    if canonicalize_devices:
      self._devices = tuple(device_util.canonicalize(d) for d in devices)
    else:
      self._devices = tuple(
          device_util.canonicalize_without_job_and_task(d) for d in devices)
    self._element_spec = element_spec
    # `self._options` intentionally made not `None` for proper serialization.
    self._options = (options if options is not None else
                     distribute_lib.InputOptions())
    self._canonicalize_devices = canonicalize_devices

  @property
  def value_type(self):
    return _SingleWorkerOwnedDatasetIterator

  def _serialize(self):
    return (self._worker, self._devices, self._element_spec, self._options,
            self._canonicalize_devices)

  def _get_multi_device_iterator_spec(self, specs):
    device_scope = device_util.canonicalize(self._worker, device_util.current())
    host_device = device_util.get_host_for_device(device_scope)
    # source_device while creating iterator governs the worker device in
    # iterator spec.
    worker = host_device
    specs.append(
        multi_device_iterator_ops.MultiDeviceIteratorSpec(
            self._devices, worker, element_spec=self._element_spec))

  @property
  def _component_specs(self):
    specs = []
    if _should_use_multi_device_iterator(self._options):
      self._get_multi_device_iterator_spec(specs)
    else:
      specs.append(iterator_ops.IteratorSpec(element_spec=self._element_spec))
    return specs

  def _to_components(self, value):
    return [value._iterator]  # pylint: disable=protected-access

  def _from_components(self, components):
    return _SingleWorkerOwnedDatasetIterator(
        dataset=None,
        worker=self._worker,
        devices=self._devices,
        components=components,
        element_spec=self._element_spec,
        options=self._options,
        canonicalize_devices=self._canonicalize_devices)

  @staticmethod
  def from_value(value):
    # pylint: disable=protected-access
    return _SingleWorkerDatasetIteratorSpec(value._worker, value._devices,
                                            value._element_spec, value._options,
                                            value._canonicalize_devices)


class _SingleWorkerOwnedDatasetIterator(_SingleWorkerDatasetIteratorBase,
                                        composite_tensor.CompositeTensor):
  """Iterator for a DistributedDataset instance."""

  def __init__(self,
               dataset=None,
               worker=None,
               devices=None,
               components=None,
               element_spec=None,
               options=None,
               canonicalize_devices=None):
    """Create iterator for the `dataset` to fetch data to worker's `devices` .

    `OwnedMultiDeviceIterator` is used to prefetch input to the devices on the
    given worker. The lifetime of this iterator is tied to the encompassing
    python object. Once we go out of scope of the python object or return from
    a tf.function the underlying iterator resource is deleted.

    Args:
      dataset: A `tf.data.Dataset` instance.
      worker: Worker on which ops should be created.
      devices: Distribute data from `dataset` to these devices.
      components: Tensor components to construct the
        _SingleWorkerOwnedDatasetIterator from.
      element_spec: A nested structure of `TypeSpec` objects that represents the
      type specification of elements of the iterator.
      options: `tf.distribute.InputOptions` used to control options on how this
      dataset is distributed.
      canonicalize_devices: Whether to canonicalize devices for workers fully or
      partially. If False, it will partially canonicalize devices by removing
      job and task.
    """
    if worker is None or devices is None:
      raise ValueError("Both `worker` and `devices` should be provided")

    error_message = ("Either `dataset` or both `components` and `element_spec` "
                     "need to be provided.")

    self._options = options
    self._canonicalize_devices = canonicalize_devices
    if dataset is None:
      if (components is None or element_spec is None):
        raise ValueError(error_message)
      self._element_spec = element_spec
      self._worker = worker
      self._devices = devices
      self._iterator = components[0]
    else:
      if (components is not None or element_spec is not None):
        raise ValueError(error_message)
      super(_SingleWorkerOwnedDatasetIterator,
            self).__init__(dataset, worker, devices, self._options)

  def _create_owned_multi_device_iterator(self):
    # If the worker devices are already canonicalized, canonicalizing again
    # would have no impact.
    # For strategies running on remote workers such as PS Strategy, the device
    # scope will be derived from current worker, if used under init_scope().
    if not ops.inside_function():
      device_scope = device_util.canonicalize(self._worker,
                                              device_util.current())
      host_device = device_util.get_host_for_device(device_scope)
    else:
      # In general, iterators should not be created within tf.functions. For
      # exact visitation guarantee solutions for parameter server training,
      # however, we do create iterators within the tf.functions that are
      # dispatched to workers. In these cases, the traced device must match the
      # runtime device. Since tracing occurs on the chief, we do not want to use
      # the current device scope, which would be the chief, but rather use the
      # relative worker device scope explicitly.
      device_scope, host_device = self._worker, self._worker
    with ops.device(device_scope):
      if self._options is not None:
        self._iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator(
            self._dataset,
            self._devices,
            source_device=host_device,
            max_buffer_size=self._options
            .experimental_per_replica_buffer_size,
            prefetch_buffer_size=self._options
            .experimental_per_replica_buffer_size)
      else:
        self._iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator(
            self._dataset, self._devices, source_device=host_device)

  def _make_iterator(self):
    """Make appropriate iterator on the dataset."""
    if not self._worker:
      raise ValueError("Worker device must be specified when creating an "
                       "owned iterator.")
    if _should_use_multi_device_iterator(self._options):
      self._create_owned_multi_device_iterator()
    else:
      with ops.device(self._worker):
        self._iterator = iter(self._dataset)

  @property
  def element_spec(self):
    return self._element_spec

  @property
  def _type_spec(self):
    return _SingleWorkerDatasetIteratorSpec(self._worker, self._devices,
                                            self._element_spec, self._options,
                                            self._canonicalize_devices)

  @property
  def output_classes(self):
    """Returns the class of each component of an element of this iterator.

    The expected values are `tf.Tensor` and `tf.SparseTensor`.

    Returns:
      A nested structure of Python `type` objects corresponding to each
      component of an element of this dataset.
    """
    return nest.map_structure(
        lambda component_spec: component_spec._to_legacy_output_classes(),  # pylint: disable=protected-access
        self._element_spec)

  @property
  def output_shapes(self):
    """Returns the shape of each component of an element of this iterator.

    Returns:
      A nested structure of `tf.TensorShape` objects corresponding to each
      component of an element of this dataset.
    """
    return nest.map_structure(
        lambda component_spec: component_spec._to_legacy_output_shapes(),  # pylint: disable=protected-access
        self._element_spec)

  @property
  def output_types(self):
    """Returns the type of each component of an element of this iterator.

    Returns:
      A nested structure of `tf.DType` objects corresponding to each component
      of an element of this dataset.
    """
    return nest.map_structure(
        lambda component_spec: component_spec._to_legacy_output_types(),  # pylint: disable=protected-access
        self._element_spec)


def _create_iterators_per_worker(worker_datasets,
                                 input_workers,
                                 options=None,
                                 canonicalize_devices=False):
  """Create a multidevice iterator on each of the workers."""
  assert isinstance(input_workers, InputWorkers)
  assert len(worker_datasets) == len(input_workers.worker_devices)
  iterators = []
  for i, worker in enumerate(input_workers.worker_devices):
    with ops.device(worker):
      worker_devices = input_workers.compute_devices_for_worker(i)
      iterator = _SingleWorkerOwnedDatasetIterator(
          dataset=worker_datasets[i],
          worker=worker,
          devices=worker_devices,
          options=options,
          canonicalize_devices=canonicalize_devices)
      iterators.append(iterator)
  return iterators


def _create_datasets_from_function_with_input_context(input_contexts,
                                                      input_workers,
                                                      dataset_fn):
  """Create device datasets per worker given a dataset function."""
  datasets = []
  for i, ctx in enumerate(input_contexts):
    worker = input_workers.worker_devices[i]
    with ops.device(worker):
      dataset = dataset_fn(ctx)
      datasets.append(dataset)
  return datasets, dataset.element_spec


# TODO(sourabhbajaj): Remove this in lieu of distributed datasets
def _get_batched_dataset(d):
  """Get the batched dataset from `d`."""
  # pylint: disable=protected-access
  if isinstance(d, dataset_ops.DatasetV1Adapter):
    d = d._dataset

  if isinstance(d, (dataset_ops.BatchDataset, batching._MapAndBatchDataset)):
    return d
  elif isinstance(d, (dataset_ops.PrefetchDataset,
                      dataset_ops._OptionsDataset)):
    return _get_batched_dataset(d._input_dataset)

  raise ValueError(
      "Unable to get batched dataset from the input dataset. `batch` "
      "`map_and_batch` need to be the last operations on the dataset. "
      "The batch operations can be followed by a prefetch.")


def _get_batched_dataset_attributes(d):
  """Get `batch_size`, `drop_remainder` of dataset."""
  # pylint: disable=protected-access
  assert isinstance(d,
                    (dataset_ops.BatchDataset, batching._MapAndBatchDataset))
  if isinstance(d, dataset_ops.BatchDataset):
    batch_size = d._batch_size
    drop_remainder = d._drop_remainder
  elif isinstance(d, batching._MapAndBatchDataset):
    batch_size = d._batch_size_t
    drop_remainder = d._drop_remainder_t
  # pylint: enable=protected-access

  if tensor_util.is_tf_type(batch_size):
    batch_size = tensor_util.constant_value(batch_size)

  if tensor_util.is_tf_type(drop_remainder):
    drop_remainder = tensor_util.constant_value(drop_remainder)

  return batch_size, drop_remainder


# TODO(sourabhbajaj): Remove this in lieu of distributed datasets
def _get_dataset_attributes(dataset):
  """Get the underlying attributes from the dataset object."""
  # pylint: disable=protected-access

  # First, get batch_size and drop_remainder from the dataset. We need
  # to walk back the dataset creation process and find the batched version in
  # order to get the attributes.
  batched_dataset = _get_batched_dataset(dataset)
  batch_size, drop_remainder = _get_batched_dataset_attributes(batched_dataset)

  # Second, prefetch buffer should be get from the original dataset.
  prefetch_buffer = None
  if isinstance(dataset, dataset_ops.PrefetchDataset):
    prefetch_buffer = dataset._buffer_size
  elif (isinstance(dataset, dataset_ops.DatasetV1Adapter)
        and isinstance(dataset._dataset, dataset_ops.PrefetchDataset)):
    prefetch_buffer = dataset._dataset._buffer_size

  return batch_size, drop_remainder, prefetch_buffer


def _should_use_multi_device_iterator(options):
  """Determine whether to use multi_device_iterator_ops."""
  if (options is None or
      options.experimental_replication_mode == InputReplicationMode.PER_WORKER
      or
      (options.experimental_replication_mode == InputReplicationMode.PER_REPLICA
       and options.experimental_fetch_to_device)):
    return True
  return False


class MultiStepContext(object):
  """A context object that can be used to capture things when running steps.

  This context object is useful when running multiple steps at a time using the
  `experimental_run_steps_on_iterator` API. For e.g. it allows the user's step
  function to specify which outputs to emit at what frequency. Currently it
  supports capturing output from the last step, as well as capturing non tensor
  outputs.  In the future it will be augmented to support other use cases such
  as output each N steps.
  """

  def __init__(self):
    """Initialize an output context.

    Returns:
      A context object.
    """
    self._last_step_outputs = {}
    self._last_step_outputs_reduce_ops = {}
    self._non_tensor_outputs = {}

  @property
  def last_step_outputs(self):
    """A dictionary consisting of outputs to be captured on last step.

    Keys in the dictionary are names of tensors to be captured, as specified
    when `set_last_step_output` is called.
    Values in the dictionary are the tensors themselves. If
    `set_last_step_output` was called with a `reduce_op` for this output,
    then the value is the reduced value.

    Returns:
      A dictionary with last step outputs.
    """
    return self._last_step_outputs

  def _set_last_step_outputs(self, outputs):
    """Replace the entire dictionary of last step outputs."""
    if not isinstance(outputs, dict):
      raise ValueError("Need a dictionary to set last_step_outputs.")
    self._last_step_outputs = outputs

  def set_last_step_output(self, name, output, reduce_op=None):
    """Set `output` with `name` to be outputted from the last step.

    Args:
      name: String, name to identify the output. Doesn't need to match tensor
        name.
      output: The tensors that should be outputted with `name`. See below for
        actual types supported.
      reduce_op: Reduction method to use to reduce outputs from multiple
        replicas. Required if `set_last_step_output` is called in a replica
        context. Optional in cross_replica_context.
        When present, the outputs from all the replicas are reduced using the
        current distribution strategy's `reduce` method. Hence, the type of
        `output` must be what's supported by the corresponding `reduce` method.
        For e.g. if using MirroredStrategy and reduction is set, output
        must be a `PerReplica` value.
        The reduce method is also recorded in a dictionary
        `_last_step_outputs_reduce_ops` for later interpreting of the
        outputs as already reduced or not.
    """
    if distribute_lib.in_cross_replica_context():
      self._last_step_outputs_reduce_ops[name] = reduce_op
      if reduce_op is None:
        self._last_step_outputs[name] = output
      else:
        distribution = distribute_lib.get_strategy()
        self._last_step_outputs[name] = distribution.reduce(reduce_op, output,
                                                            axis=None)
    else:
      assert reduce_op is not None
      def merge_fn(distribution, value):
        self._last_step_outputs[name] = distribution.reduce(reduce_op, value,
                                                            axis=None)
        # Setting this inside the `merge_fn` because all replicas share the same
        # context object, so it's more robust to set it only once (even if all
        # the replicas are trying to set the same value).
        self._last_step_outputs_reduce_ops[name] = reduce_op

      distribute_lib.get_replica_context().merge_call(
          merge_fn, args=(output,))

  @property
  def non_tensor_outputs(self):
    """A dictionary consisting of any non tensor outputs to be captured."""
    return self._non_tensor_outputs

  def set_non_tensor_output(self, name, output):
    """Set `output` with `name` to be captured as a non tensor output."""
    if distribute_lib.in_cross_replica_context():
      self._non_tensor_outputs[name] = output
    else:
      def merge_fn(distribution, value):
        # NOTE(priyag): For non tensor outputs, we simply return all the values
        # in a list as reduction doesn't make sense on non tensors.
        self._non_tensor_outputs[name] = (
            distribution.experimental_local_results(value))
      distribute_lib.get_replica_context().merge_call(
          merge_fn, args=(output,))


def _create_distributed_tensor_spec(strategy, tensor_spec):
  """Create a `tf.TypeSpec` for a given strategy and input `tensor_spec`.

  Args:
    strategy: The given `tf.distribute` strategy.
    tensor_spec: `tf.TensorSpec` of a given value. The batch dimension of the
      shape should be None if you have partial batches.

  Returns:
    A `tf.TypeSpec` that matches the values produced by a given strategy. This
    can be a `tf.TensorSpec` or a `PerRelicaSpec`.
  """
  num_replicas = len(strategy.extended.worker_devices)

  # For one device strategy that is not MultiWorkerMirroredStrategy,  return the
  # tensor_spec as is, since we don't wrap the output with PerReplica in this
  # case.
  # TODO(b/166464552): remove after we always wrap for all strategies.
  if not _always_wrap(strategy):
    return tensor_spec

  # For other cases we assume the input to tf.function is a per replica type.
  def _get_value_per_replica(tensor_spec_per_input):
    value_specs = [tensor_spec_per_input for _ in range(num_replicas)]
    return values.PerReplicaSpec(*value_specs)

  return nest.map_structure(_get_value_per_replica, tensor_spec)


def _replace_per_replica_spec(spec, i):
  """If `spec` is a `PerReplicaSpec`, then return its `i`th value_spec."""
  if isinstance(spec, values.PerReplicaSpec):
    return spec._value_specs[i]  # pylint: disable=protected-access
  else:
    return spec


def _cardinality(dataset):
  """Returns the cardinality of the dataset."""
  if context.executing_eagerly():
    with ops.device(dataset._variant_tensor.device):  # pylint: disable=protected-access
      return dataset.cardinality().numpy()
  return cardinality_lib.UNKNOWN


def _enable_get_next_as_optional(strategy, dataset, cardinality):
  """Returns whether to enable using partial batch handling."""
  # TODO(b/133073708): we currently need a flag to control the usage because
  # there is a performance difference between get_next() and
  # get_next_as_optional(). And we only enable get_next_as_optional when the
  # output shapes are not static.
  #
  # TODO(rxsang): We want to always enable the get_next_as_optional behavior
  # when user passed input_fn instead of dataset.
  if not getattr(
      strategy.extended, "enable_partial_batch_handling",
      getattr(strategy.extended, "experimental_enable_get_next_as_optional",
              False)):
    return False

  # If the dataset is infinite, we don't need to enable last partial batch
  # support. Note that we can only evaluate the cardinality of the dataset in
  # eager.
  if cardinality == cardinality_lib.INFINITE:
    return False

  return not _is_statically_shaped(
      dataset.element_spec) or strategy.extended._in_multi_worker_mode()  # pylint: disable=protected-access


def _create_per_replica(value_list, strategy):
  """Creates a PerReplica.

  For strategies other than OneDeviceStrategy, it creates a PerReplica whose
  type spec is set to the element spec of the dataset. This helps avoid
  retracing for partial batches. Retracing is problematic for multi client when
  different client retraces different time, since retracing changes the
  collective keys in the tf.function, and causes mismatches among clients.

  For single client strategies, this simply calls distribute_utils.regroup().

  Args:
    value_list: a list of values, one for each replica.
    strategy: the `tf.distribute.Strategy`.

  Returns:
    a structure of PerReplica.

  """
  # TODO(b/166464552): always wrap for all one device strategies as well.
  always_wrap = _always_wrap(strategy)
  per_replicas = distribute_utils.regroup(value_list, always_wrap=always_wrap)
  return per_replicas


def _always_wrap(strategy):
  """Returns whether to always wrap the values in a DistributedValues."""
  return strategy.extended._in_multi_worker_mode() or len(  # pylint: disable=protected-access
      strategy.extended.worker_devices) > 1


def _rebatch_as_dynamic(per_replica_spec):
  """Rebatch the spec to have a dynamic batch dimension."""
  assert isinstance(per_replica_spec, values.PerReplicaSpec), per_replica_spec

  # pylint: disable=protected-access
  def _rebatch(spec):
    # Rebatch if possible.
    try:
      return spec._unbatch()._batch(None)
    except ValueError:
      pass
    return spec

  return values.PerReplicaSpec(
      *nest.map_structure(_rebatch, per_replica_spec._value_specs))
  # pylint: enable=protected-access


def _ag_enumerate_not_implemented(s, unused_start):
  msg = (
      f"enumerate not supported with {s.__class__.__name__} types within "
      "tf.functions. Use a for loop over the dataset and keep a separate "
      "counter instead."
  )
  raise NotImplementedError(msg)


py_builtins.enumerate_registry.register(
    DistributedIterator, _ag_enumerate_not_implemented
)
py_builtins.enumerate_registry.register(
    DistributedDataset, _ag_enumerate_not_implemented
)