Input:
# 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.
# ==============================================================================
"""Python API for executing a tf.data.Dataset using a tf.data service."""

import enum
import functools
from typing import Callable

from tensorflow.core.protobuf import data_service_pb2
from tensorflow.python import tf2
from tensorflow.python.data.experimental.ops import compression_ops
from tensorflow.python.data.experimental.service import _pywrap_server_lib
from tensorflow.python.data.experimental.service import _pywrap_utils
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import options as options_lib
from tensorflow.python.data.ops import structured_function
from tensorflow.python.data.ops.options import AutoShardPolicy
from tensorflow.python.data.ops.options import ExternalStatePolicy
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import gen_experimental_dataset_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.saved_model import nested_structure_coder
from tensorflow.python.util.tf_export import tf_export

COMPRESSION_AUTO = "AUTO"
COMPRESSION_NONE = None
_PARALLEL_EPOCHS = "parallel_epochs"
_DISTRIBUTED_EPOCH = "distributed_epoch"

@tf_export("data.experimental.service.ShardingPolicy")
class ShardingPolicy(enum.IntEnum):
  """Specifies how to shard data among tf.data service workers.

  OFF: No sharding will be performed. Each worker produces the entire dataset
  without any sharding. With this mode, the best practice is to shuffle the
  dataset nondeterministically so that workers process the dataset in different
  orders. If workers are restarted or join the cluster mid-job, they will begin
  processing the dataset from the beginning.

  DYNAMIC: The input dataset is dynamically split among workers at runtime. Each
  worker gets the next split when it reads data from the dispatcher. Data is
  produced non-deterministically in this mode. Dynamic sharding works well with
  varying-sized tf.data service clusters, e.g., when you need to auto-scale your
  workers. Dynamic sharding provides at-most once visitation guarantees. No
  examples will be repeated, but some may be missed if a tf.data service worker
  gets restarted while processing a file.

  The following are static sharding policies. The semantics are similar to
  `tf.data.experimental.AutoShardPolicy`. These policies require:
  * The tf.data service cluster is configured with a fixed list of workers
    in DispatcherConfig.
  * Each client only reads from the local tf.data service worker.

  If a worker is restarted while performing static sharding, the worker will
  begin processing its shard again from the beginning.

  FILE: Shards by input files (i.e. each worker will get a fixed set of files to
  process). When this option is selected, make sure that there is at least as
  many files as workers. If there are fewer input files than workers, a runtime
  error will be raised.

  DATA: Shards by elements produced by the dataset. Each worker will process the
  whole dataset and discard the portion that is not for itself. Note that for
  this mode to correctly partition the dataset elements, the dataset needs to
  produce elements in a deterministic order.

  FILE_OR_DATA: Attempts FILE-based sharding, falling back to DATA-based
  sharding on failure.

  HINT: Looks for the presence of `shard(SHARD_HINT, ...)` which is treated as a
  placeholder to replace with `shard(num_workers, worker_index)`.
  """

  # LINT.IfChange(tf_data_service_sharding_policy)
  OFF = 0
  DYNAMIC = 1
  FILE = 2
  DATA = 3
  FILE_OR_DATA = 4
  HINT = 5
  # LINT.ThenChange()

  def _to_proto(self) -> data_service_pb2.ProcessingModeDef.ShardingPolicy:
    """Converts the policy to ProcessingModeDef proto enum."""

    if self == ShardingPolicy.OFF:
      return data_service_pb2.ProcessingModeDef.OFF
    if self == ShardingPolicy.DYNAMIC:
      return data_service_pb2.ProcessingModeDef.DYNAMIC
    if self == ShardingPolicy.FILE:
      return data_service_pb2.ProcessingModeDef.FILE
    if self == ShardingPolicy.DATA:
      return data_service_pb2.ProcessingModeDef.DATA
    if self == ShardingPolicy.FILE_OR_DATA:
      return data_service_pb2.ProcessingModeDef.FILE_OR_DATA
    if self == ShardingPolicy.HINT:
      return data_service_pb2.ProcessingModeDef.HINT
    raise ValueError(f"Unable to convert sharding policy {self!r} to proto.")

@tf_export("data.experimental.service.CrossTrainerCache")
class CrossTrainerCache:
  """Options related to the tf.data service cross trainer cache.

  This is used to enable cross-trainer cache when distributing a dataset. For
  example:

  ```
  dataset = dataset.apply(tf.data.experimental.service.distribute(
      processing_mode=tf.data.experimental.service.ShardingPolicy.OFF,
      service=FLAGS.tf_data_service_address,
      job_name="job",
      cross_trainer_cache=data_service_ops.CrossTrainerCache(
          trainer_id=trainer_id())))
  ```

  For more details, refer to
  https://www.tensorflow.org/api_docs/python/tf/data/experimental/service#sharing_tfdata_service_with_concurrent_trainers.
  """

  def __init__(self, trainer_id):
    """Constructs a CrossTrainerCache.

    Args:
      trainer_id: Each training job has a unique ID. Once a job has consumed
      data, the data remains in the cache and is re-used by jobs with different
      `trainer_id`s. Requests with the same `trainer_id` do not re-use data.

    Raises:
      ValueError if `trainer_id` is empty.
    """
    if not trainer_id:
      raise ValueError(
          "tf.data service cross-trainer cache requires a non-empty trainer ID."
      )
    self.trainer_id = trainer_id

  def _to_proto(self) -> data_service_pb2.CrossTrainerCacheOptions:
    return data_service_pb2.CrossTrainerCacheOptions(trainer_id=self.trainer_id)

def _get_validated_sharding_policy(processing_mode) -> ShardingPolicy:
  """Validates `processing_mode` and converts it to ShardingPolicy."""

  if isinstance(processing_mode, ShardingPolicy):
    return processing_mode
  if processing_mode == _PARALLEL_EPOCHS:
    return ShardingPolicy.OFF
  if processing_mode == _DISTRIBUTED_EPOCH:
    return ShardingPolicy.DYNAMIC

  raise ValueError("tf.data service processing mode should be a "
                   "`tf.data.experimental.service.ShardingPolicy`, "
                   "`\"parallel_epochs\"`, or `\"distributed_epoch\"`. Got "
                   f"{processing_mode!r}.")

def _validate_job_name(job_name) -> None:
  if job_name is None:
    return
  if not isinstance(job_name, str):
    raise ValueError("`job_name` must be a string, but `job_name` was of type "
                     f"{type(job_name)}. job_name={job_name}")
  if not job_name:
    raise ValueError("`job_name` must not be empty")

def _validate_compression(compression) -> None:
  valid_compressions = [COMPRESSION_AUTO, COMPRESSION_NONE]
  if compression not in valid_compressions:
    raise ValueError(f"Invalid `compression` argument: {compression}. "
                     f"Must be one of {valid_compressions}.")

def _get_compression_proto(
    compression) -> data_service_pb2.DataServiceMetadata.Compression:
  if compression == COMPRESSION_AUTO:
    return data_service_pb2.DataServiceMetadata.COMPRESSION_SNAPPY
  if compression == COMPRESSION_NONE:
    return data_service_pb2.DataServiceMetadata.COMPRESSION_OFF
  raise ValueError(f"Invalid `compression` argument: {compression}. "
                   f"Must be one of {[COMPRESSION_AUTO, COMPRESSION_NONE]}.")

def _to_tensor(dataset_id) -> tensor.Tensor:
  """Converts `dataset_id` to Tensor."""

  if isinstance(dataset_id, tensor.Tensor):
    return dataset_id
  if isinstance(dataset_id, str) or isinstance(dataset_id, bytes):
    return ops.convert_to_tensor(
        dataset_id, dtype=dtypes.string, name="dataset_id")
  return ops.convert_to_tensor(
      dataset_id, dtype=dtypes.int64, name="dataset_id")

def _to_string(dataset_id) -> str:
  """Converts `dataset_id` to string."""

  if isinstance(dataset_id, tensor.Tensor):
    return (dataset_id if dataset_id.dtype == dtypes.string else
            string_ops.as_string(dataset_id))
  return (dataset_id.decode()
          if isinstance(dataset_id, bytes) else str(dataset_id))

