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

import itertools

import uuid
from absl.testing import parameterized
from tensorflow.python.checkpoint import checkpoint as trackable_utils
from tensorflow.python.checkpoint import checkpoint_management as ckpt_manager
from tensorflow.python.distribute import collective_all_reduce_strategy
from tensorflow.python.distribute import combinations
from tensorflow.python.distribute import distribute_lib
from tensorflow.python.distribute import strategy_combinations
from tensorflow.python.distribute import strategy_test_lib
from tensorflow.python.distribute import test_util
from tensorflow.python.distribute import values
from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
from tensorflow.python.eager import context
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 indexed_slices
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import array_ops_stack
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variable_v1
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.util import variable_utils


def strategy_and_run_tf_function_combinations():
  # Test the combination of different strategies and whether a tf.function
  # is passed into strategy.run."""
  # TODO(b/197981388): re-enable MWMS test
  # return combinations.combine(
  #     distribution=[
  #         strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
  #     ],
  #     mode=["graph", "eager"],
  #     experimental_run_tf_function=[True, False],
  #     use_var_policy=[True, False]) +
  return combinations.combine(
      distribution=[
          strategy_combinations.tpu_strategy,
          strategy_combinations.tpu_strategy_packed_var,
      ],
      mode=["graph", "eager"],
      experimental_run_tf_function=[True],
      use_var_policy=[True, False])


