# 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.
# ==============================================================================
"""Tests for custom training loops."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl.testing import parameterized

from tensorflow.python import tf2
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import device_util
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import reduce_util
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.eager import def_function
from tensorflow.python.eager import test
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import map_fn
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.util import nest


def get_dataset_from_tensor_slices(inp_array):
  dataset = dataset_ops.DatasetV2.from_tensor_slices(inp_array)
  # TODO(b/138326910): Remove Dataset V1 version once bug resolved.
  if not tf2.enabled():
    dataset = dataset_ops.Dataset.from_tensor_slices(inp_array)
  return dataset


class AssertFlattenedMixin(object):
  """Mixin for specialized asserts."""

  def assert_equal_flattened(self, expected_results, actual_results):
    """Asserts that flattened results are equal.

    Due to the number of replicas in the strategy, the output may have a
    different structure and needs to be flattened for comparison.

    Args:
      expected_results: The results expected as a result of a computation.
      actual_results: The actual results of a computation.
    """
    self.assertEqual(len(expected_results), len(actual_results))

    for i, expected_result in enumerate(expected_results):
      final_result = []
      actual_result = actual_results[i]
      for val in actual_result:
        final_result.extend(val.numpy())
      self.assertAllEqual(expected_result, final_result)


class InputIterationTest(test.TestCase, parameterized.TestCase,
                         AssertFlattenedMixin):

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testConstantNumpyInput(self, distribution):

    @def_function.function
    def run(x):

      def computation(x):
        return math_ops.square(x)

      outputs = distribution.experimental_local_results(
          distribution.run(computation, args=(x,)))
      return outputs

    self.assertAllEqual(
        constant_op.constant(4., shape=(distribution.num_replicas_in_sync)),
        run(2.))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testStatefulExperimentalRunAlwaysExecute(self, distribution):
    with distribution.scope():
      v = variables.Variable(
          0.0, aggregation=variables.VariableAggregation.MEAN)

    @def_function.function
    def train_step():

      def assign_add():
        v.assign_add(1.0)

      distribution.run(assign_add)
      return array_ops.zeros([])

    train_step()
    self.assertAllEqual(1.0, v.numpy())

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.strategies_minus_tpu,
          mode=["eager"]))
  def testFullEager(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)

    def train_step(data):
      return math_ops.square(data)

    dist_dataset = distribution.experimental_distribute_dataset(dataset)
    results = []
    for x in dist_dataset:
      output = distribution.experimental_local_results(
          distribution.run(train_step, args=(x,)))
      results.append(output)
    self.assert_equal_flattened([[25., 36.], [49., 64.]], results)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies, mode=["eager"]))
  def testGetNextAsOptional(self, distribution):
    data = [5., 6., 7., 8.]
    dataset = get_dataset_from_tensor_slices(data).batch(2)
    dist_dataset = distribution.experimental_distribute_dataset(dataset)
    iterator = iter(dist_dataset)

    def train_step(data):
      return math_ops.square(data)

    @def_function.function
    def run(iterator):
      return distribution.experimental_local_results(
          distribution.run(
              train_step, args=(iterator.get_next_as_optional().get_value(),)))

    self.assert_equal_flattened([[25., 36.]], [run(iterator)])

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies, mode=["eager"]))
  def testGetNextAsOptionalExampleUsage(self, distribution):
    global_batch_size = 2
    steps_per_loop = 6
    dataset = dataset_ops.Dataset.range(
        8, output_type=dtypes.int32).batch(global_batch_size)
    distributed_iterator = iter(
        distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def train_fn(distributed_iterator):

      def step_fn(x):
        return x

      for _ in math_ops.range(steps_per_loop):
        optional_data = distributed_iterator.get_next_as_optional()
        if not optional_data.has_value():
          break
        distribution.run(step_fn, args=(optional_data.get_value(),))

    train_fn(distributed_iterator)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.tpu_strategies, mode=["eager"]))
  def testFullEagerTPU(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)

    def train_step(data):
      return math_ops.square(data)

    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    with self.assertRaisesRegex(NotImplementedError,
                                "does not support pure eager execution"):
      distribution.run(train_step, args=(next(input_iterator),))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testStepInFunction(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)

    @def_function.function
    def train_step(data):
      return math_ops.square(data)

    dist_dataset = distribution.experimental_distribute_dataset(dataset)
    results = []
    for x in dist_dataset:
      output = distribution.experimental_local_results(
          distribution.run(train_step, args=(x,)))
      results.append(output)
    self.assert_equal_flattened([[25., 36.], [49., 64.]], results)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testRunInFunction(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)

    def train_step(data):
      return math_ops.square(data)

    @def_function.function
    def f_train_step(input_data):
      return distribution.experimental_local_results(
          distribution.run(train_step, args=(input_data,)))

    dist_dataset = distribution.experimental_distribute_dataset(dataset)
    results = []
    for x in dist_dataset:
      output = f_train_step(x)
      results.append(output)
    self.assert_equal_flattened([[25., 36.], [49., 64.]], results)

  @combinations.generate(
      combinations.combine(
          distribution=[
              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
              strategy_combinations.tpu_strategy,
              strategy_combinations.tpu_strategy_packed_var,
          ],
          mode=["eager"]))
  def testNestedOutput(self, distribution):
    dataset = get_dataset_from_tensor_slices([0, 1, 2, 3]).batch(2)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def run(iterator):

      def computation(x):
        return [{
            "a": x - 1,
            "b": x + 1
        }]

      inputs = next(iterator)
      outputs = distribution.run(computation, args=(inputs,))
      return nest.map_structure(distribution.experimental_local_results,
                                outputs)

    results = run(input_iterator)
    for replica in range(distribution.num_replicas_in_sync):
      # The input dataset is range(4), so the replica id is same as input.
      self.assertAllEqual(results[0]["a"][replica], [replica - 1])
      self.assertAllEqual(results[0]["b"][replica], [replica + 1])

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testRunInFunctionAutoGraphApplication(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)

    def train_step(data):
      return math_ops.square(data)

    @def_function.function
    def f_train_step(input_data):
      return distribution.experimental_local_results(
          distribution.run(train_step, args=(input_data,)))

    dist_dataset = distribution.experimental_distribute_dataset(dataset)
    results = []
    for x in dist_dataset:
      output = f_train_step(x)
      results.append(output)
    self.assert_equal_flattened([[25., 36.], [49., 64.]], results)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testDatasetIterationInFunction(self, distribution):
    with distribution.scope():
      a = variables.Variable(
          1.0, aggregation=variables.VariableAggregation.ONLY_FIRST_REPLICA)

    def train_step(_):
      a.assign_add(1.0)

    @def_function.function
    def f_train_step(dist_dataset):
      number_of_steps = constant_op.constant(0.0)
      product_of_means = constant_op.constant(2.0)
      for x in dist_dataset:  # loop with values modified each iteration
        number_of_steps += 1
        product_of_means *= math_ops.cast(
            distribution.reduce("MEAN", x, axis=0), product_of_means.dtype)

      for y in dist_dataset:  # loop with no intermediate state
        distribution.run(train_step, args=(y,))

      return number_of_steps, product_of_means

    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)
    dist_dataset = distribution.experimental_distribute_dataset(dataset)

    number_of_steps, product_of_means = f_train_step(dist_dataset)
    self.assertEqual(2, number_of_steps.numpy())
    self.assertNear((2 * (5+6)/2 * (7+8)/2), product_of_means.numpy(), 1e-3)

    # We set the initial value of `a` to 1 and iterate through the dataset 2
    # times(4/2 where 4 is the number of dataset elements and 2 is the batch
    # size). Hence the final result is 3.
    self.assertEqual(3.0, (a.numpy()))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testDatasetAssertWithDynamicBatch(self, distribution):
    # Regression test for github issue 33517.
    def step_fn(data):
      assert_op = control_flow_ops.Assert(math_ops.less_equal(
          math_ops.reduce_max(data), 100.), [data])
      with ops.control_dependencies([assert_op]):
        return math_ops.square(data)

    @def_function.function
    def train(dataset):
      results = []
      iterator = iter(dataset)
      # we iterate through the loop 5 times since we have 3 elements and a
      # global batch of 2.
      for _ in range(2):
        elem = next(iterator)
        output = distribution.experimental_local_results(
            distribution.run(step_fn, args=(elem,)))
        results.append(output)
      return results

    dataset = dataset_ops.DatasetV2.from_tensor_slices([5., 6., 7.,]).batch(2)
    # TODO(b/138326910): Remove Dataset V1 version once bug resolved.
    if not tf2.enabled():
      dataset = dataset_ops.Dataset.from_tensor_slices([5., 6., 7.,]).batch(2)
    dist_dataset = distribution.experimental_distribute_dataset(dataset)
    results = train(dist_dataset)

    expected_results = [[25., 36.], [49.]]
    self.assertEqual(len(expected_results), len(results))

    # Need to expand results since output will be grouped differently depending
    # on the number of replicas.
    for i, expected_result in enumerate(expected_results):
      final_result = []
      actual_result = results[i]
      for val in actual_result:
        final_result.extend(val.numpy())
      self.assertAllEqual(expected_result, final_result)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testDistributeDatasetIteratorWithoutFunction(self, distribution):
    data = [5., 6., 7., 8.]
    input_iterator = iter(
        distribution.distribute_datasets_from_function(
            lambda _: get_dataset_from_tensor_slices(data)))

    self.assertAllEqual(
        distribution.experimental_local_results(input_iterator.get_next()),
        data[0:distribution.num_replicas_in_sync])

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]
      ))
  def testDistributeDatasetIteratorWithFunction(self, distribution):
    data = [5., 6., 7., 8.]
    input_iterator = iter(
        distribution.distribute_datasets_from_function(
            lambda _: get_dataset_from_tensor_slices(data)))

    @def_function.function
    def run(iterator):
      return distribution.experimental_local_results(iterator.get_next())

    local_results = run(input_iterator)
    self.assertAllEqual(local_results,
                        data[0:distribution.num_replicas_in_sync])
    backing_devices = [result.backing_device for result in local_results]
    self.assertAllEqual(backing_devices, distribution.extended.worker_devices)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]
      ))
  def testDistributeDatasetPrefetch(self, distribution):
    data = [5., 6., 7., 8.]
    input_iterator = iter(
        distribution.experimental_distribute_dataset(
            get_dataset_from_tensor_slices(data).batch(2)))

    local_results = distribution.experimental_local_results(
        input_iterator.get_next())

    backing_devices = [result.backing_device for result in local_results]
    self.assertAllEqual(backing_devices, distribution.extended.worker_devices)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]
      ))
  def testDistributeDatasetFunctionPrefetch(self, distribution):
    data = [5., 6., 7., 8.]
    input_iterator = iter(
        distribution.distribute_datasets_from_function(
            lambda _: get_dataset_from_tensor_slices(data)))

    local_results = distribution.experimental_local_results(
        input_iterator.get_next())

    backing_devices = [result.backing_device for result in local_results]
    self.assertAllEqual(backing_devices, distribution.extended.worker_devices)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.tpu_strategies,
          mode=["eager"]
      ))
  def testDistributeDatasetHostPrefetch(self, distribution):
    data = [5., 6., 7., 8.]
    input_iterator = iter(
        distribution.experimental_distribute_dataset(
            get_dataset_from_tensor_slices(data).batch(2),
            distribute_lib.InputOptions(
                experimental_prefetch_to_device=False)))

    local_results = distribution.experimental_local_results(
        input_iterator.get_next())

    for result in local_results:
      self.assertEqual(result.backing_device,
                       device_util.resolve("/device:CPU:0"))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.tpu_strategies,
          mode=["eager"]
      ))
  def testDistributeDatasetFunctionHostPrefetch(self, distribution):
    data = [5., 6., 7., 8.]
    input_iterator = iter(
        distribution.distribute_datasets_from_function(
            lambda _: get_dataset_from_tensor_slices(data),
            distribute_lib.InputOptions(experimental_prefetch_to_device=False)))

    local_results = distribution.experimental_local_results(
        input_iterator.get_next())

    for result in local_results:
      self.assertEqual(result.backing_device,
                       device_util.resolve("/device:CPU:0"))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]
      ))
  def testDynamicShapes(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7.]).batch(4)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def run(iterator):
      def computation(x):
        return math_ops.reduce_mean(x)
      inputs = next(iterator)
      outputs = distribution.experimental_local_results(
          distribution.run(computation, args=(inputs,)))
      return outputs

    # This assumes that there are exactly 2 replicas
    self.assertAllEqual([5.5, 7.], run(input_iterator))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]))
  def testDynamicShapesWithRunOptions(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7.]).batch(4)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))
    options = distribute_lib.RunOptions
    options.experimental_bucketizing_dynamic_shape = True

    @def_function.function
    def run(iterator):

      def computation(x):
        return math_ops.reduce_mean(x)

      inputs = next(iterator)
      outputs = distribution.experimental_local_results(
          distribution.run(
              computation, args=(inputs,), options=options))
      return outputs

    # This assumes that there are exactly 2 replicas
    self.assertAllEqual([5.5, 7.], run(input_iterator))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]))
  def testDynamicOutputsWithX64(self, distribution):
    dataset = get_dataset_from_tensor_slices(
        [5]).map(lambda x: math_ops.cast(x, dtypes.int64)).batch(2)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def run(iterator):

      def computation(x):
        return math_ops.add(x, x)

      inputs = next(iterator)
      outputs = distribution.experimental_local_results(
          distribution.run(computation, args=(inputs,)))
      return outputs

    # This assumes that there are exactly 2 replicas
    result = run(input_iterator)
    self.assertAllEqual([10], result[0])
    self.assertAllEqual([], result[1])

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]
      ))
  def testDynamicShapesWithGetNextOutsideFunction(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7.]).batch(4)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def run(inputs):
      def computation(x):
        return math_ops.reduce_mean(x)
      outputs = distribution.experimental_local_results(
          distribution.run(computation, args=(inputs,)))
      return outputs

    # This assumes that there are exactly 2 replicas
    self.assertAllEqual([5.5, 7.], run(next(input_iterator)))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]
      ))
  def testStrategyReduceWithDynamicShapes(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7.]).batch(4)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def run(iterator):
      inputs = next(iterator)
      return distribution.reduce(reduce_util.ReduceOp.MEAN, inputs, axis=0)

    self.assertAllEqual(6., run(input_iterator))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]
      ))
  def testStrategyReduceWithDynamicShapesRank2(self, distribution):
    dataset = get_dataset_from_tensor_slices(
        [[1., 1.], [1., 1.], [1., 1.]]).batch(4)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def run(iterator):
      inputs = next(iterator)
      return distribution.reduce(reduce_util.ReduceOp.MEAN, inputs, axis=0)

    self.assertAllEqual([1., 1.], run(input_iterator))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]
      ))
  def testDynamicShapesWithSizeOp(self, distribution):
    dataset = get_dataset_from_tensor_slices([5., 6., 7.]).batch(4)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def run(inputs):
      def computation(x):
        return array_ops.size_v2(x)
      outputs = distribution.experimental_local_results(
          distribution.run(computation, args=(inputs,)))
      return outputs

    # This assumes that there are exactly 2 replicas
    self.assertAllEqual([2, 1], run(next(input_iterator)))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]))
  def testSegmentSumWithDynamicNumberOfSegments(self, distribution):

    def dataset_fn(_):
      data = array_ops.zeros(5, dtype=dtypes.int32)
      dataset = get_dataset_from_tensor_slices(data)
      dataset = dataset.batch(3)
      return dataset

    input_iterator = iter(
        distribution.distribute_datasets_from_function(dataset_fn))

    @def_function.function
    def step_fn(example):
      segment_ids = array_ops.zeros_like_v2(example)
      num_segment = array_ops.shape(example)[0]
      # If number of segments is dynamic, output should be a dynamic shape.
      return math_ops.unsorted_segment_sum(example, segment_ids, num_segment)

    # This assumes that there are exactly 2 replicas
    outputs = distribution.experimental_local_results(
        distribution.run(step_fn, args=(next(input_iterator),)))
    self.assertAllEqual((3,), outputs[0].shape)
    self.assertAllEqual((2,), outputs[1].shape)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]))
  def testReshapeWithDynamicInputs(self, distribution):

    def dataset_fn(_):
      data = array_ops.zeros((5, 1, 2), dtype=dtypes.int32)
      dataset = get_dataset_from_tensor_slices(data)
      dataset = dataset.batch(3)
      return dataset

    input_iterator = iter(
        distribution.distribute_datasets_from_function(dataset_fn))

    @def_function.function
    def step_fn(example):
      # example: [<=3, 1, 2]
      # tile: [<=3, <=3, 2]
      tile = array_ops.tile(example, [1, array_ops.shape(example)[0], 1])
      # reshape1: [<=(3*3 = 9), 2]
      reshape1 = array_ops.reshape(tile, [-1, 2])

      # reshape2: [<=3, <=3, 2]
      reshape2 = array_ops.reshape(
          reshape1,
          [array_ops.shape(example)[0],
           array_ops.shape(example)[0], 2])

      # reshape3: [<=3, -1, 2]
      reshape3 = array_ops.reshape(reshape1,
                                   [array_ops.shape(example)[0], -1, 2])
      # reshape4: [-1, <=3, 2]
      reshape4 = array_ops.reshape(reshape1,
                                   [-1, array_ops.shape(example)[0], 2])
      return [reshape1, reshape2, reshape3, reshape4]

    # This assumes that there are exactly 2 replicas
    outputs = distribution.experimental_local_results(
        distribution.run(step_fn, args=(next(input_iterator),)))
    self.assertAllEqual((9, 2), outputs[0][0].values[0].shape)
    self.assertAllEqual((3, 3, 2), outputs[0][1].values[0].shape)
    self.assertAllEqual((3, 3, 2), outputs[0][2].values[0].shape)
    self.assertAllEqual((3, 3, 2), outputs[0][3].values[0].shape)

    self.assertAllEqual((4, 2), outputs[0][0].values[1].shape)
    self.assertAllEqual((2, 2, 2), outputs[0][1].values[1].shape)
    self.assertAllEqual((2, 2, 2), outputs[0][2].values[1].shape)
    self.assertAllEqual((2, 2, 2), outputs[0][3].values[1].shape)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]))
  def testDynamicShapesWithFirstReplicaNotMaximumShape(self, distribution):
    def dataset_fn(_):
      dataset1 = get_dataset_from_tensor_slices([[1., 2.], [1., 2.]])
      dataset2 = get_dataset_from_tensor_slices([[1., 2., 3.],
                                                 [1., 2., 3.]])
      dataset = dataset1.concatenate(dataset2)
      dataset = dataset.batch(2, drop_remainder=True)
      return dataset

    input_iterator = iter(
        distribution.distribute_datasets_from_function(dataset_fn))

    @def_function.function
    def run(inputs):
      def computation(x):
        return math_ops.reduce_mean(x)
      outputs = distribution.experimental_local_results(
          distribution.run(computation, args=(inputs,)))
      return outputs

    # This assumes that there are exactly 2 replicas
    self.assertAllEqual([1.5, 2.], run(next(input_iterator)))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.multidevice_strategies,
          mode=["eager"]))
  def testMapFnWithDynamicInputs(self, distribution):

    def dataset_fn(_):
      data = array_ops.zeros((20, 300, 32), dtype=dtypes.int32)
      dataset = get_dataset_from_tensor_slices(data)
      dataset = dataset.batch(16)
      return dataset

    input_iterator = iter(
        distribution.distribute_datasets_from_function(dataset_fn))

    def embedding_lookup(inputs):
      embedding_weights = array_ops.zeros((1, 128))
      flat_inputs = array_ops.reshape(inputs, [-1])
      embeddings = array_ops.gather(embedding_weights, flat_inputs)
      embeddings = array_ops.reshape(embeddings, inputs.shape.as_list() + [128])
      return embeddings

    @def_function.function
    def step_fn(example):
      return map_fn.map_fn(
          embedding_lookup, example, fn_output_signature=dtypes.float32)

    # This assumes that there are exactly 2 replicas
    outputs = distribution.experimental_local_results(
        distribution.run(step_fn, args=(next(input_iterator),)))
    self.assertAllEqual((16, 300, 32, 128), outputs[0].shape)
    self.assertAllEqual((4, 300, 32, 128), outputs[1].shape)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testDatasetDistributeEvenlyDivisibleDrop(self, distribution):
    # If the batch size is evenly divisible by the number of workers and we set
    # drop_remainder=True on the dataset, then DistributedIterator will use a
    # different (and more efficient) code path which avoids some control flow
    # ops.
    dataset = get_dataset_from_tensor_slices([5., 6.]).batch(
        2, drop_remainder=True)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    data = next(input_iterator)

    expected_result = [5., 6.]
    final_result = []
    actual_result = distribution.experimental_local_results(data)
    for val in actual_result:
      final_result.extend(val)
    self.assertAllEqual(expected_result, final_result)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testDatasetDistributeNotDivisibleDrop(self, distribution):
    # If each batch is not evenly divisible by the number of workers,
    # the remainder will be dropped.
    dataset = get_dataset_from_tensor_slices([5., 6.]).batch(
        1, drop_remainder=True)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    data = next(input_iterator)

    expected_result = [5.]
    final_result = []
    actual_result = distribution.experimental_local_results(data)
    for val in actual_result:
      final_result.extend(val)
    self.assertAllEqual(expected_result, final_result)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testDatasetDistributeEvenlyDivisibleNoDrop(self, distribution):
    # Setting drop_remainder=False on the dataset causes DistributedIterator
    # to use get_next_as_optional(), even if the batched dataset is evenly
    # divisible by the number of workers.
    dataset = get_dataset_from_tensor_slices([5., 6.]).batch(
        2, drop_remainder=False)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    data = next(input_iterator)

    expected_result = [5., 6.]
    final_result = []
    actual_result = distribution.experimental_local_results(data)
    for val in actual_result:
      final_result.extend(val)
    self.assertAllEqual(expected_result, final_result)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testDatasetPartialBatchWithMixedOutputs(self, distribution):
    # Dynamic output size with a mix of static and dynamic outputs
    dataset = get_dataset_from_tensor_slices([5.]).batch(2)
    input_iterator = iter(distribution.experimental_distribute_dataset(dataset))

    @def_function.function
    def run(iterator):

      def computation(x):
        # Fixed size output with a dynamic sized output.
        return array_ops.zeros([3]), math_ops.square(x)

      return distribution.run(
          computation, args=(next(iterator),))

    results = run(input_iterator)

    # First result is fixed for all replicas.
    for replica_id in range(distribution.num_replicas_in_sync):
      self.assertAllEqual([0., 0., 0.],
                          distribution.experimental_local_results(
                              results[0])[replica_id])
    # Only first replica has distributed dataset computation.
    self.assertAllEqual([25.],
                        distribution.experimental_local_results(results[1])[0])
    # Other replicas have no distributed dataset computation.
    for replica_id in range(1, distribution.num_replicas_in_sync):
      self.assertAllEqual([],
                          distribution.experimental_local_results(
                              results[1])[replica_id])

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testIterationInsideFunction(self, distribution):

    def step_fn(data):
      return math_ops.square(data)

    @def_function.function
    def train(dataset):
      results = []
      iterator = iter(dataset)
      # we iterate through the loop 2 times since we have 4 elements and a
      # global batch of 2.
      for _ in range(2):
        elem = next(iterator)
        output = distribution.experimental_local_results(
            distribution.run(step_fn, args=(elem,)))
        results.append(output)
      return results

    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)
    dist_dataset = distribution.experimental_distribute_dataset(dataset)
    results = train(dist_dataset)
    self.assert_equal_flattened([[25., 36.], [49., 64.]], results)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testIterationOutsideFunction(self, distribution):

    def train_step(data):
      return math_ops.square(data)

    @def_function.function
    def f_train_step(input_data):
      return distribution.experimental_local_results(
          distribution.run(train_step, args=(input_data,)))

    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)
    dist_dataset = distribution.experimental_distribute_dataset(dataset)
    iterator = iter(dist_dataset)
    results = []
    # we iterate through the loop 2 times since we have 4 elements and a
    # global batch of 2.
    for _ in range(2):
      output = f_train_step(next(iterator))
      results.append(output)
    self.assert_equal_flattened([[25., 36.], [49., 64.]], results)

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testMultiDeviceDataCapturedFunction(self, distribution):
    inputs = constant_op.constant([2., 3.])
    dataset = lambda _: dataset_ops.Dataset.from_tensor_slices(inputs).repeat(5)
    input_iterator = iter(
        distribution.distribute_datasets_from_function(dataset))
    with distribution.scope():
      var = variables.Variable(1.0)

    @def_function.function
    def train_step(input_iterator):

      def func(inputs):
        return math_ops.square(inputs) + var

      per_replica_outputs = distribution.run(
          func, (next(input_iterator),))
      mean = distribution.reduce(
          reduce_util.ReduceOp.MEAN, per_replica_outputs, axis=None)
      for _ in dataset_ops.Dataset.range(1):
        per_replica_outputs = distribution.run(
            func, (next(input_iterator),))
        mean = distribution.reduce(
            reduce_util.ReduceOp.MEAN, per_replica_outputs, axis=None)
      return mean

    with distribution.scope():
      if distribution.num_replicas_in_sync == 1:
        self.assertAlmostEqual(10.0, self.evaluate(train_step(input_iterator)))
      else:
        self.assertAlmostEqual(7.5, self.evaluate(train_step(input_iterator)))

  @combinations.generate(
      combinations.combine(
          distribution=strategy_combinations.all_strategies,
          mode=["eager"]
      ))
  def testDatasetOutOfRange(self, distribution):
    with distribution.scope():
      a = variables.Variable(
          0.0, aggregation=variables.VariableAggregation.SUM)

    def train_step(val):
      a.assign_add(math_ops.reduce_sum(val))

    @def_function.function
    def f_train_step(iterator):
      distribution.run(train_step, args=(next(iterator),))
      return a

    dataset = get_dataset_from_tensor_slices([5., 6., 7., 8.]).batch(2)
    dist_dataset = distribution.experimental_distribute_dataset(dataset)

    iterator = iter(dist_dataset)
    with self.assertRaises(errors.OutOfRangeError):
      for _ in range(100):
        f_train_step(iterator)

    self.assertAlmostEqual(26.0, a.numpy())


if __name__ == "__main__":
  test.main()