class _DataServiceDatasetV2(dataset_ops.DatasetSource):
  """A `Dataset` that reads elements from the tf.data service."""

  def __init__(self,
               dataset_id,
               processing_mode,
               address,
               element_spec,
               protocol,
               data_transfer_protocol,
               job_name=None,
               consumer_index=None,
               num_consumers=None,
               max_outstanding_requests=None,
               task_refresh_interval_hint_ms=None,
               cross_trainer_cache=None,
               target_workers="AUTO"):
    """Constructs a _DataServiceDatasetV2.

    Args:
      dataset_id: The dataset id for the dataset to read from.
      processing_mode: A `tf.data.experimental.service.ShardingPolicy`
        specifying how to shard the dataset among tf.data workers. See
        `tf.data.experimental.service.ShardingPolicy` for details. For backwards
        compatibility, `processing_mode` may also be set to the strings
        `"parallel_epochs"` or `"distributed_epoch"`, which are respectively
        equivalent to `ShardingPolicy.OFF` and `ShardingPolicy.DYNAMIC`.
      address: The tf.data service address, e.g. "localhost:5000".
      element_spec: The dataset element spec for the dataset to read from.
      protocol: The protocol to use for communicating with the tf.data service,
        e.g. "grpc".
      data_transfer_protocol: (Optional.) The protocol to use for transferring
        data with the tf.data service. By default, data is transferred using
        gRPC.
      job_name: (Optional.) The name of the job. If provided, it must be a
        non-empty string or Tensor. This argument makes it possible for multiple
        datasets to share the same job. The default behavior is that the dataset
        creates anonymous, exclusively owned jobs.
      consumer_index: (Optional.) The index of the consumer in the range from
        `0` to `num_consumers`. Must be specified alongside `num_consumers`.
        When specified, consumers will read from the job in a strict round-robin
        order, instead of the default first-come-first-served order.
      num_consumers: (Optional.) The number of consumers which will consume from
        the job. Must be specified alongside `consumer_index`. When specified,
        consumers will read from the job in a strict round-robin order, instead
        of the default first-come-first-served order. When `num_consumers` is
        specified, the dataset must have infinite cardinality to prevent a
        producer from running out of data early and causing consumers to go out
        of sync.
      max_outstanding_requests: (Optional.) A limit on how many elements may be
        requested at the same time. You can use this option to control the
        amount of memory used, since `distribute` won't use more than
        `element_size` * `max_outstanding_requests` of memory.
      task_refresh_interval_hint_ms: (Optional.) A hint for how often to query
        the dispatcher for task changes.
      cross_trainer_cache: (Optional.) If a `CrossTrainerCache` object is
        provided, dataset iteration will be shared across concurrently running
        trainers. See
        https://www.tensorflow.org/api_docs/python/tf/data/experimental/service#sharing_tfdata_service_with_concurrent_trainers
        for details.
      target_workers: (Optional.) Which workers to read from. If `"AUTO"`,
        tf.data runtime decides which workers to read from. If `"ANY"`, reads
        from any tf.data service workers. If `"LOCAL"`, only reads from local
        in-processs tf.data service workers. `"AUTO"` works well for most cases,
        while users can specify other targets. For example, `"LOCAL"` helps
        avoid RPCs and data copy if every TF worker colocates with a tf.data
        service worker. Consumers of a shared job must use the same
        `target_workers`. Defaults to `"AUTO"`.
    """
    if consumer_index is None != num_consumers is None:
      raise ValueError(
          "Must either set both `consumer_index` and `num_consumers`, "
          "or neither. ",
          f"consumer_index={consumer_index}, num_consumers={num_consumers}")
    if num_consumers is not None and job_name is None:
      raise ValueError("`job_name` must be set when setting `num_consumers`. "
                       f"num_consumers was set to {num_consumers}.")

    processing_mode_def = data_service_pb2.ProcessingModeDef(
        sharding_policy=_get_validated_sharding_policy(
            processing_mode)._to_proto())
    if job_name is None:
      job_name = ""
    if max_outstanding_requests is None:
      max_outstanding_requests = dataset_ops.AUTOTUNE
    if task_refresh_interval_hint_ms is None:
      task_refresh_interval_hint_ms = dataset_ops.AUTOTUNE

    self._dataset_id = _to_tensor(dataset_id)
    self._processing_mode = ops.convert_to_tensor(
        processing_mode_def.SerializeToString(),
        dtype=dtypes.string,
        name="processing_mode")
    self._address = ops.convert_to_tensor(
        address, dtype=dtypes.string, name="address")
    self._protocol = ops.convert_to_tensor(
        protocol, dtype=dtypes.string, name="protocol")
    self._job_name = ops.convert_to_tensor(
        job_name, dtype=dtypes.string, name="job_name")
    self._consumer_index = ops.convert_to_tensor(
        -1 if consumer_index is None else consumer_index,
        dtype=dtypes.int64,
        name="consumer_index")
    self._num_consumers = ops.convert_to_tensor(
        -1 if num_consumers is None else num_consumers,
        dtype=dtypes.int64,
        name="num_consumers")
    self._max_outstanding_requests = ops.convert_to_tensor(
        max_outstanding_requests,
        dtype=dtypes.int64,
        name="max_outstanding_requests")
    self._element_spec = element_spec
    uncompress_func = structured_function.StructuredFunctionWrapper(
        lambda x: compression_ops.uncompress(x, output_spec=element_spec),
        transformation_name="DataServiceDataset.uncompress()",
        input_structure=tensor.TensorSpec(shape=(), dtype=dtypes.variant))
    cross_trainer_cache_options = (
        cross_trainer_cache._to_proto().SerializeToString()
        if cross_trainer_cache else None)

    compat_kwargs = {}
    if data_transfer_protocol is not None:
      compat_kwargs["data_transfer_protocol"] = data_transfer_protocol

    # If `uncompress` is `True`, the dataset will query the servers to find
    # out the actual compression used. It is always set to `True` the first
    # time the graph is built, and set to false when serializing, so we will
    # uncompress at most once.
    uncompress = True
    variant_tensor = gen_experimental_dataset_ops.data_service_dataset_v4(
        dataset_id=self._dataset_id,
        processing_mode=self._processing_mode,
        address=self._address,
        protocol=self._protocol,
        job_name=self._job_name,
        consumer_index=self._consumer_index,
        num_consumers=self._num_consumers,
        max_outstanding_requests=self._max_outstanding_requests,
        task_refresh_interval_hint_ms=task_refresh_interval_hint_ms,
        iteration_counter=(
            gen_experimental_dataset_ops.dummy_iteration_counter()),
        target_workers=target_workers,
        uncompress=uncompress,
        uncompress_fn=uncompress_func.function,
        cross_trainer_cache_options=cross_trainer_cache_options,
        **compat_kwargs,
        **self._flat_structure)
    super(_DataServiceDatasetV2, self).__init__(variant_tensor)

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

class _DataServiceDatasetV1(dataset_ops.DatasetV1Adapter):
  """A `Dataset` that executes its input through the tf.data service."""

  @functools.wraps(_DataServiceDatasetV2.__init__)
  def __init__(self, dataset_id, processing_mode, address, element_spec,
               protocol, data_transfer_protocol, job_name, consumer_index,
               num_consumers, max_outstanding_requests,
               task_refresh_interval_hint_ms, cross_trainer_cache,
               target_workers):

    self._wrapped = _DataServiceDatasetV2(
        dataset_id=dataset_id,
        processing_mode=processing_mode,
        address=address,
        element_spec=element_spec,
        protocol=protocol,
        data_transfer_protocol=data_transfer_protocol,
        job_name=job_name,
        consumer_index=consumer_index,
        num_consumers=num_consumers,
        max_outstanding_requests=max_outstanding_requests,
        task_refresh_interval_hint_ms=task_refresh_interval_hint_ms,
        cross_trainer_cache=cross_trainer_cache,
        target_workers=target_workers)
    super(_DataServiceDatasetV1, self).__init__(self._wrapped)

if tf2.enabled():
  _DataServiceDataset = _DataServiceDatasetV2
else:
  _DataServiceDataset = _DataServiceDatasetV1

def _parse_service(service) -> tuple[str, str]:
  """Converts a tf.data service string into a (protocol, address) tuple.

  Args:
    service: A string in the format "protocol://address" or just "address". If
      the string is only an address, the default protocol will be used.

  Returns:
    The (protocol, address) tuple
  """
  if not isinstance(service, str):
    raise ValueError("`service` must be a string, but `service` was of type "
                     f"{type(service)}. service={service}")
  if not service:
    raise ValueError("`service` must not be empty")
  parts = service.split("://")
  if len(parts) == 2:
    protocol, address = parts
  elif len(parts) == 1:
    address = parts[0]
    protocol = _pywrap_utils.TF_DATA_DefaultProtocol()
  else:
    raise ValueError("Malformed `service` string has multiple '://': "
                     f"{service}.")
  # TODO(aaudibert): Considering validating reachability of address here.
  return (protocol, address)

def _distribute(
    processing_mode,
    service,
    job_name=None,
    consumer_index=None,
    num_consumers=None,
    max_outstanding_requests=None,
    task_refresh_interval_hint_ms=None,
    data_transfer_protocol=None,
    compression="AUTO",
    cross_trainer_cache=None,
    target_workers="AUTO",
) -> Callable[dataset_ops.Dataset, dataset_ops.Dataset]:
  """A transformation that moves dataset processing to the tf.data service.

  This transformation is similar to `distribute`, but supports additional
  parameters which we do not yet want to add to the public Python API.

  Args:
    processing_mode: A `tf.data.experimental.service.ShardingPolicy` specifying
      how to shard the dataset among tf.data workers. See
      `tf.data.experimental.service.ShardingPolicy` for details. For backwards
      compatibility, `processing_mode` may also be set to the strings
      `"parallel_epochs"` or `"distributed_epoch"`, which are respectively
      equivalent to `ShardingPolicy.OFF` and `ShardingPolicy.DYNAMIC`.
    service: A string or a tuple indicating how to connect to the tf.data
      service. If it's a string, it should be in the format
      `[<protocol>://]<address>`, where `<address>` identifies the dispatcher
        address and `<protocol>` can optionally be used to override the default
        protocol to use. If it's a tuple, it should be (protocol, address).
    job_name: (Optional.) The name of the job. If provided, it must be a
      non-empty string. This argument makes it possible for multiple datasets to
      share the same job. The default behavior is that the dataset creates
      anonymous, exclusively owned jobs.
    consumer_index: (Optional.) The index of the consumer in the range from `0`
      to `num_consumers`. Must be specified alongside `num_consumers`. When
      specified, consumers will read from the job in a strict round-robin order,
      instead of the default first-come-first-served order.
    num_consumers: (Optional.) The number of consumers which will consume from
      the job. Must be specified alongside `consumer_index`. When specified,
      consumers will read from the job in a strict round-robin order, instead of
      the default first-come-first-served order. When `num_consumers` is
      specified, the dataset must have infinite cardinality to prevent a
      producer from running out of data early and causing consumers to go out of
      sync.
    max_outstanding_requests: (Optional.) A limit on how many elements may be
      requested at the same time. You can use this option to control the amount
      of memory used, since `distribute` won't use more than `element_size` *
      `max_outstanding_requests` of memory.
    task_refresh_interval_hint_ms: (Optional.) A hint for how often to query the
      dispatcher for task changes.
    data_transfer_protocol: (Optional.) The protocol to use for transferring
      data with the tf.data service. By default, data is transferred using gRPC.
    compression: How to compress the dataset's elements before transferring them
      over the network. "AUTO" leaves the decision of how to compress up to the
      tf.data service runtime. `None` indicates not to compress.
    cross_trainer_cache: (Optional.) If a `CrossTrainerCache` object is
      provided, dataset iteration will be shared across concurrently running
      trainers. See
      https://www.tensorflow.org/api_docs/python/tf/data/experimental/service#sharing_tfdata_service_with_concurrent_trainers
      for details.
    target_workers: (Optional.) Which workers to read from. If `"AUTO"`, tf.data
      runtime decides which workers to read from. If `"ANY"`, reads from any
      tf.data service workers. If `"LOCAL"`, only reads from local in-processs
      tf.data service workers. `"AUTO"` works well for most cases, while users
      can specify other targets. For example, `"LOCAL"` helps avoid RPCs and
      data copy if every TF worker colocates with a tf.data service worker.
      Consumers of a shared job must use the same `target_workers`. Defaults to
      `"AUTO"`.

  Returns:
    Dataset: A `Dataset` of the elements produced by the data service.
  """
  processing_mode = _get_validated_sharding_policy(processing_mode)
  _validate_compression(compression)

  def _apply_fn(dataset) -> dataset_ops.Dataset:  # pylint: disable=missing-docstring
    dataset_id = _register_dataset(service, dataset, compression=compression)
    return _from_dataset_id(
        processing_mode,
        service,
        dataset_id,
        dataset.element_spec,
        job_name=job_name,
        consumer_index=consumer_index,
        num_consumers=num_consumers,
        max_outstanding_requests=max_outstanding_requests,
        task_refresh_interval_hint_ms=task_refresh_interval_hint_ms,
        data_transfer_protocol=data_transfer_protocol,
        cross_trainer_cache=cross_trainer_cache,
        target_workers=target_workers)

  return _apply_fn