def strategy_with_var_policy():
  return combinations.combine(
      distribution=[
          strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
          # TODO(b/197981388): re-enable MWMS test
          # strategy_combinations.multi_worker_mirrored_2x1_cpu,
          # strategy_combinations.multi_worker_mirrored_2x1_gpu,
          strategy_combinations.tpu_strategy,
          strategy_combinations.tpu_strategy_packed_var,
      ],
      mode=["graph", "eager"],
      use_var_policy=[True, False])


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

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAssign(self, distribution, experimental_run_tf_function):

    def assign(fn, v, update_value, cross_replica):
      update_fn = lambda: getattr(v, fn)(update_value)
      if cross_replica:
        return update_fn()
      else:
        if experimental_run_tf_function:
          update_fn = def_function.function(update_fn)
        return test_util.gather(distribution, distribution.run(update_fn))

    updates = [("assign", 1.), ("assign_add", 1.), ("assign_sub", -1.)]
    aggregations = [
        variables_lib.VariableAggregation.NONE,
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    options = list(
        x for x in itertools.product(updates, aggregations, [True, False]))
    for update, aggregation, cross_replica in options:
      # assign in replica context with SUM does not make sense cause you can
      # just do value * num replicas error is 1. is not a distributed value and
      # is unsupported for aggregation SUM
      if (not cross_replica and aggregation ==
          variables_lib.VariableAggregation.SUM):
        continue
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      fn, update_value = update
      self.evaluate(assign(fn, v, update_value, cross_replica))
      for component in v._values:
        self.assertAllEqual(self.evaluate(component.read_value()),
                            self.evaluate(array_ops.ones_like(component)))

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAssignOnWriteVar(self, distribution, experimental_run_tf_function):

    with distribution.scope():
      v_to_assign = variable_v1.VariableV1(
          2., aggregation=variables_lib.VariableAggregation.MEAN)
      v_to_assign_sub = variable_v1.VariableV1(
          -2., aggregation=variables_lib.VariableAggregation.MEAN)

    def assign(fn, v, update_value, cross_replica):
      update_fn = lambda: getattr(v, fn)(update_value)
      if cross_replica:
        return update_fn()
      else:
        if experimental_run_tf_function:
          update_fn = def_function.function(update_fn)
        return test_util.gather(distribution, distribution.run(update_fn))

    updates = [("assign", v_to_assign), ("assign_add", v_to_assign),
               ("assign_sub", v_to_assign_sub)]
    aggregations = [
        variables_lib.VariableAggregation.NONE,
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    options = list(
        x for x in itertools.product(updates, aggregations, [True, False]))
    for update, aggregation, cross_replica in options:
      # assign in replica context with SUM does not make sense cause you can
      # just do value * num replicas error is 1. is not a distributed value and
      # is unsupported for aggregation SUM
      if aggregation == variables_lib.VariableAggregation.SUM:
        continue
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      fn, update_value = update
      self.evaluate(assign(fn, v, update_value, cross_replica))
      for component in v._values:
        self.assertAllEqual(2.0, self.evaluate(component.read_value()))

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAssignPerReplicaVal(self, distribution, experimental_run_tf_function):

    if strategy_test_lib.is_tpu_strategy(distribution):
      self.skipTest("Assigning PerReplica values is not supported. See"
                    " sponge/80ba41f8-4220-4516-98ce-bbad48f9f11a.")

    with distribution.scope():
      per_replica_value = values.PerReplica(
          [constant_op.constant(2.0),
           constant_op.constant(2.0)])
      per_replica_sub_value = values.PerReplica(
          [constant_op.constant(-2.0),
           constant_op.constant(-2.0)])

    def assign(fn, v, update_value, cross_replica):
      update_fn = lambda: getattr(v, fn)(update_value)
      if cross_replica:
        return update_fn()
      else:
        if experimental_run_tf_function:
          update_fn = def_function.function(update_fn)
        return test_util.gather(distribution, distribution.run(update_fn))

    updates = [("assign", per_replica_value), ("assign_add", per_replica_value),
               ("assign_sub", per_replica_sub_value)]
    # We don't support assigning PerReplica valus to vars in replica context
    # with aggregation=NONE.
    aggregations = [
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    options = list(
        x for x in itertools.product(updates, aggregations, [True, False]))
    for update, aggregation, cross_replica in options:
      # assign in replica context with SUM does not make sense cause you can
      # just do value * num replicas error is 1. is not a distributed value and
      # is unsupported for aggregation SUM
      if cross_replica:
        # We don't support assigning PerReplica values to MirroredVariables in
        # cross replica context
        continue
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      fn, update_value = update
      self.evaluate(assign(fn, v, update_value, cross_replica))
      if aggregation == variables_lib.VariableAggregation.SUM:
        expected = 4.0
      else:
        expected = 2.0
      for component in v._values:
        self.assertAllEqual(expected, self.evaluate(component.read_value()))

  @combinations.generate(strategy_with_var_policy())
  def testValueInReplicaContext(self, distribution):
    with distribution.scope():
      v = variables_lib.Variable(
          1., aggregation=variables_lib.VariableAggregation.MEAN)
      self.evaluate(variables_lib.global_variables_initializer())

      @def_function.function
      def f():
        with ops.control_dependencies([v.assign_add(1.)]):
          return v.value()

      results = self.evaluate(
          test_util.gather(distribution, distribution.run(f)))
      for value in results:
        self.assertEqual(2., value)

  @combinations.generate(strategy_with_var_policy())
  def testValueInReplicaContextAssignDirectValue(self, distribution,
                                                 use_var_policy):
    with distribution.scope():
      v = variables_lib.Variable(
          1., aggregation=variables_lib.VariableAggregation.MEAN)
      self.evaluate(variables_lib.global_variables_initializer())

      @def_function.function
      def f():
        with ops.control_dependencies([v.assign_add(1.)]):
          return v.value()

      results = self.evaluate(
          test_util.gather(distribution, distribution.run(f)))
      for value in results:
        self.assertEqual(2., value)

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testReadValueInReplicaContext(self, distribution,
                                    experimental_run_tf_function):
    aggregations = [
        variables_lib.VariableAggregation.NONE,
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    for aggregation in aggregations:
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      if experimental_run_tf_function:
        read_var_fn = def_function.function(v.read_value)
      else:
        read_var_fn = v.read_value
      results = self.evaluate(
          test_util.gather(distribution, distribution.run(read_var_fn)))
      for component, value in zip(v._values, results):
        self.assertAllEqual(self.evaluate(component.read_value()), value)

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testReadValueInCrossReplicaContext(self, distribution,
                                         experimental_run_tf_function):
    aggregations = [
        variables_lib.VariableAggregation.NONE,
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    for aggregation in aggregations:
      with distribution.scope():
        v = variable_v1.VariableV1(
            2.,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())

      if experimental_run_tf_function:
        read_var_fn = def_function.function(v.read_value)
      else:
        read_var_fn = v.read_value

      results = read_var_fn()
      for component in v._values:
        self.assertEqual(self.evaluate(component.read_value()),
                         self.evaluate(results))

  @combinations.generate(strategy_with_var_policy())
  def testAssignOutOfScope(self, distribution):
    with distribution.scope():
      mirrored = variables_lib.Variable(1.)
    self.evaluate(mirrored.assign(3.))
    self.assertEqual(self.evaluate(mirrored.read_value()), 3.)
    for component in mirrored.values:
      self.assertEqual(self.evaluate(component.read_value()), 3.)

  @combinations.generate(strategy_with_var_policy())
  def testInitializedToSameValueInsideEagerRun(self, distribution):
    if not context.executing_eagerly(): self.skipTest("eager only test")
    if isinstance(distribution.extended,
                  collective_all_reduce_strategy.CollectiveAllReduceExtended):
      self.skipTest("Test for more than 1 device per worker only.")
    v = [None]

    @def_function.function
    def step():

      def f():
        if v[0] is None:
          v[0] = variables_lib.Variable(random_ops.random_normal([]))

      distribution.run(f)

    context.set_global_seed(None)
    step()
    vals = self.evaluate(v[0].values)
    self.assertAllEqual(vals[0], vals[1])

  @combinations.generate(strategy_with_var_policy())
  def testAggregationOnlyFirstReplica(self, distribution):
    if isinstance(distribution.extended,
                  collective_all_reduce_strategy.CollectiveAllReduceExtended):
      self.skipTest("b/212945803")
    with distribution.scope():
      v = variable_v1.VariableV1(
          15.,
          synchronization=variables_lib.VariableSynchronization.ON_WRITE,
          aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA)
    self.evaluate(variables_lib.global_variables_initializer())

    @def_function.function
    def assign():
      ctx = distribute_lib.get_replica_context()
      replica_id = ctx.replica_id_in_sync_group
      return v.assign(math_ops.cast(replica_id, dtypes.float32))

    per_replica_results = self.evaluate(
        test_util.gather(distribution, distribution.run(assign)))
    # The per-replica values should always match the first replicas value.
    self.assertAllEqual(
        array_ops.zeros(distribution.num_replicas_in_sync, dtypes.float32),
        per_replica_results)

  @combinations.generate(strategy_with_var_policy())
  def testInitScope(self, distribution):
    if not context.executing_eagerly(): self.skipTest("eager only")

    class C(object):
      pass

    obj = C()
    obj.w = None
    obj.v = None

    @def_function.function
    def assign():
      with ops.init_scope():
        if obj.w is None:
          obj.w = variables_lib.Variable(
              0., aggregation=variables_lib.VariableAggregation.MEAN)
          obj.v = variables_lib.Variable(
              obj.w.read_value(),
              aggregation=variables_lib.VariableAggregation.MEAN)
          self.evaluate(variables_lib.global_variables_initializer())

      return obj.v.assign_add(2.)

    per_replica_results = self.evaluate(
        test_util.gather(distribution, distribution.run(assign)))
    self.assertAllEqual([2., 2.], per_replica_results)

  @combinations.generate(strategy_with_var_policy())
  def testOperatorOverride(self, distribution):

    if not context.executing_eagerly() and isinstance(
        distribution.extended,
        collective_all_reduce_strategy.CollectiveAllReduceExtended):
      self.skipTest("b/212954197")

    with distribution.scope():
      v = variable_v1.VariableV1(
          1, aggregation=variables_lib.VariableAggregation.SUM)
      self.evaluate(variables_lib.global_variables_initializer())

    self.assertEqual(2, self.evaluate(v + 1))

    @def_function.function
    def add():
      return v + 1

    per_replica_results = self.evaluate(
        test_util.gather(distribution, distribution.run(add)))
    self.assertAllEqual([2, 2], per_replica_results)

  @combinations.generate(
      combinations.combine(
          strategy=[
              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
              strategy_combinations.tpu_strategy,
              strategy_combinations.tpu_strategy_packed_var,
              strategy_combinations.multi_worker_mirrored_2x1_cpu,
              strategy_combinations.multi_worker_mirrored_2x1_gpu,
          ],
          mode=["eager"],
          use_var_policy=[True, False]))
  def testSaveAndRestoreOnWrite(self, strategy):
    aggregation = [
        variable_scope.VariableAggregation.NONE,
        variable_scope.VariableAggregation.ONLY_FIRST_REPLICA,
        variable_scope.VariableAggregation.SUM,
        variable_scope.VariableAggregation.MEAN
    ]
    for agg in aggregation:
      v_normal_restore = variables_lib.Variable(1.0)
      v_normal_save = variables_lib.Variable(3.0)
      with strategy.scope():
        v_on_write = variables_lib.Variable(2.0, aggregation=agg)

        # Save ONWRITE Restore ONWRITE
        # Save
        ckpt = trackable_utils.Checkpoint(var=v_on_write)
        manager = ckpt_manager.CheckpointManager(
            ckpt, "/tmp/ckpt_" + str(uuid.uuid4()), max_to_keep=None)
        manager.save()
        # Restore
        ckpt.restore(manager.latest_checkpoint)
        self.assertEqual(2.0, self.evaluate(v_on_write._values[0]))
        self.assertEqual(2.0, self.evaluate(v_on_write.read_value()))

        # Save Mirrored Restore Normal
        # We've already saved Mirrored, so we only need to restore normal
        ckpt_normal = trackable_utils.Checkpoint(var=v_normal_restore)
        ckpt_normal.restore(manager.latest_checkpoint)
        self.assertEqual(2.0, self.evaluate(v_on_write._values[0]))
        self.assertEqual(2.0, self.evaluate(v_normal_restore.read_value()))

        # Save Normal Restore Mirrored
        # Save
        ckpt = trackable_utils.Checkpoint(var=v_normal_save)
        manager_2 = ckpt_manager.CheckpointManager(
            ckpt, "/tmp/ckptckpt_" + str(uuid.uuid4()), max_to_keep=None)
        manager_2.save()
        # Restore
        ckpt_on_write = trackable_utils.Checkpoint(var=v_on_write)
        ckpt_on_write.restore(manager_2.latest_checkpoint)
        self.assertEqual(3.0, self.evaluate(v_on_write._values[0]))
        self.assertEqual(3.0, self.evaluate(v_on_write.read_value()))


ms_combination = combinations.combine(
    distribution=[strategy_combinations.mirrored_strategy_with_gpu_and_cpu],
    mode=["graph", "eager"])
tpu_combination = combinations.combine(
    distribution=[strategy_combinations.tpu_strategy_packed_var],
    mode=["graph", "eager"])


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

  @combinations.generate(ms_combination)
  def testScatterSub(self, distribution):
    with distribution.scope():
      v = variables_lib.Variable(
          [0., 0., 0.], aggregation=variables_lib.VariableAggregation.MEAN)
    self.evaluate(v.initializer)

    @def_function.function
    def scatter_sub():
      ctx = distribute_lib.get_replica_context()
      replica_id = ctx.replica_id_in_sync_group
      value = indexed_slices.IndexedSlices(
          values=array_ops_stack.stack([
              math_ops.cast(replica_id, dtypes.float32),
              math_ops.cast(replica_id + 1, dtypes.float32)
          ]),
          indices=array_ops_stack.stack([replica_id, replica_id + 1]),
          dense_shape=(3,))
      return v.scatter_sub(value)

    per_replica_results = self.evaluate(
        distribution.experimental_local_results(
            distribution.run(scatter_sub)))
    self.assertAllEqual([[0., -1., -1.], [0., -1., -1.]], per_replica_results)

  @combinations.generate(ms_combination)
  def testScatterAdd(self, distribution):
    with distribution.scope():
      v = variables_lib.Variable(
          [0, 0, 0], aggregation=variables_lib.VariableAggregation.SUM)
    self.evaluate(v.initializer)

    @def_function.function
    def scatter_add():
      ctx = distribute_lib.get_replica_context()
      replica_id = ctx.replica_id_in_sync_group
      value = indexed_slices.IndexedSlices(
          values=array_ops_stack.stack([replica_id, replica_id + 1]),
          indices=array_ops_stack.stack([replica_id, replica_id + 1]),
          dense_shape=(3,))
      return v.scatter_add(value)

    per_replica_results = self.evaluate(
        test_util.gather(distribution, distribution.run(scatter_add)))
    self.assertAllEqual([[0, 2, 2], [0, 2, 2]], per_replica_results)

  @combinations.generate(ms_combination)
  def testScatterDiv(self, distribution):
    with distribution.scope():
      v = variables_lib.Variable(
          [1, 6, 1], aggregation=variables_lib.VariableAggregation.SUM)
    self.evaluate(v.initializer)

    @def_function.function
    def scatter_div():
      ctx = distribute_lib.get_replica_context()
      replica_id = ctx.replica_id_in_sync_group
      value = indexed_slices.IndexedSlices(
          values=array_ops.reshape(replica_id + 2, [1]),
          indices=array_ops.reshape(replica_id, [1]),
          dense_shape=(3,))
      return v.scatter_div(value)

    per_replica_results = self.evaluate(
        test_util.gather(distribution, distribution.run(scatter_div)))
    self.assertAllEqual([[0, 2, 1], [0, 2, 1]], per_replica_results)

  @combinations.generate(ms_combination)
  def testScatterMul(self, distribution):
    with distribution.scope():
      v = variables_lib.Variable(
          [2., 1., 1.], aggregation=variables_lib.VariableAggregation.MEAN)
    self.evaluate(v.initializer)

    @def_function.function
    def scatter_mul():
      ctx = distribute_lib.get_replica_context()
      replica_id = ctx.replica_id_in_sync_group
      value = indexed_slices.IndexedSlices(
          values=array_ops.reshape(
              math_ops.cast(replica_id + 2, dtypes.float32), [1]),
          indices=array_ops.reshape(replica_id, [1]),
          dense_shape=(3,))
      return v.scatter_mul(value)

    per_replica_results = self.evaluate(
        test_util.gather(distribution, distribution.run(scatter_mul)))
    self.assertAllClose([[2., 1.5, 1.], [2., 1.5, 1.]], per_replica_results)

  @combinations.generate(ms_combination)
  def testScatterMin(self, distribution):
    with distribution.scope():
      v1 = variables_lib.Variable(
          [0, 2, 0], aggregation=variables_lib.VariableAggregation.SUM)
      v2 = variables_lib.Variable(
          [0, 2, 0],
          aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA)
    self.evaluate(variables_lib.global_variables_initializer())

    @def_function.function
    def scatter_min(v):
      value = indexed_slices.IndexedSlices(
          values=array_ops.identity([1]),
          indices=array_ops.identity([1]),
          dense_shape=(3,))
      return v.scatter_min(value)

    with self.assertRaisesRegex(NotImplementedError, "scatter_min.*"):
      self.evaluate(
          test_util.gather(distribution,
                           distribution.run(scatter_min, args=(v1,))))

    per_replica_results = self.evaluate(
        test_util.gather(distribution,
                         distribution.run(scatter_min, args=(v2,))))
    self.assertAllClose([[0, 1, 0], [0, 1, 0]], per_replica_results)

  @combinations.generate(ms_combination)
  def testScatterMax(self, distribution):
    with distribution.scope():
      v1 = variables_lib.Variable(
          [0, 0, 0], aggregation=variables_lib.VariableAggregation.SUM)
      v2 = variables_lib.Variable(
          [0, 0, 0],
          aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA)
    self.evaluate(variables_lib.global_variables_initializer())

    @def_function.function
    def scatter_max(v):
      value = indexed_slices.IndexedSlices(
          values=array_ops.identity([1]),
          indices=array_ops.identity([0]),
          dense_shape=(3,))
      return v.scatter_max(value)

    with self.assertRaisesRegex(NotImplementedError, "scatter_max.*"):
      self.evaluate(
          test_util.gather(distribution,
                           distribution.run(scatter_max, args=(v1,))))

    per_replica_results = self.evaluate(
        test_util.gather(distribution,
                         distribution.run(scatter_max, args=(v2,))))
    self.assertAllClose([[1, 0, 0], [1, 0, 0]], per_replica_results)

  @combinations.generate(ms_combination)
  def testScatterUpdate(self, distribution):
    with distribution.scope():
      v1 = variables_lib.Variable(
          [0, 0, 0], aggregation=variables_lib.VariableAggregation.SUM)
      v2 = variables_lib.Variable(
          [0, 0, 0],
          aggregation=variables_lib.VariableAggregation.ONLY_FIRST_REPLICA)
    self.evaluate(variables_lib.global_variables_initializer())

    @def_function.function
    def scatter_update(v):
      value = indexed_slices.IndexedSlices(
          values=array_ops.identity([3]),
          indices=array_ops.identity([1]),
          dense_shape=(3,))
      return v.scatter_update(value)

    with self.assertRaisesRegex(NotImplementedError, "scatter_update.*"):
      self.evaluate(
          test_util.gather(distribution,
                           distribution.run(scatter_update, args=(v1,))))

    per_replica_results = self.evaluate(
        test_util.gather(distribution,
                         distribution.run(scatter_update, args=(v2,))))
    self.assertAllClose([[0, 3, 0], [0, 3, 0]], per_replica_results)

  @combinations.generate(ms_combination + tpu_combination)
  def testScatterOpsWithNoneAggregation(self, distribution):

    def assert_close(v, op, delta, expect):
      scatter_op = getattr(v, op)

      @def_function.function
      def scatter_xxx():
        return scatter_op(delta)

      per_replica_results = self.evaluate(
          variable_utils.convert_variables_to_tensors(
              distribution.experimental_local_results(
                  distribution.run(scatter_xxx))))
      self.assertAllClose([expect, expect], per_replica_results)

    with distribution.scope():
      v = variables_lib.Variable(
          [4.], aggregation=variables_lib.VariableAggregation.NONE)
    self.evaluate(variables_lib.global_variables_initializer())

    delta = indexed_slices.IndexedSlices(
        values=array_ops.identity([2.]),
        indices=array_ops.identity([0]),
        dense_shape=(1,))

    assert_close(v, "scatter_sub", delta, [2.])
    assert_close(v, "scatter_add", delta, [4.])
    assert_close(v, "scatter_max", delta, [4.])
    assert_close(v, "scatter_min", delta, [2.])
    assert_close(v, "scatter_mul", delta, [4.])
    assert_close(v, "scatter_div", delta, [2.])
    assert_close(v, "scatter_update", delta, [2.])

  @combinations.generate(ms_combination + tpu_combination)
  def testScatterOpsInCrossReplicaContext(self, distribution):
    with distribution.scope():
      v1 = variables_lib.Variable(
          [1, 1, 1], aggregation=variables_lib.VariableAggregation.SUM)
      v2 = variables_lib.Variable([1, 1, 1])
    self.evaluate(variables_lib.global_variables_initializer())

    value = indexed_slices.IndexedSlices(
        values=array_ops.identity([2]),
        indices=array_ops.identity([0]),
        dense_shape=(3,))
    with distribution.scope():
      self.evaluate(v1.scatter_add(value))
      self.assertAllEqual([3, 1, 1], self.evaluate(v1.read_value()))

      self.evaluate(v2.scatter_min(value))
      self.assertAllEqual([1, 1, 1], self.evaluate(v2.read_value()))


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

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAssign(self, distribution, experimental_run_tf_function):

    def assign(fn, v, update_value, cross_replica):
      update_fn = lambda: getattr(v, fn)(update_value)
      if cross_replica:
        return update_fn()
      else:
        if experimental_run_tf_function:
          update_fn = def_function.function(update_fn)
        return test_util.gather(distribution, distribution.run(update_fn))

    updates = [("assign", 1.), ("assign_add", 1.), ("assign_sub", -1.)]
    aggregations = [
        variables_lib.VariableAggregation.NONE,
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    options = list(
        x for x in itertools.product(updates, aggregations, [True, False]))
    for update, aggregation, cross_replica in options:
      # VariableAggregation.SUM in cross-replica mode is tested below,
      # VariableAggregation.NONE in cross-replica mode is not supported.
      if cross_replica and aggregation in [
          variables_lib.VariableAggregation.SUM,
          variables_lib.VariableAggregation.NONE,
      ]:
        continue
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            synchronization=variables_lib.VariableSynchronization.ON_READ,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      fn, update_value = update
      self.evaluate(assign(fn, v, update_value, cross_replica))
      for component in v._values:
        self.assertAllEqual(self.evaluate(component.read_value()),
                            self.evaluate(array_ops.ones_like(component)))

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAssignOnReadVar(self, distribution, experimental_run_tf_function):

    with distribution.scope():
      v_to_assign = variable_v1.VariableV1(
          2., aggregation=variables_lib.VariableAggregation.MEAN)
      v_to_assign_sub = variable_v1.VariableV1(
          -2., aggregation=variables_lib.VariableAggregation.MEAN)

    def assign(fn, v, update_value, cross_replica):
      update_fn = lambda: getattr(v, fn)(update_value)
      if cross_replica:
        return update_fn()
      else:
        if experimental_run_tf_function:
          update_fn = def_function.function(update_fn)
        return test_util.gather(distribution, distribution.run(update_fn))

    updates = [("assign", v_to_assign), ("assign_add", v_to_assign),
               ("assign_sub", v_to_assign_sub)]
    expected_cross_replica = {
        variables_lib.VariableAggregation.SUM: 1.0,
        variables_lib.VariableAggregation.MEAN: 2.0,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA: 2.0
    }
    expected_replica = {
        variables_lib.VariableAggregation.SUM: 2.0,
        variables_lib.VariableAggregation.MEAN: 2.0,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA: 2.0
    }
    # aggregation=NONE is not supported for OnReadVariables.
    aggregations = [
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    options = list(
        x for x in itertools.product(updates, aggregations, [True, False]))
    for update, aggregation, cross_replica in options:
      # assign in replica context with SUM does not make sense cause you can
      # just do value * num replicas error is 1. is not a distributed value and
      # is unsupported for aggregation SUM
      if aggregation == variables_lib.VariableAggregation.SUM:
        continue
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      fn, update_value = update
      self.evaluate(assign(fn, v, update_value, cross_replica))
      if cross_replica:
        for component in v._values:
          self.assertAllEqual(expected_cross_replica.get(aggregation),
                              self.evaluate(component.read_value()))
      else:
        for component in v._values:
          self.assertAllEqual(expected_replica.get(aggregation),
                              self.evaluate(component.read_value()))

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAssignPerReplicaVal(self, distribution, experimental_run_tf_function):

    if strategy_test_lib.is_tpu_strategy(distribution):
      self.skipTest("Assigning PerReplica values is not supported. See"
                    " sponge/80ba41f8-4220-4516-98ce-bbad48f9f11a.")

    self.skipTest("We don't support assiging PerReplica values in cross "
                  "replica context or replica context. see error in "
                  "sponge/2b2e54c1-eda6-4534-82e1-c73b1dcd517f.")

    with distribution.scope():
      per_replica_value = values.PerReplica(
          [constant_op.constant(2.0),
           constant_op.constant(2.0)])

    def assign(fn, v, update_value, cross_replica):
      update_fn = lambda: getattr(v, fn)(update_value)
      if cross_replica:
        return update_fn()
      else:
        if experimental_run_tf_function:
          update_fn = def_function.function(update_fn)
        return test_util.gather(distribution, distribution.run(update_fn))

    updates = [("assign", per_replica_value)]
    # We don't support assigning PerReplica valus to vars in replica context
    # with aggregation=NONE.
    aggregations = [
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    options = list(
        x for x in itertools.product(updates, aggregations, [True, False]))
    for update, aggregation, cross_replica in options:
      # assign in replica context with SUM does not make sense cause you can
      # just do value * num replicas error is 1. is not a distributed value and
      # is unsupported for aggregation SUM
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            synchronization=variables_lib.VariableSynchronization.ON_READ,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      fn, update_value = update
      # with self.assertRaisesRegex(ValueError, "Attempt to convert a value "):
      self.evaluate(assign(fn, v, update_value, cross_replica))
      if aggregation == variables_lib.VariableAggregation.SUM:
        expected = 4.0
      else:
        expected = 2.0
      for component in v._values:
        self.assertAllEqual(expected, self.evaluate(component.read_value()))

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAssignDtypeConversion(self, distribution,
                                experimental_run_tf_function):

    def assign(fn, v, update_value, cross_replica):
      update_fn = lambda: getattr(v, fn)(update_value)
      if cross_replica:
        return update_fn()
      else:
        if experimental_run_tf_function:
          update_fn = def_function.function(update_fn)
        return test_util.gather(distribution, distribution.run(update_fn))

    updates = [("assign", 1), ("assign_add", 1), ("assign_sub", -1)]
    aggregations = [
        variables_lib.VariableAggregation.NONE,
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    options = list(
        x for x in itertools.product(updates, aggregations, [True, False]))
    for update, aggregation, cross_replica in options:
      # VariableAggregation.SUM in cross-replica mode is tested below,
      # VariableAggregation.NONE in cross-replica mode is not supported.
      if cross_replica and aggregation in [
          variables_lib.VariableAggregation.SUM,
          variables_lib.VariableAggregation.NONE,
      ]:
        continue
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            synchronization=variables_lib.VariableSynchronization.ON_READ,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      fn, update_value = update
      self.evaluate(assign(fn, v, update_value, cross_replica))
      for component in v._values:
        self.assertAllEqual(self.evaluate(component.read_value()),
                            self.evaluate(array_ops.ones_like(component)))

  @combinations.generate(strategy_with_var_policy())
  def testAssignWithAggregationSum(self, distribution):
    with distribution.scope():
      v = variable_v1.VariableV1(
          0.,
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=variables_lib.VariableAggregation.SUM)
    self.evaluate(variables_lib.global_variables_initializer())
    self.evaluate(v.assign(1. * distribution.num_replicas_in_sync))
    for component in v._values:
      self.assertAllEqual(self.evaluate(component.read_value()),
                          self.evaluate(array_ops.ones_like(component)))

  @combinations.generate(strategy_with_var_policy())
  def testAssignAddSubWithAggregationSum(self, distribution):
    with distribution.scope():
      v = variable_v1.VariableV1(
          0.,
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=variables_lib.VariableAggregation.SUM)
    self.evaluate(variables_lib.global_variables_initializer())
    with self.assertRaisesRegex(
        ValueError, "SyncOnReadVariable does not support "):
      self.evaluate(v.assign_add(1.))
    with self.assertRaisesRegex(
        ValueError, "SyncOnReadVariable does not support "):
      self.evaluate(v.assign_sub(1.))

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testReadValueInReplicaContext(self, distribution,
                                    experimental_run_tf_function):
    aggregations = [
        variables_lib.VariableAggregation.NONE,
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    for aggregation in aggregations:
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            synchronization=variables_lib.VariableSynchronization.ON_READ,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())
      if experimental_run_tf_function:
        read_var_fn = def_function.function(v.read_value)
      else:
        read_var_fn = v.read_value
      results = self.evaluate(
          test_util.gather(distribution, distribution.run(read_var_fn)))
      for component, value in zip(v._values, results):
        self.assertAllEqual(self.evaluate(component.read_value()), value)

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testReadValueInCrossReplicaContext(self, distribution,
                                         experimental_run_tf_function):
    aggregations = [
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    for aggregation in aggregations:
      if strategy_test_lib.is_tpu_strategy(distribution):
        resolver = tpu_cluster_resolver.TPUClusterResolver("")
        tpu_cluster_resolver.initialize_tpu_system(resolver)
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            synchronization=variables_lib.VariableSynchronization.ON_READ,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())

      def assign(v=v):
        ctx = distribute_lib.get_replica_context()
        replica_id = ctx.replica_id_in_sync_group
        return v.assign(math_ops.cast(replica_id, dtypes.float32))

      if experimental_run_tf_function:
        assign = def_function.function(assign)

      self.evaluate(test_util.gather(distribution, distribution.run(assign)))
      num_replicas = distribution.num_replicas_in_sync
      sum_of_replica_values = num_replicas * (num_replicas - 1) / 2.
      if aggregation == variables_lib.VariableAggregation.SUM:
        expected = sum_of_replica_values
      elif aggregation == variables_lib.VariableAggregation.MEAN:
        expected = sum_of_replica_values / num_replicas
      else:
        expected = 0
      self.assertEqual(expected, self.evaluate(v.read_value()), aggregation)
      self.assertEqual(expected, self.evaluate(v.value()), aggregation)
      self.assertEqual(expected, self.evaluate(v), aggregation)
      self.assertEqual(expected, self.evaluate(array_ops.identity(v)),
                       aggregation)

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAllReduce(self, distribution, experimental_run_tf_function):
    with distribution.scope():
      v = variable_v1.VariableV1(
          2.,
          synchronization=variables_lib.VariableSynchronization.ON_WRITE,
          aggregation=variables_lib.VariableAggregation.MEAN)
    self.evaluate(variables_lib.global_variables_initializer())

    def all_reduce():
      ctx = distribute_lib.get_replica_context()
      replica_id = ctx.replica_id_in_sync_group
      return ctx.all_reduce("SUM", v) + math_ops.cast(replica_id,
                                                      dtypes.float32)

    if experimental_run_tf_function:
      all_reduce = def_function.function(all_reduce)

    per_replica_results = self.evaluate(
        test_util.gather(distribution, distribution.run(all_reduce)))
    expected_result = []
    for i in range(distribution.num_replicas_in_sync):
      expected_result.append(2.0 * distribution.num_replicas_in_sync +
                             1.0 * i)
    self.assertAllEqual(per_replica_results, tuple(expected_result))

  @combinations.generate(strategy_and_run_tf_function_combinations())
  def testAssignPerReplicaBeforeRead(self, distribution,
                                     experimental_run_tf_function):
    aggregations = [
        variables_lib.VariableAggregation.SUM,
        variables_lib.VariableAggregation.MEAN,
        variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
    ]
    for aggregation in aggregations:
      with distribution.scope():
        v = variable_v1.VariableV1(
            0.,
            synchronization=variables_lib.VariableSynchronization.ON_READ,
            aggregation=aggregation)
      self.evaluate(variables_lib.global_variables_initializer())

      def assign(var=v):
        ctx = distribute_lib.get_replica_context()
        replica_id = ctx.replica_id_in_sync_group
        return var.assign(math_ops.cast(replica_id, dtypes.float32))

      if experimental_run_tf_function:
        assign = def_function.function(assign)

      per_replica_results = self.evaluate(
          test_util.gather(distribution, distribution.run(assign)))
      expected_result = []
      for i in range(distribution.num_replicas_in_sync):
        expected_result.append(1.0 * i)
      self.assertAllEqual(per_replica_results, tuple(expected_result))

  @combinations.generate(strategy_with_var_policy())
  def testReadValueWithAggregationNoneInCrossReplicaContext(self, distribution):
    with distribution.scope():
      v = variable_v1.VariableV1(
          0.,
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=variables_lib.VariableAggregation.NONE)
    self.evaluate(variables_lib.global_variables_initializer())
    with self.assertRaisesRegex(
        ValueError, "Could not convert from .* VariableAggregation\\.NONE"):
      self.evaluate(v.read_value())

  @combinations.generate(strategy_with_var_policy())
  def testInitializedToSameValueInsideEagerRun(self, distribution):
    if not context.executing_eagerly(): self.skipTest("eager only")
    if isinstance(distribution.extended,
                  collective_all_reduce_strategy.CollectiveAllReduceExtended):
      self.skipTest("Test for more than 1 device per worker only.")

    v = [None]
    @def_function.function
    def step():
      def f():
        if v[0] is None:
          v[0] = variables_lib.Variable(
              random_ops.random_normal([]),
              synchronization=variables_lib.VariableSynchronization.ON_READ)

      distribution.run(f)

    context.set_global_seed(None)
    step()
    vals = self.evaluate(v[0].values)
    self.assertAllEqual(vals[0], vals[1])

  @combinations.generate(strategy_with_var_policy())
  def testOperatorOverride(self, distribution):

    with distribution.scope():
      v = variable_v1.VariableV1(
          0.0,
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=variables_lib.VariableAggregation.MEAN)
      self.evaluate(variables_lib.global_variables_initializer())

      @def_function.function
      def assign():
        ctx = distribute_lib.get_replica_context()
        replica_id = ctx.replica_id_in_sync_group
        return v.assign(math_ops.cast(replica_id, dtypes.float32))

      # Assign different replicas with different values.
      self.evaluate(test_util.gather(distribution, distribution.run(assign)))
      self.assertEqual(1.5, self.evaluate(v + 1))

      @def_function.function
      def add():
        return v + 1

      per_replica_results = self.evaluate(
          test_util.gather(distribution, distribution.run(add)))
      self.assertAllEqual([1, 2], per_replica_results)

  @combinations.generate(
      combinations.combine(
          strategy=[
              strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
              strategy_combinations.tpu_strategy,
              strategy_combinations.tpu_strategy_packed_var,
              strategy_combinations.multi_worker_mirrored_2x1_cpu,
              strategy_combinations.multi_worker_mirrored_2x1_gpu,
          ],
          mode=["eager"],
          use_var_policy=[True, False]))
  def testSaveAndRestoreOnRead(self, strategy):
    aggregation = [variable_scope.VariableAggregation.SUM,
                   variable_scope.VariableAggregation.MEAN]
    for agg in aggregation:
      v_normal_restore = variables_lib.Variable(1.0)
      v_normal_save = variables_lib.Variable(2.0)

      with strategy.scope():
        v_on_read = variables_lib.Variable(
            1.0, synchronization=variable_scope.VariableSynchronization.ON_READ,
            aggregation=agg)

        @def_function.function
        def assign_fn():
          cluster_resolver = strategy.cluster_resolver
          replica_ctx = distribute_lib.get_replica_context()
          if ((cluster_resolver and cluster_resolver.task_type == "worker") or
              math_ops.equal(replica_ctx.replica_id_in_sync_group,
                             constant_op.constant(1))):
            v_on_read.assign(3.)  # pylint:disable=cell-var-from-loop
          else:
            v_on_read.assign(4.)  # pylint:disable=cell-var-from-loop

        strategy.run(assign_fn)

        # Save ONREAD, restore ONREAD
        # Saves v[0] + v[1] = 7 for SUM and 3.5 for MEAN.
        ckpt = trackable_utils.Checkpoint(var=v_on_read)
        manager = ckpt_manager.CheckpointManager(
            ckpt, "/tmp/ckpt_" + str(uuid.uuid4()), max_to_keep=None)
        manager.save()
        # Restores a value of 7/2 = 3.5 for SUM and 3.5 for MEAN.
        ckpt.restore(manager.latest_checkpoint)
        self.assertEqual(3.5, self.evaluate(v_on_read._values[0]))

        # Save ONREAD, restore normal
        ckpt_normal = trackable_utils.Checkpoint(var=v_normal_restore)
        ckpt_normal.restore(manager.latest_checkpoint)
        if agg == variable_scope.VariableAggregation.SUM:
          self.assertEqual(7.0, self.evaluate(v_normal_restore.read_value()))
        else:
          self.assertEqual(3.5, self.evaluate(v_normal_restore.read_value()))

        # Save normal, restore ONREAD
        ckpt = trackable_utils.Checkpoint(var=v_normal_save)
        manager = ckpt_manager.CheckpointManager(
            ckpt, "/tmp/ckpt_" + str(uuid.uuid4()), max_to_keep=None)
        manager.save()
        # Restores a value of 2/2 = 1.0 for SUM and 2.0 for MEAN.
        ckpt_on_read = trackable_utils.Checkpoint(var=v_on_read)
        ckpt_on_read.restore(manager.latest_checkpoint)
        if agg == variable_scope.VariableAggregation.SUM:
          self.assertEqual(1.0, self.evaluate(v_on_read._values[0]))
        else:
          self.assertEqual(2.0, self.evaluate(v_on_read._values[0]))


@combinations.generate(
    combinations.combine(
        distribution=[
            strategy_combinations.mirrored_strategy_with_gpu_and_cpu,
            strategy_combinations.multi_worker_mirrored_2x1_cpu,
            strategy_combinations.multi_worker_mirrored_2x1_gpu,
        ],
        aggregation=[
            variables_lib.VariableAggregation.MEAN,
            variables_lib.VariableAggregation.SUM,
            variables_lib.VariableAggregation.ONLY_FIRST_REPLICA,
        ],
        mode=["graph", "eager"],
        use_var_policy=[True, False]))
class SyncOnReadScatterReplicaTest(test.TestCase, parameterized.TestCase):

  def testScatterSub(self, distribution, aggregation):
    with distribution.scope():
      v = variables_lib.Variable(
          [1., 1., 1.],
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=aggregation)
    self.evaluate(v.initializer)

    delta = values.PerReplica([
        indexed_slices.IndexedSlices(
            values=[[0.], [1.]], indices=[0, 1], dense_shape=(3,)),
        indexed_slices.IndexedSlices(
            values=[[1.], [2.]], indices=[1, 2], dense_shape=(3,)),
    ])

    with self.assertRaises(NotImplementedError):
      self.evaluate(distribution.run(v.scatter_sub, args=(delta,)))

  def testScatterAdd(self, distribution, aggregation):
    with distribution.scope():
      v = variables_lib.Variable(
          [1., 1., 1.],
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=aggregation)
    self.evaluate(v.initializer)

    delta = values.PerReplica([
        indexed_slices.IndexedSlices(
            values=[[0.], [1.]], indices=[0, 1], dense_shape=(3,)),
        indexed_slices.IndexedSlices(
            values=[[1.], [2.]], indices=[1, 2], dense_shape=(3,)),
    ])

    with self.assertRaises(NotImplementedError):
      self.evaluate(distribution.run(v.scatter_add, args=(delta,)))

  def testScatterDiv(self, distribution, aggregation):
    with distribution.scope():
      v = variables_lib.Variable(
          [2., 6., 1.],
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=aggregation)
    self.evaluate(v.initializer)

    delta = values.PerReplica([
        indexed_slices.IndexedSlices(
            values=[[2.], [2.]], indices=[0, 1], dense_shape=(3,)),
        indexed_slices.IndexedSlices(
            values=[[3.], [3.]], indices=[1, 2], dense_shape=(3,)),
    ])

    with self.assertRaises(NotImplementedError):
      self.evaluate(distribution.run(v.scatter_div, args=(delta,)))

  def testScatterMul(self, distribution, aggregation):
    with distribution.scope():
      v = variables_lib.Variable(
          [2., 1., 1.],
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=aggregation)
    self.evaluate(v.initializer)

    delta = values.PerReplica([
        indexed_slices.IndexedSlices(
            values=[[2.], [3.]], indices=[0, 1], dense_shape=(3,)),
        indexed_slices.IndexedSlices(
            values=[[4.], [5.]], indices=[1, 2], dense_shape=(3,)),
    ])

    with self.assertRaises(NotImplementedError):
      self.evaluate(distribution.run(v.scatter_mul, args=(delta,)))

  def testScatterMin(self, distribution, aggregation):
    with distribution.scope():
      v = variables_lib.Variable(
          [3., 4., 5.],
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=aggregation)
    self.evaluate(v.initializer)

    delta = values.PerReplica([
        indexed_slices.IndexedSlices(
            values=[[1.], [8.]], indices=[0, 1], dense_shape=(3,)),
        indexed_slices.IndexedSlices(
            values=[[9.], [2.]], indices=[1, 2], dense_shape=(3,)),
    ])

    with self.assertRaises(NotImplementedError):
      self.evaluate(distribution.run(v.scatter_min, args=(delta,)))

  def testScatterMax(self, distribution, aggregation):
    with distribution.scope():
      v = variables_lib.Variable(
          [3., 4., 5.],
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=aggregation)
    self.evaluate(v.initializer)

    delta = values.PerReplica([
        indexed_slices.IndexedSlices(
            values=[[1.], [8.]], indices=[0, 1], dense_shape=(3,)),
        indexed_slices.IndexedSlices(
            values=[[9.], [2.]], indices=[1, 2], dense_shape=(3,)),
    ])

    with self.assertRaises(NotImplementedError):
      self.evaluate(distribution.run(v.scatter_max, args=(delta,)))

  def testScatterUpdate(self, distribution, aggregation):
    with distribution.scope():
      v = variables_lib.Variable(
          [0., 0., 0.],
          synchronization=variables_lib.VariableSynchronization.ON_READ,
          aggregation=aggregation)
    self.evaluate(v.initializer)

    delta = values.PerReplica([
        indexed_slices.IndexedSlices(
            values=[[1.], [2.]], indices=[0, 1], dense_shape=(3,)),
        indexed_slices.IndexedSlices(
            values=[[3.], [4.]], indices=[1, 2], dense_shape=(3,)),
    ])

    with self.assertRaises(NotImplementedError):
      self.evaluate(distribution.run(v.scatter_min, args=(delta,)))


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