@tf_export("data.experimental.service.distribute")
def distribute(
    processing_mode,
    service,
    job_name=None,
    consumer_index=None,
    num_consumers=None,
    max_outstanding_requests=None,
    data_transfer_protocol=None,
    compression="AUTO",
    cross_trainer_cache=None,
    target_workers="AUTO",
) -> Callable[dataset_ops.Dataset, dataset_ops.Dataset]:
  """A transformation that moves dataset processing to the tf.data service.

  When you iterate over a dataset containing the `distribute` transformation,
  the tf.data service creates a "job" which produces data for the dataset
  iteration.

  The tf.data service uses a cluster of workers to prepare data for training
  your model.
  The `processing_mode` argument to `tf.data.experimental.service.distribute`
  describes how to leverage multiple workers to process the input dataset.
  Currently, there are two processing modes to choose from: "distributed_epoch"
  and "parallel_epochs".

  "distributed_epoch" means that the dataset will be split across all tf.data
  service workers.
  The dispatcher produces "splits" for the dataset and sends them to workers for
  further processing. For example, if a dataset begins with a list of filenames,
  the dispatcher will iterate through the filenames and send the filenames to
  tf.data workers, which will perform the rest of the dataset transformations on
  those files. "distributed_epoch" is useful when your model needs to see each
  element of the dataset exactly once, or if it needs to see the data in a
  generally-sequential order. "distributed_epoch" only works for datasets with
  splittable sources, such as `Dataset.from_tensor_slices`,
  `Dataset.list_files`, or `Dataset.range`.

  "parallel_epochs" means that the entire input dataset will be processed
  independently by each of the tf.data service workers.
  For this reason, it is important to shuffle data (e.g. filenames)
  non-deterministically, so that each worker will process the elements of the
  dataset in a different order. "parallel_epochs" can be used to distribute
  datasets that aren't splittable.

  With two workers, "parallel_epochs" will produce every element of the dataset
  twice:

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> # Start two workers
  >>> workers = [
  ...     tf.data.experimental.service.WorkerServer(
  ...         tf.data.experimental.service.WorkerConfig(
  ...             dispatcher_address=dispatcher_address)) for _ in range(2)
  ... ]
  >>> dataset = tf.data.Dataset.range(10)
  >>> dataset = dataset.apply(tf.data.experimental.service.distribute(
  ...     processing_mode="parallel_epochs", service=dispatcher.target))
  >>> print(sorted(list(dataset.as_numpy_iterator())))
  [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9]

  "distributed_epoch", on the other hand, will still produce each element once:

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> workers = [
  ...     tf.data.experimental.service.WorkerServer(
  ...         tf.data.experimental.service.WorkerConfig(
  ...             dispatcher_address=dispatcher_address)) for _ in range(2)
  ... ]
  >>> dataset = tf.data.Dataset.range(10)
  >>> dataset = dataset.apply(tf.data.experimental.service.distribute(
  ...     processing_mode="distributed_epoch", service=dispatcher.target))
  >>> print(sorted(list(dataset.as_numpy_iterator())))
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  When using `apply(tf.data.experimental.service.distribute(...))`, the dataset
  before the `apply` transformation executes within the tf.data service, while
  the operations after `apply` happen within the local process.

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> workers = [
  ...     tf.data.experimental.service.WorkerServer(
  ...         tf.data.experimental.service.WorkerConfig(
  ...             dispatcher_address=dispatcher_address)) for _ in range(2)
  ... ]
  >>> dataset = tf.data.Dataset.range(5)
  >>> dataset = dataset.map(lambda x: x*x)
  >>> dataset = dataset.apply(
  ...    tf.data.experimental.service.distribute("parallel_epochs",
  ...                                            dispatcher.target))
  >>> dataset = dataset.map(lambda x: x+1)
  >>> print(sorted(list(dataset.as_numpy_iterator())))
  [1, 1, 2, 2, 5, 5, 10, 10, 17, 17]

  In the above example, the dataset operations (before applying the `distribute`
  function on the elements) will be executed on the tf.data workers,
  and the elements are provided over RPC. The remaining transformations
  (after the call to `distribute`) will be executed locally. The dispatcher
  and the workers will bind to usused free ports (which are chosen at random),
  in order to communicate with each other. However, to bind them to specific
  ports, the `port` parameter can be passed.

  The `job_name` argument allows jobs to be shared across multiple
  datasets. Instead of each dataset creating its own job, all
  datasets with the same `job_name` will consume from the same job. A new job
  will be created for each iteration of the dataset (with each repetition of
  `Dataset.repeat` counting as a new iteration). Suppose the `DispatchServer`
  is serving on `localhost:5000` and two training workers (in either a single
  client or multi-client setup) iterate over the below dataset, and there is a
  single tf.data worker:

  ```
  range5_dataset = tf.data.Dataset.range(5)
  dataset = range5_dataset.apply(tf.data.experimental.service.distribute(
      "parallel_epochs", "localhost:5000", job_name="my_job_name"))
  for iteration in range(3):
    print(list(dataset))
  ```

  The elements of each job will be split between the two processes, with
  elements being consumed by the processes on a first-come first-served basis.
  One possible result is that process 1 prints

  ```
  [0, 2, 4]
  [0, 1, 3]
  [1]
  ```

  and process 2 prints

  ```
  [1, 3]
  [2, 4]
  [0, 2, 3, 4]
  ```

  Job names must not be re-used across different training jobs within the
  lifetime of the tf.data service. In general, the tf.data service is expected
  to live for the duration of a single training job.
  To use the tf.data service with multiple training jobs, make sure to use
  different job names to avoid conflicts. For example, suppose a training job
  calls `distribute` with `job_name="job"` and reads until end of input. If
  another independent job connects to the same tf.data service and tries to read
  from `job_name="job"`, it will immediately receive end of input, without
  getting any data.

  **Coordinated data read**

  By default, when multiple consumers read from the same job, they receive data
  on a first-come first-served basis. In some use cases, it is advantageous to
  coordinate the consumers. At each step, consumers read data from the same
  worker.

  For example, the tf.data service can be used to coordinate example sizes
  across a cluster during synchronous training, so that during each step all
  replicas train on similar-sized elements. To achieve this, define a dataset
  which generates rounds of `num_consumers` consecutive similar-sized batches,
  then enable coordinated reads by setting `consumer_index` and `num_consumers`.

  NOTE: To keep consumers in sync, round robin data consumption requires that
  the dataset have infinite cardinality. You can get this by adding `.repeat()`
  at the end of the dataset definition.

  **Keras and Distribution Strategies**

  The dataset produced by the `distribute` transformation can be passed to
  Keras' `Model.fit` or Distribution Strategy's
  `tf.distribute.Strategy.experimental_distribute_dataset` like any other
  `tf.data.Dataset`. We recommend setting a `job_name` on the call to
  `distribute` so that if there are multiple workers, they read data from the
  same job. Note that the autosharding normally performed by
  `experimental_distribute_dataset` will be disabled when setting a `job_name`,
  since sharing the job already results in splitting data across the workers.
  When using a shared job, data will be dynamically balanced across workers, so
  that they reach end of input about the same time. This results in better
  worker utilization than with autosharding, where each worker processes an
  independent set of files, and some workers may run out of data earlier than
  others.

  Args:
    processing_mode: A `tf.data.experimental.service.ShardingPolicy` specifying
      how to shard the dataset among tf.data workers. See
      `tf.data.experimental.service.ShardingPolicy` for details. For backwards
      compatibility, `processing_mode` may also be set to the strings
      `"parallel_epochs"` or `"distributed_epoch"`, which are respectively
      equivalent to `ShardingPolicy.OFF` and `ShardingPolicy.DYNAMIC`.
    service: A string or a tuple indicating how to connect to the tf.data
      service. If it's a string, it should be in the format
      `[<protocol>://]<address>`, where `<address>` identifies the dispatcher
        address and `<protocol>` can optionally be used to override the default
        protocol to use. If it's a tuple, it should be (protocol, address).
    job_name: (Optional.) The name of the job. If provided, it must be a
      non-empty string. This argument makes it possible for multiple datasets to
      share the same job. The default behavior is that the dataset creates
      anonymous, exclusively owned jobs.
    consumer_index: (Optional.) The index of the consumer in the range from `0`
      to `num_consumers`. Must be specified alongside `num_consumers`. When
      specified, consumers will read from the job in a strict round-robin order,
      instead of the default first-come-first-served order.
    num_consumers: (Optional.) The number of consumers which will consume from
      the job. Must be specified alongside `consumer_index`. When specified,
      consumers will read from the job in a strict round-robin order, instead of
      the default first-come-first-served order. When `num_consumers` is
      specified, the dataset must have infinite cardinality to prevent a
      producer from running out of data early and causing consumers to go out of
      sync.
    max_outstanding_requests: (Optional.) A limit on how many elements may be
      requested at the same time. You can use this option to control the amount
      of memory used, since `distribute` won't use more than `element_size` *
      `max_outstanding_requests` of memory.
    data_transfer_protocol: (Optional.) The protocol to use for transferring
      data with the tf.data service. By default, data is transferred using gRPC.
    compression: How to compress the dataset's elements before transferring them
      over the network. "AUTO" leaves the decision of how to compress up to the
      tf.data service runtime. `None` indicates not to compress.
    cross_trainer_cache: (Optional.) If a `CrossTrainerCache` object is
      provided, dataset iteration will be shared across concurrently running
      trainers. See
      https://www.tensorflow.org/api_docs/python/tf/data/experimental/service#sharing_tfdata_service_with_concurrent_trainers
      for details.
    target_workers: (Optional.) Which workers to read from. If `"AUTO"`, tf.data
      runtime decides which workers to read from. If `"ANY"`, reads from any
      tf.data service workers. If `"LOCAL"`, only reads from local in-processs
      tf.data service workers. `"AUTO"` works well for most cases, while users
      can specify other targets. For example, `"LOCAL"` helps avoid RPCs and
      data copy if every TF worker colocates with a tf.data service worker.
      Consumers of a shared job must use the same `target_workers`. Defaults to
      `"AUTO"`.

  Returns:
    Dataset: A `Dataset` of the elements produced by the data service.
  """
  _validate_job_name(job_name)
  return _distribute(
      processing_mode=processing_mode,
      service=service,
      job_name=job_name,
      consumer_index=consumer_index,
      num_consumers=num_consumers,
      max_outstanding_requests=max_outstanding_requests,
      data_transfer_protocol=data_transfer_protocol,
      compression=compression,
      cross_trainer_cache=cross_trainer_cache,
      target_workers=target_workers)

def _register_dataset(
    service, dataset, compression, dataset_id=None) -> tensor.Tensor:
  """Registers a dataset with the tf.data service.

  This transformation is similar to `register_dataset`, but supports additional
  parameters which we do not yet want to add to the public Python API.

  Args:
    service: A string or a tuple indicating how to connect to the tf.data
      service. If it's a string, it should be in the format
      `[<protocol>://]<address>`, where `<address>` identifies the dispatcher
        address and `<protocol>` can optionally be used to override the default
        protocol to use. If it's a tuple, it should be (protocol, address).
    dataset: A `tf.data.Dataset` to register with the tf.data service.
    compression: How to compress the dataset's elements before transferring them
      over the network. "AUTO" leaves the decision of how to compress up to the
      tf.data service runtime. `None` indicates not to compress.
    dataset_id: (Optional.) By default, tf.data service generates a unique
      (string) ID for each registered dataset. If a `dataset_id` is provided, it
      will use the specified ID. If a dataset with a matching ID already exists,
      no new dataset is registered. This is useful if multiple training jobs
      want to (re)use the same dataset for training. In this case, they can
      register the dataset with the same dataset ID.

  Returns:
    A scalar string tensor representing the dataset ID.
  """
  _validate_compression(compression)

  if isinstance(service, tuple):
    protocol, address = service
  else:
    protocol, address = _parse_service(service)
  external_state_policy = dataset.options().experimental_external_state_policy
  if external_state_policy is None:
    external_state_policy = ExternalStatePolicy.WARN

  encoded_spec = None
  if context.executing_eagerly():
    encoded_spec = nested_structure_coder.encode_structure(
        dataset.element_spec).SerializeToString()

  if compression == COMPRESSION_AUTO:
    dataset = dataset.map(
        lambda *x: compression_ops.compress(x),
        num_parallel_calls=dataset_ops.AUTOTUNE)
  dataset = dataset._apply_debug_options()  # pylint: disable=protected-access

  metadata = data_service_pb2.DataServiceMetadata(
      element_spec=encoded_spec,
      compression=_get_compression_proto(compression))

  return gen_experimental_dataset_ops.register_dataset_v2(
      dataset._variant_tensor,  # pylint: disable=protected-access
      address=address,
      protocol=protocol,
      external_state_policy=external_state_policy.value,
      requested_dataset_id=dataset_id,
      metadata=metadata.SerializeToString())

@tf_export("data.experimental.service.register_dataset")
def register_dataset(
    service, dataset, compression="AUTO", dataset_id=None) -> tensor.Tensor:
  """Registers a dataset with the tf.data service.

  `register_dataset` registers a dataset with the tf.data service so that
  datasets can be created later with
  `tf.data.experimental.service.from_dataset_id`. This is useful when the
  dataset
  is registered by one process, then used in another process. When the same
  process is both registering and reading from the dataset, it is simpler to use
  `tf.data.experimental.service.distribute` instead.

  If the dataset is already registered with the tf.data service,
  `register_dataset` returns the already-registered dataset's id.

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> worker = tf.data.experimental.service.WorkerServer(
  ...     tf.data.experimental.service.WorkerConfig(
  ...         dispatcher_address=dispatcher_address))
  >>> dataset = tf.data.Dataset.range(10)
  >>> dataset_id = tf.data.experimental.service.register_dataset(
  ...     dispatcher.target, dataset)
  >>> dataset = tf.data.experimental.service.from_dataset_id(
  ...     processing_mode="parallel_epochs",
  ...     service=dispatcher.target,
  ...     dataset_id=dataset_id,
  ...     element_spec=dataset.element_spec)
  >>> print(list(dataset.as_numpy_iterator()))
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  Args:
    service: A string or a tuple indicating how to connect to the tf.data
      service. If it's a string, it should be in the format
      `[<protocol>://]<address>`, where `<address>` identifies the dispatcher
        address and `<protocol>` can optionally be used to override the default
        protocol to use. If it's a tuple, it should be (protocol, address).
    dataset: A `tf.data.Dataset` to register with the tf.data service.
    compression: (Optional.) How to compress the dataset's elements before
      transferring them over the network. "AUTO" leaves the decision of how to
      compress up to the tf.data service runtime. `None` indicates not to
      compress.
    dataset_id: (Optional.) By default, tf.data service generates a unique
      (string) ID for each registered dataset. If a `dataset_id` is provided, it
      will use the specified ID. If a dataset with a matching ID already exists,
      no new dataset is registered. This is useful if multiple training jobs
      want to (re)use the same dataset for training. In this case, they can
      register the dataset with the same dataset ID.

  Returns:
    A scalar string tensor representing the dataset ID.
  """
  return _register_dataset(service, dataset, compression, dataset_id)

def _from_dataset_id(processing_mode,
                     service,
                     dataset_id,
                     element_spec,
                     job_name=None,
                     consumer_index=None,
                     num_consumers=None,
                     max_outstanding_requests=None,
                     task_refresh_interval_hint_ms=None,
                     data_transfer_protocol=None,
                     cross_trainer_cache=None,
                     target_workers="AUTO") -> dataset_ops.Dataset:
  """Creates a dataset which reads data from the tf.data service.

  This transformation is similar to `from_dataset_id`, but supports additional
  parameters which we do not yet want to add to the public Python API.

  Args:
    processing_mode: A `tf.data.experimental.service.ShardingPolicy` specifying
      how to shard the dataset among tf.data workers. See
      `tf.data.experimental.service.ShardingPolicy` for details. For backwards
      compatibility, `processing_mode` may also be set to the strings
      `"parallel_epochs"` or `"distributed_epoch"`, which are respectively
      equivalent to `ShardingPolicy.OFF` and `ShardingPolicy.DYNAMIC`.
    service: A string or a tuple indicating how to connect to the tf.data
      service. If it's a string, it should be in the format
      `[<protocol>://]<address>`, where `<address>` identifies the dispatcher
        address and `<protocol>` can optionally be used to override the default
        protocol to use. If it's a tuple, it should be (protocol, address).
    dataset_id: The id of the dataset to read from. This id is returned by
      `register_dataset` when the dataset is registered with the tf.data
      service.
    element_spec: A nested structure of `tf.TypeSpec`s representing the type of
      elements produced by the dataset. This argument is only required inside a
      tf.function. Use `tf.data.Dataset.element_spec` to get the element spec
      for a given dataset.
    job_name: (Optional.) The name of the job. If provided, it must be a
      non-empty string or tensor. This argument makes it possible for multiple
      datasets to share the same job. The default behavior is that the dataset
      creates anonymous, exclusively owned jobs.
    consumer_index: (Optional.) The index of the consumer in the range from `0`
      to `num_consumers`. Must be specified alongside `num_consumers`. When
      specified, consumers will read from the job in a strict round-robin order,
      instead of the default first-come-first-served order.
    num_consumers: (Optional.) The number of consumers which will consume from
      the job. Must be specified alongside `consumer_index`. When specified,
      consumers will read from the job in a strict round-robin order, instead of
      the default first-come-first-served order. When `num_consumers` is
      specified, the dataset must have infinite cardinality to prevent a
      producer from running out of data early and causing consumers to go out of
      sync.
    max_outstanding_requests: (Optional.) A limit on how many elements may be
      requested at the same time. You can use this option to control the amount
      of memory used, since `distribute` won't use more than `element_size` *
      `max_outstanding_requests` of memory.
    task_refresh_interval_hint_ms: (Optional.) A hint for how often to query the
      dispatcher for task changes.
    data_transfer_protocol: (Optional.) The protocol to use for transferring
      data with the tf.data service. By default, data is transferred using gRPC.
    cross_trainer_cache: (Optional.) If a `CrossTrainerCache` object is
      provided, dataset iteration will be shared across concurrently running
      trainers. See
      https://www.tensorflow.org/api_docs/python/tf/data/experimental/service#sharing_tfdata_service_with_concurrent_trainers
      for details.
    target_workers: (Optional.) Which workers to read from. If `"AUTO"`, tf.data
      runtime decides which workers to read from. If `"ANY"`, reads from any
      tf.data service workers. If `"LOCAL"`, only reads from local in-processs
      tf.data service workers. `"AUTO"` works well for most cases, while users
      can specify other targets. For example, `"LOCAL"` helps avoid RPCs and
      data copy if every TF worker colocates with a tf.data service worker.
      Consumers of a shared job must use the same `target_workers`. Defaults to
      `"AUTO"`.

  Returns:
    A `tf.data.Dataset` which reads from the tf.data service.
  """
  def _get_element_spec():
    """Fetches the element spec from the server."""
    data_service_metadata = None
    dataset_id_val = tensor_util.constant_value(dataset_id)
    try:
      data_service_metadata = (
          _pywrap_server_lib.TF_DATA_GetDataServiceMetadataByID(
              dataset_id_val, address, protocol
          )
      )
    except NotImplementedError as err:
      raise ValueError(
          "The tf.data service is running an earlier version of TensorFlow "
          "that requires specifying `element_spec` as an argument to "
          "`from_dataset_id`. Please either supply an element spec or update "
          "the tf.data service to the latest version.") from err
    except RuntimeError:
      # This error results from dataset ID not found. A more appropriate error
      # will be raised when the dataset is created.
      pass

    if not data_service_metadata or not data_service_metadata.element_spec:
      dataset_id_val = tensor_util.constant_value(dataset_id)
      raise ValueError(
          f"Failed to fetch element spec for dataset id {dataset_id_val} from "
          "tf.data service. If the dataset was registered in graph mode or "
          "inside a tf.function, the `element_spec` must be specified as an "
          "argument to `from_dataset_id`.")

    struct_pb = nested_structure_coder.struct_pb2.StructuredValue()
    struct_pb.ParseFromString(data_service_metadata.element_spec)
    return nested_structure_coder.decode_proto(struct_pb)

  processing_mode = _get_validated_sharding_policy(processing_mode)
  if isinstance(service, tuple):
    protocol, address = service
  else:
    protocol, address = _parse_service(service)
  if job_name is not None:
    if not isinstance(job_name, str) and not isinstance(
        job_name, tensor.Tensor):
      raise ValueError(
          "`job_name` must be a string or Tensor, but `job_name` was of type "
          f"{type(job_name)}. job_name={job_name}.")

  if not element_spec:
    if not context.executing_eagerly():
      raise ValueError(
          "In graph mode `element_spec` must be provided manually.")
    element_spec = _get_element_spec()

  dataset = _DataServiceDataset(
      dataset_id=dataset_id,
      processing_mode=processing_mode,
      address=address,
      element_spec=element_spec,
      protocol=protocol,
      data_transfer_protocol=data_transfer_protocol,
      job_name=job_name,
      consumer_index=consumer_index,
      num_consumers=num_consumers,
      max_outstanding_requests=max_outstanding_requests,
      task_refresh_interval_hint_ms=task_refresh_interval_hint_ms,
      cross_trainer_cache=cross_trainer_cache,
      target_workers=target_workers)

  # Disable autosharding for shared jobs.
  if job_name is not None:
    options = options_lib.Options()
    options.experimental_distribute.auto_shard_policy = AutoShardPolicy.OFF
    dataset = dataset.with_options(options)
  return dataset

@tf_export("data.experimental.service.from_dataset_id")
def from_dataset_id(processing_mode,
                    service,
                    dataset_id,
                    element_spec=None,
                    job_name=None,
                    consumer_index=None,
                    num_consumers=None,
                    max_outstanding_requests=None,
                    data_transfer_protocol=None,
                    cross_trainer_cache=None,
                    target_workers="AUTO") -> dataset_ops.Dataset:
  """Creates a dataset which reads data from the tf.data service.

  This is useful when the dataset is registered by one process, then used in
  another process. When the same process is both registering and reading from
  the dataset, it is simpler to use `tf.data.experimental.service.distribute`
  instead.

  Before using `from_dataset_id`, the dataset must have been registered with the
  tf.data service using `tf.data.experimental.service.register_dataset`.
  `register_dataset` returns a dataset id for the registered dataset. That is
  the `dataset_id` which should be passed to `from_dataset_id`.

  The `element_spec` argument indicates the `tf.TypeSpec`s for the elements
  produced by the dataset. Currently `element_spec` must be explicitly
  specified, and match the dataset registered under `dataset_id`. `element_spec`
  defaults to `None` so that in the future we can support automatically
  discovering the `element_spec` by querying the tf.data service.

  `tf.data.experimental.service.distribute` is a convenience method which
  combines `register_dataset` and `from_dataset_id` into a dataset
  transformation.
  See the documentation for `tf.data.experimental.service.distribute` for more
  detail about how `from_dataset_id` works.

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> worker = tf.data.experimental.service.WorkerServer(
  ...     tf.data.experimental.service.WorkerConfig(
  ...         dispatcher_address=dispatcher_address))
  >>> dataset = tf.data.Dataset.range(10)
  >>> dataset_id = tf.data.experimental.service.register_dataset(
  ...     dispatcher.target, dataset)
  >>> dataset = tf.data.experimental.service.from_dataset_id(
  ...     processing_mode="parallel_epochs",
  ...     service=dispatcher.target,
  ...     dataset_id=dataset_id,
  ...     element_spec=dataset.element_spec)
  >>> print(list(dataset.as_numpy_iterator()))
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  Args:
    processing_mode: A `tf.data.experimental.service.ShardingPolicy` specifying
      how to shard the dataset among tf.data workers. See
      `tf.data.experimental.service.ShardingPolicy` for details. For backwards
      compatibility, `processing_mode` may also be set to the strings
      `"parallel_epochs"` or `"distributed_epoch"`, which are respectively
      equivalent to `ShardingPolicy.OFF` and `ShardingPolicy.DYNAMIC`.
    service: A string or a tuple indicating how to connect to the tf.data
      service. If it's a string, it should be in the format
      `[<protocol>://]<address>`, where `<address>` identifies the dispatcher
        address and `<protocol>` can optionally be used to override the default
        protocol to use. If it's a tuple, it should be (protocol, address).
    dataset_id: The id of the dataset to read from. This id is returned by
      `register_dataset` when the dataset is registered with the tf.data
      service.
    element_spec: A nested structure of `tf.TypeSpec`s representing the type of
      elements produced by the dataset. This argument is only required inside a
      tf.function. Use `tf.data.Dataset.element_spec` to get the element spec
      for a given dataset.
    job_name: (Optional.) The name of the job. If provided, it must be a
      non-empty string. This argument makes it possible for multiple datasets to
      share the same job. The default behavior is that the dataset creates
      anonymous, exclusively owned jobs.
    consumer_index: (Optional.) The index of the consumer in the range from `0`
      to `num_consumers`. Must be specified alongside `num_consumers`. When
      specified, consumers will read from the job in a strict round-robin order,
      instead of the default first-come-first-served order.
    num_consumers: (Optional.) The number of consumers which will consume from
      the job. Must be specified alongside `consumer_index`. When specified,
      consumers will read from the job in a strict round-robin order, instead of
      the default first-come-first-served order. When `num_consumers` is
      specified, the dataset must have infinite cardinality to prevent a
      producer from running out of data early and causing consumers to go out of
      sync.
    max_outstanding_requests: (Optional.) A limit on how many elements may be
      requested at the same time. You can use this option to control the amount
      of memory used, since `distribute` won't use more than `element_size` *
      `max_outstanding_requests` of memory.
    data_transfer_protocol: (Optional.) The protocol to use for transferring
      data with the tf.data service. By default, data is transferred using gRPC.
    cross_trainer_cache: (Optional.) If a `CrossTrainerCache` object is
      provided, dataset iteration will be shared across concurrently running
      trainers. See
      https://www.tensorflow.org/api_docs/python/tf/data/experimental/service#sharing_tfdata_service_with_concurrent_trainers
      for details.
    target_workers: (Optional.) Which workers to read from. If `"AUTO"`, tf.data
      runtime decides which workers to read from. If `"ANY"`, reads from any
      tf.data service workers. If `"LOCAL"`, only reads from local in-processs
      tf.data service workers. `"AUTO"` works well for most cases, while users
      can specify other targets. For example, `"LOCAL"` helps avoid RPCs and
      data copy if every TF worker colocates with a tf.data service worker.
      Consumers of a shared job must use the same `target_workers`. Defaults to
      `"AUTO"`.

  Returns:
    A `tf.data.Dataset` which reads from the tf.data service.
  """
  _validate_job_name(job_name)
  if job_name is not None:
    job_name = string_ops.string_join(
        ["dataset_id=", _to_string(dataset_id), job_name], "/")

  return _from_dataset_id(
      processing_mode=processing_mode,
      service=service,
      dataset_id=dataset_id,
      element_spec=element_spec,
      job_name=job_name,
      consumer_index=consumer_index,
      num_consumers=num_consumers,
      max_outstanding_requests=max_outstanding_requests,
      data_transfer_protocol=data_transfer_protocol,
      cross_trainer_cache=cross_trainer_cache,
      target_workers=target_workers)

Output: _validate_job_name

Output: _distribute

Output: _parse_service

Output: _from_dataset_id

Output: _register_dataset

Output: _get_validated_sharding_policy

Output: _to_string

Output: _to_tensor

Output: _validate_compression

Output: _get_compression_proto

Output: ShardingPolicy

Output: CrossTrainerCache


Input:
# Copyright 2019 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.
# ==============================================================================

import functools
import gc
import weakref

from absl.testing import parameterized
import numpy as np

from tensorflow.python import pywrap_tfe
from tensorflow.python.distribute import mirrored_strategy
from tensorflow.python.eager import backprop
from tensorflow.python.eager import def_function
from tensorflow.python.eager import forwardprop
from tensorflow.python.eager import forwardprop_util
from tensorflow.python.eager import record
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import custom_gradient
from tensorflow.python.ops import gradient_checker_v2
from tensorflow.python.ops import map_fn
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_impl
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops.parallel_for import control_flow_ops
from tensorflow.python.ops.unconnected_gradients import UnconnectedGradients
from tensorflow.python.platform import test
from tensorflow.python.util import nest

_X11_35_DERIVATIVES = [
    1.1**3.5, 3.5 * 1.1**2.5, 3.5 * 2.5 * 1.1**1.5, 3.5 * 2.5 * 1.5 * 1.1**0.5
]

# TODO(allenl): Move this somewhere useful once forward gradients are stable.
def _jvp(f, primals, tangents):
  """Compute the jacobian of `f` at `primals` multiplied by `tangents`."""
  with forwardprop.ForwardAccumulator(primals, tangents) as acc:
    primals_out = f(*primals)
  return primals_out, acc.jvp(
      primals_out, unconnected_gradients=UnconnectedGradients.ZERO)

def _jacfwd(f, primals):
  """Compute the jacobian of `f` at `primals` using forward-mode autodiff."""
  jac_flat = []
  flat_primals = nest.flatten(primals)
  tangent_mask = [
      array_ops.zeros_like(primal, dtype=primal.dtype)
      for primal in flat_primals
  ]
  for primal_index, primal in enumerate(flat_primals):
    primal_vector = array_ops.reshape(primal, [-1])
    primal_vector_length = array_ops.size(primal_vector)
    jac_columns = []
    for element_index in math_ops.range(primal_vector_length):
      mask = array_ops.one_hot(
          element_index, primal_vector_length, dtype=primal.dtype)
      tangent_mask[primal_index] = array_ops.reshape(mask,
                                                     array_ops.shape(primal))
      jac_columns.append(
          nest.map_structure(
              functools.partial(array_ops.reshape, shape=[-1]),
              _jvp(f, primals, nest.pack_sequence_as(primals,
                                                     tangent_mask))[1]))
    jac_flat.append(array_ops_stack.stack(jac_columns, axis=1))
    tangent_mask[primal_index] = array_ops.zeros_like(primal)
  return nest.pack_sequence_as(primals, jac_flat)

def _jvp_batch(f, primal, tangents):
  tf_function = def_function.function(f)

  return control_flow_ops.vectorized_map(
      functools.partial(_jvp, tf_function, primal), tangents)

def _jvp_batch_matmul(f, primals, tangent_batch):
  """Compute the jacobian of `f` at `primals` multiplied by `tangents`."""
  jac_fwd = _jacfwd(f, primals)

  def jac_mul(tangent):
    flat_tangent = array_ops.reshape(tangent, shape=[-1])
    tangent_vector = array_ops.expand_dims(flat_tangent, 1)
    jvp_vector = math_ops.matmul(jac_fwd, tangent_vector)
    return array_ops.reshape(jvp_vector, tangent.shape)

  return control_flow_ops.vectorized_map(jac_mul, tangent_batch)

def _grad(f, argnums=0):
  """Return a function which computes the gradient of `f`."""

  def _f(*params):
    with backprop.GradientTape() as tape:
      tape.watch(params)
      primals_out = f(*params)
    return tape.gradient(
        primals_out,
        params[argnums],
        unconnected_gradients=UnconnectedGradients.ZERO)

  return _f

def _gradfwd(f, argnums=0, f_out_dtypes=dtypes.float32):
  """Return a function which computes the gradient of `f` in forward mode."""

  def _f(*params):

    def _single_jvp(param_mask):
      with forwardprop.ForwardAccumulator(
          primals=[params[argnums]], tangents=param_mask) as acc:
        primals_out = f(*params)
      return acc.jvp(primals_out)

    # Building up a function to run with pfor takes a bit too long since we're
    # only running it a handful of times.
    return _vectorize_parameters(
        _single_jvp, [params[argnums]], use_pfor=False, dtype=f_out_dtypes)

  return _f

def _hvp(f, primals, tangents):
  """Compute a forward-over-back Hessian-vector product."""
  with forwardprop.ForwardAccumulator(primals, tangents) as acc:
    with backprop.GradientTape() as tape:
      tape.watch(primals)
      f_out = f(*primals)
      f_out.shape.assert_is_compatible_with([])
    return acc.jvp(tape.gradient(f_out, primals))

def _vectorize_parameters(f, params, use_pfor, dtype):
  """Loop over `params`, providing a one-hot mask to `f` for each."""
  parameter_sizes = [array_ops.size(param) for param in params]
  total_size = math_ops.add_n(parameter_sizes)

  def _wrapper(index):
    full_onehot = array_ops.one_hot(index, total_size)
    split_onehot = array_ops.split(full_onehot, parameter_sizes)
    tangents = [
        array_ops.reshape(v, array_ops.shape(param))
        for param, v in zip(params, split_onehot)
    ]
    return f(tangents)

  if use_pfor:
    return control_flow_ops.vectorized_map(_wrapper, math_ops.range(total_size))

  return map_fn.map_fn(_wrapper, math_ops.range(total_size), dtype)

def _forward_over_back_hessian(f, params, use_pfor, dtype=None):
  """Computes the full Hessian matrix for the scalar-valued f(*params).

  Args:
    f: A function taking `params` and returning a scalar.
    params: A possibly nested structure of tensors.
    use_pfor: If true, uses `tf.vectorized_map` calls instead of looping.
    dtype: Required if `use_pfor=False`. A possibly nested structure of dtypes
      (e.g. `tf.float32`) matching the structure of `f`'s returns.

  Returns:
    A possibly nested structure of matrix slices corresponding to `params`. Each
    slice has shape [P, p_s] where `p_s` is the number of parameters (`tf.size`)
    in the corresponding element of `params` and `P` is the total number of
    parameters (`sum_s(p_s)`). The full matrix can be obtained by concatenating
    along the second axis.
  """
  return _vectorize_parameters(
      functools.partial(_hvp, f, params),
      params,
      use_pfor=use_pfor,
      dtype=dtype)

def _test_gradients(testcase,
                    f,
                    primals,
                    order,
                    delta=1e-3,
                    rtol=1e-2,
                    atol=1e-6,
                    srtol=1e-6,
                    satol=1e-6):
  """Tests forward/backward jacobians of `f`'s [0, `order`)-order gradients."""
  if order < 1:
    raise ValueError(
        "`order` should be a positive integer, got '{}'.".format(order))
  if order > 1:
    _test_gradients(
        testcase=testcase,
        f=_grad(f),
        primals=primals,
        order=order - 1,
        delta=delta,
        rtol=rtol,
        atol=atol,
        srtol=srtol,
        satol=satol)
  sym_jac_back, num_jac = gradient_checker_v2.compute_gradient(
      f, primals, delta=delta)
  testcase.assertAllClose(num_jac, sym_jac_back, rtol=rtol, atol=atol)
  sym_jac_fwd = _jacfwd(f, primals)
  testcase.assertAllClose(num_jac, sym_jac_fwd, rtol=rtol, atol=atol)
  # And the symbolic computations should be much closer.
  testcase.assertAllClose(sym_jac_back, sym_jac_fwd, rtol=srtol, atol=satol)

@test_util.with_eager_op_as_function
class ForwardpropTest(test.TestCase, parameterized.TestCase):

  def testJVPFunction(self):
    add_outputs = (constant_op.constant(4.),)
    vp, = forwardprop._jvp_dispatch(
        op_name="Add",
        attr_tuple=(),
        inputs=(constant_op.constant(1.), constant_op.constant(3.)),
        outputs=add_outputs,
        tangents=(
            constant_op.constant(1.),
            constant_op.constant(5.),
        ))
    self.assertAllClose(1. + 5., self.evaluate(vp))

    mul_outputs = (constant_op.constant([20.]),)
    vp, = forwardprop._jvp_dispatch(
        op_name="Mul",
        attr_tuple=(),
        inputs=(constant_op.constant([4.]), constant_op.constant([5.])),
        outputs=mul_outputs,
        tangents=(
            constant_op.constant([2.]),
            constant_op.constant([3.]),
        ))
    self.assertAllClose([2. * 5. + 3. * 4.], self.evaluate(vp))

  def testJVPFunctionWithBatchOfTangents(self):
    add_outputs = (constant_op.constant(4.),)
    jvp_flat = forwardprop._jvp_dispatch(
        op_name="Add",
        attr_tuple=(),
        inputs=(constant_op.constant(1.), constant_op.constant(3.)),
        outputs=add_outputs,
        tangents=(
            constant_op.constant([1., 2., 3.]),
            constant_op.constant([4., 5., 6.]),
        ),
        use_batch=True)

    # Using evaluate and asserting with just a list works too
    # but the output is more explicit this way
    self.assertAllClose([constant_op.constant([1. + 4., 2. + 5., 3. + 6.])],
                        jvp_flat)

    mul_outputs = (constant_op.constant([20.]),)
    jvp_flat = forwardprop._jvp_dispatch(
        op_name="Mul",
        attr_tuple=(),
        inputs=(constant_op.constant([4.]), constant_op.constant([5.])),
        outputs=mul_outputs,
        tangents=(
            constant_op.constant([[1.], [0.], [1.]]),
            constant_op.constant([[0.], [1.], [1.]]),
        ),
        use_batch=True)
    self.assertAllClose([constant_op.constant([[5.], [4.], [5. + 4.]])],
                        jvp_flat)

  def testJVPFunctionRaisesError(self):
    sum_outputs = (constant_op.constant(6.),)

    with self.assertRaisesRegex(ValueError, r".*was expected to be of shape*"):
      forwardprop._jvp_dispatch(
          op_name="Add",
          attr_tuple=(),
          inputs=(constant_op.constant(2.), constant_op.constant(4.)),
          outputs=sum_outputs,
          tangents=(constant_op.constant([1., 2.]),
                    constant_op.constant([[1.], [2.]])),
          use_batch=True)

  def testNonDifferentiableOpWithInputTangent(self):
    x = constant_op.constant(1.)
    with forwardprop.ForwardAccumulator(x, 2.) as acc1:
      with forwardprop.ForwardAccumulator(x, 2.) as acc2:
        y = array_ops.zeros_like(x)
      self.assertIsNone(acc1.jvp(y))
    self.assertIsNone(acc2.jvp(y))

  def testRunFunctionsEagerly(self):
    try:
      original_setting = def_function.functions_run_eagerly()
      def_function.run_functions_eagerly(True)
      x = constant_op.constant(1.)
      with forwardprop.ForwardAccumulator(x, 2.) as acc:
        y = x * 3.
      self.assertAllClose(6., acc.jvp(y))
    finally:
      def_function.run_functions_eagerly(original_setting)

  def testJVPFunctionUsedByAccumulatorForOps(self):
    previous_fn = forwardprop._jvp_dispatch
    try:
      x = constant_op.constant(1.)
      with forwardprop.ForwardAccumulator(x, 2.) as acc:
        y = x + x
        pywrap_tfe.TFE_Py_RegisterJVPFunction(
            lambda *args, **kwargs: [constant_op.constant(-15.)])
        z = x + x
      self.assertAllClose(4., acc.jvp(y))
      self.assertAllClose(-15., acc.jvp(z))
    finally:
      pywrap_tfe.TFE_Py_RegisterJVPFunction(previous_fn)

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testFunctionCacheLimited(self):
    # Every time this loop is executed, it will create a slightly larger Tensor
    # and push it through Add's gradient.
    # We run TRACE_COUNT_LIMIT x 2 so that it is tested with both
    # experimental_relax_shapes on and off.
    for execution_count in range(forwardprop._TRACE_COUNT_LIMIT*2):
      x = array_ops.zeros([execution_count])
      with forwardprop.ForwardAccumulator(x, array_ops.ones_like(x)) as acc:
        y = x + x
      self.assertAllClose(2. * array_ops.ones_like(x), acc.jvp(y))

  def testVariableUnwatchedZero(self):
    v = variables.Variable([[1.]])
    x = constant_op.constant(1.)
    xt = constant_op.constant(2.)
    with forwardprop.ForwardAccumulator(x, xt) as acc:
      pass
    self.assertIsNone(acc.jvp(v))
    self.assertAllClose([[0.]], acc.jvp(v, unconnected_gradients="zero"))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testFunctionReturnsResource(self):
    v = variables.Variable([[1.]])
    x = constant_op.constant(1.)
    xt = constant_op.constant(2.)

    @def_function.function
    def f(a):
      return a, v.handle

    with forwardprop.ForwardAccumulator(x, xt) as acc:
      y, _ = f(x)
    self.assertAllClose(2., acc.jvp(y))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testMultipleWatchesAdd(self):
    x = constant_op.constant(-2.)
    with self.assertRaisesRegex(ValueError, "multiple times"):
      with forwardprop.ForwardAccumulator([x, x], [1., 2.]):
        pass
    with forwardprop.ForwardAccumulator([x], [3.]) as acc:
      self.assertAllClose(3., acc.jvp(x))
      acc._watch(x, constant_op.constant(10.))
      self.assertAllClose(13., acc.jvp(x))
      acc._watch(x, constant_op.constant(11.))
      self.assertAllClose(24., acc.jvp(x))
      y = constant_op.constant(3.) * x
    self.assertAllClose(24., acc.jvp(x))
    self.assertAllClose(24. * 3., acc.jvp(y))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testReenter(self):
    x = constant_op.constant(-2.)
    with forwardprop.ForwardAccumulator(x, 1.5) as acc:
      self.assertAllClose(1.5, acc.jvp(x))
      y = 4. * x
      self.assertAllClose(6., acc.jvp(y))
      with self.assertRaisesRegex(ValueError, "already recording"):
        with acc:
          pass
    z = 4. * x
    self.assertIsNone(acc.jvp(z))
    with acc:
      yy = y * y
    self.assertAllClose(6. * -8. * 2., acc.jvp(yy))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testDeadTensorsJVPCleared(self):
    x = array_ops.ones([100])
    x_weak = weakref.ref(x)
    grad_tensor = constant_op.constant(array_ops.zeros([100]))
    grad_tensor_weak = weakref.ref(grad_tensor)
    with forwardprop.ForwardAccumulator(x, grad_tensor) as acc:
      derived_tensor = constant_op.constant(2.) * x
      del grad_tensor
      self.assertAllClose(array_ops.zeros([100]), acc.jvp(x))
      del x
      self.assertIsNone(x_weak())
      self.assertIsNone(grad_tensor_weak())
      derived_tensor_weak = weakref.ref(derived_tensor)
      derived_tensor_grad = acc.jvp(derived_tensor)
      derived_tensor_grad_weak = weakref.ref(derived_tensor_grad)
      del derived_tensor
      del derived_tensor_grad
      self.assertIsNone(derived_tensor_weak())
      self.assertIsNone(derived_tensor_grad_weak())

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testJVPManual(self):
    primal, tangent = _jvp(math_ops.sin, (constant_op.constant(0.1),),
                           (constant_op.constant(0.2),))
    self.assertAllClose(math_ops.sin(0.1), primal)
    self.assertAllClose(math_ops.cos(0.1) * 0.2, tangent)

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testNumericHigherOrder(self):

    def f(x):
      pointwise = math_ops.sin(x) * math_ops.tan(x)
      return math_ops.reduce_prod(
          pointwise + math_ops.reduce_sum(pointwise), axis=1)

    _test_gradients(
        self,
        f,
        [constant_op.constant([[2.0, 3.0], [1.0, 4.0]])],
        order=3,
        srtol=1e-6,
        satol=1e-3,
    )

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testNumericHigherOrderFloat64(self):

    def f(x):
      pointwise = math_ops.sin(x) * math_ops.tan(x)
      return math_ops.reduce_prod(
          pointwise + math_ops.reduce_sum(pointwise), axis=1)

    _test_gradients(
        self,
        f,
        [constant_op.constant([[2.0, 3.0], [1.0, 4.0]], dtype=dtypes.float64)],
        order=3)

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testCustomGradient(self):

    @custom_gradient.custom_gradient
    def f(x):

      def grad(dy):
        return dy * math_ops.cos(x)

      return np.sin(x.numpy()), grad

    _test_gradients(self, f, [constant_op.constant([1., 2.])], order=3)

  # TODO(allenl): investigate why assert_no_new_pyobjects_executing_eagerly()
  # fails around this test?
  def testExceptionCustomGradientRecomputeGradForward(self):

    @custom_gradient.recompute_grad
    def f(x):
      return math_ops.reduce_prod(math_ops.tanh(x)**2)

    with self.assertRaisesRegex(NotImplementedError,
                                "recompute_grad tried to transpose"):
      primals = [constant_op.constant([1.])]
      sym_jac_fwd = _jacfwd(f, primals)

  def testExceptionInCustomGradientNotSwallowed(self):

    @custom_gradient.custom_gradient
    def f(unused_x):

      def grad(unused_dy):
        raise ValueError("test_error_string")

      return 1., grad

    c = constant_op.constant(1.)
    d = constant_op.constant(2.)
    with forwardprop.ForwardAccumulator(c, d):
      with self.assertRaisesRegex(ValueError, "test_error_string"):
        f(c)

  @parameterized.named_parameters([("EluM5", -0.5, nn_ops.elu),
                                   ("EluP5", [0.5], nn_ops.elu),
                                   ("SwishP5", 0.5, nn_impl.swish),
                                   ("SwishM5", [-0.5], nn_impl.swish)])
  def testElementwiseNNOps(self, value, op_fn):
    _test_gradients(self, op_fn, [constant_op.constant(value)], order=3)

  def testFusedBatchNormGradsInference(self):

    x_shape = [4, 10, 10, 2]
    increment = 3. / math_ops.reduce_prod(
        constant_op.constant(x_shape, dtype=dtypes.float32))
    x = array_ops.reshape(math_ops.range(-2., 1., increment), x_shape)
    scale = constant_op.constant([1., 1.1])
    offset = constant_op.constant([-0.5, -0.6])
    mean = constant_op.constant([-1.3, 1.4])
    variance = constant_op.constant([0.7, 0.9])
    epsilon = 0.001

    def _bn_fused(x_arg, scale_arg, offset_arg):
      return nn_impl.fused_batch_norm(
          x_arg,
          scale_arg,
          offset_arg,
          mean,
          variance,
          epsilon=epsilon,
          is_training=False)[0]

    _test_gradients(self, _bn_fused, [x, scale, offset], order=2, atol=1e-2)

  def testPushPopAccumulatorState(self):
    # Note that this example is somewhat contrived. push_forwardprop_state is
    # probably only useful in practice for building functions that compute jvps
    # alongside their usual outputs.
    c = constant_op.constant(1.)
    d = constant_op.constant(2.)
    with forwardprop.ForwardAccumulator(c, d) as acc:

      @custom_gradient.custom_gradient
      def f(x):
        y = math_ops.sin(x.numpy())

        def grad(dy):
          with forwardprop_util.push_forwardprop_state():
            x_copy = constant_op.constant(x.numpy())
            acc._watch(x_copy, dy)
            y_copy = math_ops.sin(x_copy)
          return dy * acc.jvp(y_copy)

        return y, grad

      output = f(c)
      self.assertAllClose(d * math_ops.cos(c), acc.jvp(output))

  @parameterized.named_parameters([
      ("Order{}".format(order), order, expected)
      for order, expected in enumerate(_X11_35_DERIVATIVES)
  ])
  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testHigherOrderPureForward(self, order, expected):

    def _forwardgrad(f):

      def _compute_forwardgrad(primal):
        tangent = constant_op.constant(1.)
        with forwardprop.ForwardAccumulator(primal, tangent) as acc:
          primal_out = f(primal)
        return acc.jvp(primal_out)

      return _compute_forwardgrad

    def _forward(x):
      return x**3.5

    f = _forward
    primal = constant_op.constant(1.1)
    for _ in range(order):
      f = _forwardgrad(f)
    self.assertAllClose(expected, f(primal))

  @parameterized.named_parameters([("Function", def_function.function),
                                   ("NoFunction", lambda f: f)])
  def testGradPureForward(self, decorator):

    @decorator
    def f(x):
      return x**3.5

    primal = constant_op.constant(1.1)
    with forwardprop.ForwardAccumulator(primal,
                                        constant_op.constant(1.)) as outer_acc:
      with forwardprop.ForwardAccumulator(primal,
                                          constant_op.constant(1.)) as acc:
        primal_out = f(primal)
    inner_jvp = acc.jvp(primal_out)
    outer_jvp = outer_acc.jvp(inner_jvp)
    self.assertAllClose(1.1**3.5, primal_out)
    self.assertAllClose(3.5 * 1.1**2.5, inner_jvp)
    self.assertAllClose(3.5 * 2.5 * 1.1**1.5, outer_jvp)
    self.assertIsNone(acc.jvp(outer_acc.jvp(primal_out)))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testJVPPacking(self):
    two = constant_op.constant(2.)
    primal_in = constant_op.constant(1.)
    inner_jvp = constant_op.constant(3.)
    with forwardprop.ForwardAccumulator(
        [primal_in, inner_jvp],
        [constant_op.constant(2.),
         constant_op.constant(4.)]) as outer_acc:
      with forwardprop.ForwardAccumulator(primal_in, inner_jvp) as inner_acc:
        packed_input_indices, packed_input_tangents = (
            forwardprop_util.pack_tangents([primal_in]))
        self.assertAllClose([3., 2., 4.], packed_input_tangents)
        expected_indices = (
            # inner_acc watches primal_in
            (
                (0, 1),),
            # outer_acc watches primal_in and inner_jvp
            ((0, 2), (1, 3)))
        self.assertAllEqual(expected_indices, packed_input_indices)
        primal_out = primal_in * two
        self.assertAllClose(6., inner_acc.jvp(primal_out))
        self.assertAllClose(4., outer_acc.jvp(primal_out))
        self.assertAllClose(8., outer_acc.jvp(inner_acc.jvp(primal_out)))
        packed_output_indices, packed_output_tangents = (
            forwardprop_util.pack_tangents([primal_out]))
        self.assertAllClose([6., 4., 8.], packed_output_tangents)
        self.assertAllEqual(expected_indices, packed_output_indices)

  def testFunctionGradInFunctionPureForward(self):

    @def_function.function
    def take_gradients():

      @def_function.function
      def f(x):
        return x**3.5

      primal = constant_op.constant(1.1)
      with forwardprop.ForwardAccumulator(
          primal, constant_op.constant(1.)) as outer_acc:
        with forwardprop.ForwardAccumulator(primal,
                                            constant_op.constant(1.)) as acc:
          primal_out = f(primal)
      inner_jvp = acc.jvp(primal_out)
      outer_jvp = outer_acc.jvp(inner_jvp)
      self.assertIsNone(acc.jvp(outer_acc.jvp(primal_out)))
      return primal_out, inner_jvp, outer_jvp

    primal_out, inner_jvp, outer_jvp = take_gradients()
    self.assertAllClose(1.1**3.5, primal_out)
    self.assertAllClose(3.5 * 1.1**2.5, inner_jvp)
    self.assertAllClose(3.5 * 2.5 * 1.1**1.5, outer_jvp)

  def testFunctionGrad(self):

    @def_function.function
    def f(x):
      return math_ops.reduce_prod(math_ops.tanh(x)**2)

    _test_gradients(self, f, [constant_op.constant([1., 2.])], order=3)

  def testReusingJVP(self):
    m1 = random_ops.random_uniform((256, 2096))
    m2 = array_ops.identity(m1)
    tangent1 = random_ops.random_uniform((256, 2096))
    tangent2 = random_ops.random_uniform((256, 2096))
    matmul = def_function.function(math_ops.matmul)

    with forwardprop.ForwardAccumulator(
        primals=[m1, m2], tangents=[tangent1, tangent2]) as acc:
      result1 = matmul(m1, m1, transpose_b=True)
      result2 = matmul(m2, m2, transpose_b=True)

    def _expected(mat, tangent):
      return (math_ops.matmul(tangent, mat, transpose_b=True) +
              math_ops.matmul(mat, tangent, transpose_b=True))

    self.assertAllClose(result1, result2)
    self.assertAllClose(_expected(m1, tangent1), acc.jvp(result1))
    self.assertAllClose(_expected(m2, tangent2), acc.jvp(result2))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testHVPMemory(self):

    def fun(x):
      return math_ops.reduce_prod(math_ops.tanh(x)**2)

    primals = constant_op.constant([1., 2., 3.])
    tangents = constant_op.constant([3., 4., 5.])
    _hvp(fun, (primals,), (tangents,))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testHVPCorrectness(self):

    def fun(x):
      return math_ops.reduce_prod(math_ops.tanh(x)**2)

    primals = constant_op.constant([1., 2., 3.])
    tangents = constant_op.constant([3., 4., 5.])
    forwardback_hvp_eager, = _hvp(fun, (primals,), (tangents,))
    forwardback_hvp_function, = def_function.function(_hvp)(fun, (primals,),
                                                            (tangents,))

    with backprop.GradientTape(persistent=True) as g:
      g.watch(primals)
      with backprop.GradientTape() as gg:
        gg.watch(primals)
        out = fun(primals)
      grad = array_ops_stack.unstack(gg.gradient(out, primals))
    hessian = []
    for i in range(3):
      hessian.append(g.gradient(grad[i], primals))
    hessian = array_ops_stack.stack(hessian, axis=0)
    backback_hvp = math_ops.tensordot(hessian, tangents, axes=1)

    self.assertAllClose(backback_hvp, forwardback_hvp_eager)
    self.assertAllClose(backback_hvp, forwardback_hvp_function)

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testShouldRecordAndStopRecord(self):
    c = constant_op.constant(1.)
    c_tangent = constant_op.constant(2.)
    with forwardprop.ForwardAccumulator(c, c_tangent) as acc:
      with backprop.GradientTape() as tape:
        self.assertFalse(record.should_record_backprop([c]))
        self.assertEqual(1, pywrap_tfe.TFE_Py_TapeSetPossibleGradientTypes([c]))
        tape.watch(c)
        self.assertEqual(2, pywrap_tfe.TFE_Py_TapeSetPossibleGradientTypes([c]))
        self.assertTrue(record.should_record_backprop([c]))
        with record.stop_recording():
          self.assertEqual(0,
                           pywrap_tfe.TFE_Py_TapeSetPossibleGradientTypes([c]))
          self.assertFalse(record.should_record_backprop([c]))
          d = c * 2.
        self.assertEqual(2, pywrap_tfe.TFE_Py_TapeSetPossibleGradientTypes([c]))
        self.assertTrue(record.should_record_backprop([c]))
        self.assertFalse(record.should_record_backprop([d]))
        self.assertIsNone(acc.jvp(d))
      self.assertIsNone(tape.gradient(d, c))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testRecordingSelectively(self):
    c = constant_op.constant(1.)
    c_tangent = constant_op.constant(2.)
    with forwardprop.ForwardAccumulator(c, c_tangent) as acc:
      with backprop.GradientTape(persistent=True) as tape:
        tape.watch(c)
        with record.stop_recording():
          two = constant_op.constant(2.)
          d = c * two
          three = constant_op.constant(3.)
          e = c * three
        self.assertIsNone(acc.jvp(d))
        self.assertIsNone(acc.jvp(e))
        self.assertIsNone(tape.gradient(d, c))
        self.assertIsNone(tape.gradient(e, c))
        record.record_operation_forwardprop_only(
            "CustomForwardMul", [d], [c, two], lambda dd: (two * dd, c * dd),
            None)
        record.record_operation_backprop_only("CustomBackwardMul", [e],
                                              [c, three], lambda de:
                                              (three * de, c * de))
        self.assertAllClose(4., acc.jvp(d))
        self.assertIsNone(acc.jvp(e))
        self.assertIsNone(tape.gradient(d, c))
        self.assertAllClose(3., tape.gradient(e, c))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testOpWithNoTrainableOutputs(self):
    v = variables.Variable(1.)
    with forwardprop.ForwardAccumulator(v, 11.):
      v.assign_sub(0.5)
      self.assertAllClose(0.5, self.evaluate(v))

  # TODO(b/141025187): Add a no_new_pyobjects decorator.
  def testVariableReadInFunction(self):
    v = variables.Variable(1.)
    with forwardprop.ForwardAccumulator(v, 11.) as acc:

      @def_function.function
      def f():
        return v.read_value(), 2. * v.read_value()

      result = f()
      self.assertAllClose((1.0, 2.), result)
      self.assertAllClose((11., 22.), acc.jvp(result))

  @parameterized.named_parameters([("ForwardPropFirst", True),
                                   ("TapeFirst", False)])
  def testForwardOverBackwardMemoryEfficiency(self, forward_prop_first):
    # Watching depends on nesting, not creation order
    c = constant_op.constant(1.)
    if forward_prop_first:
      forward_accumulator = forwardprop.ForwardAccumulator(c, .1)
      gradient_tape = backprop.GradientTape()
    else:
      gradient_tape = backprop.GradientTape()
      forward_accumulator = forwardprop.ForwardAccumulator(c, .1)
    try:
      gc.disable()
      with gradient_tape as tape:
        # Adding and removing the tape multiple times in different nesting
        # patterns does not affect watch ordering.
        pass
      with forward_accumulator as acc:
        with gradient_tape as tape:
          tape.watch(c)
          d = math_ops.cos(c)
          self.assertFalse(record.should_record_backprop((acc.jvp(d),)))
          e = math_ops.cos(acc.jvp(d))
          math_ops.cos(e)
          weak_e = weakref.ref(e)
          del e
          self.assertIsNone(weak_e())
        self.assertIsNone(tape.gradient(acc.jvp(d), c))
    finally:
      gc.enable()

  @parameterized.named_parameters([("ForwardPropFirst", True),
                                   ("TapeFirst", False)])
  def testBackwardOverForward(self, forward_prop_first):
    c = constant_op.constant(1.)
    # Watching depends on nesting, not creation order
    if forward_prop_first:
      forward_accumulator = forwardprop.ForwardAccumulator(c, .1)
      gradient_tape = backprop.GradientTape()
    else:
      gradient_tape = backprop.GradientTape()
      forward_accumulator = forwardprop.ForwardAccumulator(c, .1)
    with gradient_tape as tape:
      with forward_accumulator as acc:
        tape.watch(c)
        d = math_ops.cos(c)
        self.assertTrue(record.should_record_backprop((acc.jvp(d),)))
      self.assertAllClose(-.1 * math_ops.cos(1.), tape.gradient(acc.jvp(d), c))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testRecordingWithJVPIndices(self):
    c = constant_op.constant(1.)
    with forwardprop.ForwardAccumulator(c, 10.) as acc:
      packed_input_tangents = forwardprop_util.pack_tangents([c]).tangents
      self.assertAllClose([10.], packed_input_tangents)
      d = constant_op.constant(2.)
      d_tangent = constant_op.constant(3.)
      record.record_operation_forwardprop_only("FunctionWithInlineJVPs",
                                               [d] + [d_tangent],
                                               [c] + packed_input_tangents,
                                               None, (((0, 1),),))
      self.assertAllClose(3., acc.jvp(d))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testSpecialForwardFunctionUsed(self):
    c = constant_op.constant(1.)
    d = constant_op.constant(2.)
    e = constant_op.constant(3.)
    with forwardprop.ForwardAccumulator(c, 10.) as acc:
      record.record_operation("ForwardIsSpecial", [d], [c], None,
                              lambda jvp: [-2. * jvp])
      self.assertAllClose(-20., acc.jvp(d))
      record.record_operation("ForwardIsSpecial2", [], [], None, lambda: [])
      record.record_operation("ForwardIsSpecial3", [e], [d], None,
                              lambda x: [x])
      self.assertAllClose(-20., acc.jvp(e))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testVariableWatched(self):
    v = variables.Variable([1., 2., 3.])
    with forwardprop.ForwardAccumulator(v, constant_op.constant([.1, -.2,
                                                                 .3])) as acc:
      self.assertAllClose([.1, -.2, .3], acc.jvp(v))
      x = v * 2.
      self.assertAllClose([.2, -.4, .6], acc.jvp(x))
      x2 = v + .1
      self.assertAllClose([.1, -.2, .3], acc.jvp(x2))

  def testUnconnectedGradients(self):
    x = constant_op.constant(-1.)
    with forwardprop.ForwardAccumulator(x, 0.1) as acc:
      self.assertAllClose(0.1, acc.jvp(x, unconnected_gradients="zero"))
      self.assertAllClose(0.1, acc.jvp(x, unconnected_gradients="none"))
      y = constant_op.constant(-2.)
      self.assertAllClose(0.0, acc.jvp(y, unconnected_gradients="zero"))
      self.assertIsNone(acc.jvp(y, unconnected_gradients="none"))

  # TODO(kkb): One weakref instance is created with warmup_iters=2,
  # investigate.
  @test_util.assert_no_new_pyobjects_executing_eagerly(warmup_iters=3)
  def testVariableWatchedFunction(self):

    class _Model(module.Module):

      def __init__(self):
        self._v = None

      @def_function.function
      def compute_jvps(self):
        if self._v is None:
          self._v = variables.Variable([1., 2., 3.])
        with forwardprop.ForwardAccumulator(self._v,
                                            constant_op.constant([.1, -.2,
                                                                  .3])) as acc:
          x = self._v * 2.
          x2 = self._v + .1
        return acc.jvp((self._v, x, x2))

    model = _Model()
    v_jvp, x_jvp, x2_jvp = model.compute_jvps()
    self.assertAllClose([.1, -.2, .3], v_jvp)
    self.assertAllClose([.2, -.4, .6], x_jvp)
    self.assertAllClose([.1, -.2, .3], x2_jvp)

  def testIndexSlicesGrad(self):
    x = constant_op.constant([1.])

    with forwardprop.ForwardAccumulator(x, constant_op.constant([3.])) as acc:
      y = array_ops.gather(x, 0)
    self.assertAllClose(3., acc.jvp(y))

  def testIndexSlicesGradInFunction(self):

    @def_function.function
    def f(a):
      return array_ops.gather(a, 0)

    x = constant_op.constant([1.])

    with forwardprop.ForwardAccumulator(x, constant_op.constant([3.])) as acc:
      y = f(x)
    self.assertAllClose(3., acc.jvp(y))

  # NOTE: assert_no_new_pyobjects_executing_eagerly fails flakily on this
  # test... could be something wrong with the test decorator, or some sort of
  # nondeterministic caching.
  def testMirroredVariableWatched(self):

    def _replicated(input_tangent):
      with forwardprop.ForwardAccumulator(v, input_tangent) as acc:
        self.assertAllClose([.1, -.2, .3], acc.jvp(v))
        x = v * 2.
        self.assertAllClose([.2, -.4, .6], acc.jvp(x))
        x2 = v + .1
        self.assertAllClose([.1, -.2, .3], acc.jvp(x2))

    strategy = mirrored_strategy.MirroredStrategy()
    with strategy.scope():
      v = variables.Variable([1., 2., 3.])
      strategy.run(_replicated, args=(constant_op.constant([.1, -.2, .3]),))

  # TODO(b/141025187): Add a no_new_pyobjects decorator.
  def testArgumentUnused(self):
    v = constant_op.constant(1.)
    with forwardprop.ForwardAccumulator(v, 11.) as acc:

      @def_function.function
      def _f(x):
        del x
        return constant_op.constant(1.)

      result = _f(v)
      self.assertAllClose(1.0, result)
      self.assertIsNone(acc.jvp(result))

@def_function.function
def _has_loop(iters, y):
  ret = 0.
  for i in math_ops.range(iters):
    ret += y * math_ops.cast(i, dtypes.float32)
  return ret

@def_function.function
def _has_cond(k, y):
  if k > 1:
    ret = 3. * y
  else:
    ret = 0.
  return ret

@def_function.function
def _fprop_while(iters, y):
  with forwardprop.ForwardAccumulator(y, 1.) as acc:
    ret = 0.
    for i in math_ops.range(iters):
      ret += y * math_ops.cast(i, dtypes.float32)
  return acc.jvp(ret)

@def_function.function
def _fprop_cond(k, y):
  with forwardprop.ForwardAccumulator(y, 1.) as acc:
    if k > 1:
      ret = 3. * y
    else:
      ret = 0.
  return acc.jvp(ret)

class ControlFlowTests(test.TestCase):

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testOfFunctionWhile(self):
    y = constant_op.constant(1.)
    with forwardprop.ForwardAccumulator(y, 1.) as acc:
      self.assertAllClose(10., acc.jvp(_has_loop(constant_op.constant(5), y)))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testOfFunctionCond(self):
    y = constant_op.constant(1.)
    with forwardprop.ForwardAccumulator(y, 1.) as acc:
      self.assertAllClose(3., acc.jvp(_has_cond(constant_op.constant(5), y)))
      self.assertAllClose(0., acc.jvp(_has_cond(constant_op.constant(0), y)))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testInFunctionWhile(self):
    self.assertAllClose(
        10., _fprop_while(constant_op.constant(5), constant_op.constant(1.)))

  @test_util.assert_no_new_pyobjects_executing_eagerly()
  def testInFunctionCond(self):
    self.assertAllClose(
        3., _fprop_cond(constant_op.constant(5), constant_op.constant(1.)))
    self.assertAllClose(
        0., _fprop_cond(constant_op.constant(0), constant_op.constant(1.)))

class HessianTests(test.TestCase, parameterized.TestCase):

  def testHessian1D(self):
    # Note: stolen from ops/gradients_test.py
    m = 4
    rng = np.random.RandomState([1, 2, 3])
    mat_value = rng.randn(m, m).astype("float32")
    x_value = rng.randn(m).astype("float32")
    hess_value = mat_value + mat_value.T
    mat = variables.Variable(mat_value)

    def _f(x):
      return math_ops.reduce_sum(x[:, None] * mat * x[None, :])

    hessian_eager, = _forward_over_back_hessian(
        _f, [constant_op.constant(x_value)],
        use_pfor=False,
        dtype=[dtypes.float32])
    self.assertAllClose(hess_value, hessian_eager)
    hessian_function, = def_function.function(_forward_over_back_hessian)(
        _f, [constant_op.constant(x_value)],
        use_pfor=False,
        dtype=[dtypes.float32])
    self.assertAllClose(hess_value, hessian_function)
    hessian_pfor, = def_function.function(_forward_over_back_hessian)(
        _f, [constant_op.constant(x_value)],
        use_pfor=True,
        dtype=[dtypes.float32])
    self.assertAllClose(hess_value, hessian_pfor)

class BatchTests(test.TestCase, parameterized.TestCase):

  @parameterized.parameters([(math_ops.sin, (2, 3), 5),
                             (math_ops.sin, (2, 3, 4), 10)])
  def testJVPBatchCorrectness(self, f, primal_shape, batch_size):
    primals = [random_ops.random_uniform(primal_shape)]
    tangent_batch = [random_ops.random_uniform([batch_size, *primal_shape])]
    self.assertAllClose(
        _jvp_batch(f, primals, tangent_batch)[1],
        _jvp_batch_matmul(f, primals, *tangent_batch))

  def testBatchCorrectness(self):
    x = constant_op.constant(2.0)
    y = constant_op.constant(5.0)
    tangents = (
        constant_op.constant([1., 0., 1.]),
        constant_op.constant([0., 1., 1.]),
    )
    with forwardprop.ForwardAccumulator._batch_accumulator((x, y),
                                                           tangents) as acc:
      z = x * y
    self.assertAllClose(acc.jvp(z), constant_op.constant([5.0, 2.0, 7.0]))

  @parameterized.named_parameters([("ForwardPropFirst", True),
                                   ("TapeFirst", False)])
  def testBatchBackwardOverForward(self, forward_prop_first):
    x = constant_op.constant(1.)
    tangents = random_ops.random_normal(shape=[10], seed=1)
    expected = [-t * math_ops.cos(1.) for t in tangents]
    if forward_prop_first:
      batch_acc = forwardprop.ForwardAccumulator._batch_accumulator(x, tangents)
      gradient_tape = backprop.GradientTape(persistent=True)
    else:
      gradient_tape = backprop.GradientTape(persistent=True)
      batch_acc = forwardprop.ForwardAccumulator._batch_accumulator(x, tangents)
    with gradient_tape as tape:
      with batch_acc as acc:
        tape.watch(x)
        y = math_ops.cos(x)
        self.assertTrue(record.should_record_backprop((acc.jvp(y),)))
        jvps = acc.jvp(y)
      d2y_dx2 = [tape.gradient(dy_dx, x) for dy_dx in jvps]
    self.assertAllClose(expected, d2y_dx2)

if __name__ == "__main__":
  # TODO(allenl): Also test with 1.x-style graph mode.
  ops.enable_eager_execution()
  test.main()

Output: _grad

Output: _vectorize_parameters

Output: _hvp

Output: _forward_over_back_hessian

Output: _test_gradients

Output: _gradfwd

Output: _jvp_batch

Output: _jvp_batch_matmul

Output: _jacfwd

Output: _jvp

Output: ControlFlowTests

Output: ForwardpropTest

Output: BatchTests

Output: HessianTests

Output: _Model