| # 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. |
| # ============================================================================== |
| """Tests for common methods in strategy classes.""" |
| |
| from absl.testing import parameterized |
| |
| from tensorflow.python.data.ops import dataset_ops |
| from tensorflow.python.distribute import central_storage_strategy |
| from tensorflow.python.distribute import combinations |
| from tensorflow.python.distribute import distribute_lib |
| from tensorflow.python.distribute import mirrored_strategy |
| from tensorflow.python.distribute import strategy_combinations |
| from tensorflow.python.distribute import test_util |
| from tensorflow.python.distribute import tpu_strategy |
| from tensorflow.python.distribute.collective_all_reduce_strategy import CollectiveAllReduceStrategy |
| from tensorflow.python.eager import def_function |
| from tensorflow.python.framework import constant_op |
| from tensorflow.python.framework import dtypes |
| from tensorflow.python.framework import errors |
| from tensorflow.python.framework import indexed_slices |
| from tensorflow.python.framework import test_util as tf_test_util |
| from tensorflow.python.ops import array_ops |
| from tensorflow.python.ops import gradients_impl |
| from tensorflow.python.platform import test |
| from tensorflow.python.util import nest |
| |
| |
| @tf_test_util.with_eager_op_as_function |
| @combinations.generate( |
| combinations.combine( |
| strategy=[ |
| strategy_combinations.default_strategy, |
| strategy_combinations.one_device_strategy, |
| strategy_combinations.one_device_strategy_gpu, |
| strategy_combinations.central_storage_strategy_with_two_gpus, |
| strategy_combinations.central_storage_strategy_with_gpu_and_cpu, |
| strategy_combinations.mirrored_strategy_with_one_cpu, |
| strategy_combinations.mirrored_strategy_with_one_gpu, |
| strategy_combinations.mirrored_strategy_with_two_gpus, |
| strategy_combinations.mirrored_strategy_with_two_cpus, |
| strategy_combinations.mirrored_strategy_with_gpu_and_cpu, |
| strategy_combinations.multi_worker_mirrored_2x2_gpu, |
| strategy_combinations.multi_worker_mirrored_2x1_cpu, |
| strategy_combinations.multi_worker_mirrored_2x1_gpu, |
| ], |
| mode=['eager'], |
| pure_eager=[True, False]) + combinations.combine( |
| strategy=[ |
| strategy_combinations.tpu_strategy, |
| strategy_combinations.tpu_strategy_packed_var, |
| strategy_combinations.tpu_strategy_one_step, |
| strategy_combinations.cloud_tpu_strategy, |
| ], |
| mode=['eager'], |
| pure_eager=[False])) |
| class GatherTest(test.TestCase, parameterized.TestCase): |
| |
| def _gather_same_shape_and_verify(self, value_on_replica, axis, pure_eager, |
| strategy): |
| distributed_values = strategy.experimental_distribute_values_from_function( |
| lambda _: array_ops.identity(value_on_replica)) |
| |
| def run(): |
| return strategy.gather(distributed_values, axis=axis) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| all_results = [ |
| value_on_replica for _ in range(strategy.num_replicas_in_sync) |
| ] |
| expected_result = array_ops.concat(all_results, axis=axis) |
| self.assertAllEqual(expected_result, run().numpy()) |
| |
| def testGatherPerReplicaDense1D0Axis(self, strategy, pure_eager): |
| """A DistributedValues object with two tensors of shape [3] on each replica gathers to a tensor of [6].""" |
| single_value = constant_op.constant([1, 2, 3]) |
| axis = 0 |
| self._gather_same_shape_and_verify(single_value, axis, pure_eager, strategy) |
| |
| def testGatherPerReplicaDense2D0Axis(self, strategy, pure_eager): |
| """A DistributedValues object with two tensors of [1, 3] on each replica gathers along 0th dim to a tensor of [2, 3].""" |
| single_value = constant_op.constant([[1, 2, 3]]) |
| axis = 0 |
| self._gather_same_shape_and_verify(single_value, axis, pure_eager, strategy) |
| |
| def testGatherPerReplicaDense2D1Axis(self, strategy, pure_eager): |
| """A DistributedValues object with two tensors of [1, 3] on each replica gathers along 1st dim to a tensor of [1, 6].""" |
| single_value = constant_op.constant([[1, 2, 3]]) |
| axis = 1 |
| self._gather_same_shape_and_verify(single_value, axis, pure_eager, strategy) |
| |
| def testGatherPerReplicaDense3D0Axis(self, strategy, pure_eager): |
| """A DistributedValues object with two tensors of [1, 2, 2] on each replica gathers along 0th dim to a tensor of [2, 2, 2].""" |
| single_value = constant_op.constant([[[1, 2], [1, 2]]]) |
| axis = 0 |
| self._gather_same_shape_and_verify(single_value, axis, pure_eager, strategy) |
| |
| def testGatherPerReplicaDense3D1Axis(self, strategy, pure_eager): |
| """A DistributedValues object with two tensors of [1, 2, 2] on each replica gathers along 1nd dimension to a tensor of [1, 4, 2].""" |
| single_value = constant_op.constant([[[1, 2], [1, 2]]]) |
| axis = 1 |
| self._gather_same_shape_and_verify(single_value, axis, pure_eager, strategy) |
| |
| def testGatherPerReplicaDense3D2Axis(self, strategy, pure_eager): |
| """A DistributedValues object with two tensors of [1, 2, 2] on each replica gathers along 2nd dimension to a tensor of [1, 2, 4].""" |
| single_value = constant_op.constant([[[1, 2], [1, 2]]]) |
| axis = 2 |
| self._gather_same_shape_and_verify(single_value, axis, pure_eager, strategy) |
| |
| def testGatherDiffShapeAtAxis0(self, strategy, pure_eager): |
| """Different `Axis`-th (0) dimension: shape [1, 1], [2, 1] -> [3, 1].""" |
| |
| def value_fn(ctx): |
| return constant_op.constant( |
| 1, shape=(ctx.replica_id_in_sync_group + 1, 1)) |
| |
| distributed_values = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| axis = 0 |
| |
| def run(): |
| return strategy.gather(distributed_values, axis=axis) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| expected_result = constant_op.constant( |
| 1, shape=(sum(range(strategy.num_replicas_in_sync + 1)), 1)) |
| |
| self.assertAllEqual(expected_result, run().numpy()) |
| |
| def testGatherDiffShapeAtAxis1(self, strategy, pure_eager): |
| """Different `Axis`-th (non-0) dimension: shape [1, 1], [1, 2] -> [1, 3].""" |
| |
| def value_fn(ctx): |
| return constant_op.constant( |
| 1, shape=(1, ctx.replica_id_in_sync_group + 1)) |
| |
| distributed_values = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| axis = 1 |
| |
| def run(): |
| return strategy.gather(distributed_values, axis=axis) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| expected_result = constant_op.constant( |
| 1, shape=(1, sum(range(strategy.num_replicas_in_sync + 1)))) |
| |
| self.assertAllEqual(expected_result, run().numpy()) |
| |
| def testGatherRaiseDiffShapeAtNonAxis(self, strategy, pure_eager): |
| """Different at non-`axis`-th dimension : [1, 1], [1, 2], 0th -> raise error.""" |
| if isinstance(strategy, CollectiveAllReduceStrategy |
| ) and _get_num_replicas_per_client(strategy) > 1: |
| self.skipTest('b/167331966') |
| |
| if strategy.num_replicas_in_sync <= 1: |
| self.skipTest('Test for more than 1 replica only.') |
| |
| def value_fn(ctx): |
| return constant_op.constant( |
| 1, shape=(1, ctx.replica_id_in_sync_group + 1)) |
| |
| distributed_values = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| axis = 0 |
| |
| def run(): |
| return strategy.gather(distributed_values, axis=axis) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| if isinstance(strategy, CollectiveAllReduceStrategy): |
| with self.assertRaisesRegex(errors.InvalidArgumentError, |
| r'Shape mismatch'): |
| run() |
| elif isinstance(strategy, |
| (mirrored_strategy.MirroredStrategy, |
| central_storage_strategy.CentralStorageStrategy)): |
| with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError), |
| r'Dimension \d in both shapes must be equal'): |
| run() |
| |
| def testGatherRaiseSparse(self, strategy, pure_eager): |
| dense_shape = [5, 2] |
| t0 = _make_indexed_slices( |
| values=[[1., 2.]], indices=[2], dense_shape=dense_shape) |
| |
| def run(value): |
| return strategy.gather(value, axis=0) |
| |
| with self.assertRaisesRegex( |
| NotImplementedError, |
| r'gather does not support IndexedSlices'): |
| if pure_eager: |
| run(t0) |
| else: |
| def_function.function(run)(t0) |
| |
| def testGatherRaiseDifferentRank(self, strategy, pure_eager): |
| """Different rank: [1,], [1, 2] -> raise error.""" |
| if strategy.num_replicas_in_sync <= 1: |
| self.skipTest('Test for more than 1 replicas.') |
| if isinstance(strategy, CollectiveAllReduceStrategy |
| ) and _get_num_replicas_per_client(strategy) > 1: |
| self.skipTest('b/167331966') |
| def value_fn(ctx): |
| return array_ops.ones(shape=(range(1, ctx.replica_id_in_sync_group + 2))) |
| |
| distributed_values = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| axis = 0 |
| |
| def run(): |
| return strategy.gather(distributed_values, axis=axis) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| if isinstance(strategy, CollectiveAllReduceStrategy): |
| with self.assertRaisesRegex(errors.InvalidArgumentError, |
| r'Shape mismatch'): |
| run() |
| elif isinstance( |
| strategy, |
| (mirrored_strategy.MirroredStrategy, |
| central_storage_strategy.CentralStorageStrategy)): |
| if pure_eager: |
| with self.assertRaises(errors.InvalidArgumentError) as e: |
| run() |
| # Different error message depending on whether collective ops is used. |
| self.assertRegexMatch( |
| str(e.exception), |
| ['Ranks of all input tensors should match', 'Shape mismatch']) |
| else: |
| with self.assertRaises((errors.InvalidArgumentError, ValueError)) as e: |
| run() |
| self.assertRegexMatch( |
| str(e.exception), |
| [r'Shape must be rank \d but is rank \d', 'Shape mismatch']) |
| elif _is_tpu_strategy(strategy) and pure_eager: |
| with self.assertRaisesRegex(ValueError, |
| r'Dimension \d in both shapes must be equal'): |
| run() |
| else: |
| with self.assertRaisesRegex(ValueError, |
| r'Shape must be rank \d but is rank \d'): |
| run() |
| |
| # Ideally, here we should split them into another test class, AllGatherTest. |
| # But doing that makes two initialize_tpu_system() calls and one of them times |
| # out, on Kokoro. Integrating two into one avoids it. |
| def _all_gather_same_shape_and_verify(self, value_on_replica, axis, |
| pure_eager, strategy): |
| per_replica_value = strategy.experimental_distribute_values_from_function( |
| lambda _: array_ops.identity(value_on_replica)) |
| |
| def replica_fn(per_replica_value): |
| ctx = distribute_lib.get_replica_context() |
| local_value = array_ops.identity(per_replica_value) |
| return ctx.all_gather(local_value, axis=axis) |
| |
| if not pure_eager: |
| replica_fn = def_function.function(replica_fn) |
| |
| result = strategy.experimental_local_results( |
| strategy.run(replica_fn, args=(per_replica_value,))) |
| |
| all_value = [value_on_replica for _ in range(strategy.num_replicas_in_sync)] |
| expect = array_ops.concat(all_value, axis=axis) |
| expected_result = [expect] * _get_num_replicas_per_client(strategy) |
| |
| self.assertAllClose(expected_result, result) |
| |
| def testAllGatherPerReplicaDense1D0Axis(self, strategy, pure_eager): |
| """all_gather(..., axis=0,...) a DistributedValues with a Tensor of shape (3,) on two replica returns a PerReplica of tensor(s) with shape (6,).""" |
| single_value = constant_op.constant([1, 2, 3], dtype=dtypes.float32) |
| axis = 0 |
| self._all_gather_same_shape_and_verify(single_value, axis, pure_eager, |
| strategy) |
| |
| def testAllGatherPerReplicaDense2D0Axis(self, strategy, pure_eager): |
| """all_gather(..., axis=0,...) a DistributedValues with a Tensor of shape (1,3) on two replica returns PerReplica of tensor(s) with shape (2,3).""" |
| single_value = constant_op.constant([[1, 2, 3]]) |
| axis = 0 |
| self._all_gather_same_shape_and_verify(single_value, axis, pure_eager, |
| strategy) |
| |
| def testAllGatherPerReplicaDense2D1Axis(self, strategy, pure_eager): |
| """all_gather(..., axis=1,...) a DistributedValues with a Tensor of shape (1,3) on two replica returns PerReplica of tensor(s) with shape (1,6).""" |
| single_value = constant_op.constant([[1, 2, 3]]) |
| axis = 1 |
| self._all_gather_same_shape_and_verify(single_value, axis, pure_eager, |
| strategy) |
| |
| def testAllGatherPerReplicaDense3D0Axis(self, strategy, pure_eager): |
| """all_gather(..., axis=0,...) a DistributedValues with a Tensor of shape (1,2,2) on two replica returns PerReplica of tensor(s) with shape (2,2,2).""" |
| single_value = constant_op.constant([[[1, 2], [1, 2]]]) |
| axis = 0 |
| self._all_gather_same_shape_and_verify(single_value, axis, pure_eager, |
| strategy) |
| |
| def testAllGatherPerReplicaDense3D1Axis(self, strategy, pure_eager): |
| """all_gather(..., axis=1,...) a DistributedValues with a Tensor of shape (1,2,2) on two replica returns PerReplica of tensor(s) with shape (1,4,2).""" |
| single_value = constant_op.constant([[[1, 2], [1, 2]]]) |
| axis = 1 |
| self._all_gather_same_shape_and_verify(single_value, axis, pure_eager, |
| strategy) |
| |
| def testAllGatherPerReplicaDense3D2Axis(self, strategy, pure_eager): |
| """all_gather(..., axis=2,...) a DistributedValues with a Tensor of shape (1,2,2) on two replica returns PerReplica of tensor(s) with shape (1,2,4).""" |
| single_value = constant_op.constant([[[1, 2], [1, 2]]]) |
| axis = 2 |
| self._all_gather_same_shape_and_verify(single_value, axis, pure_eager, |
| strategy) |
| |
| def testAllGatherDiffValueTPU(self, strategy, pure_eager): |
| # Test for TPU only since it can't be tested via testAllGatherDiffShape* |
| if not _is_tpu_strategy(strategy): |
| self.skipTest('Test for TPU only. For other strategies case already' |
| ' covered in other tests') |
| |
| data = [[1], [2], [3], [4], [5], [6], [7], [8]] |
| |
| axis = 0 |
| dataset = dataset_ops.DatasetV2.from_tensor_slices(data).batch(8) |
| input_iterator = iter(strategy.experimental_distribute_dataset(dataset)) |
| |
| @def_function.function |
| def replica_fn(per_replica_value): |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather(array_ops.identity(per_replica_value), axis=axis) |
| |
| result = strategy.experimental_local_results( |
| strategy.run(replica_fn, args=(next(input_iterator),))) |
| |
| expected_result = [data] * _get_num_replicas_per_client(strategy) |
| self.assertAllClose(expected_result, result) |
| |
| def testAllGatherDiffShapeAtAxis0(self, strategy, pure_eager): |
| """Different `Axis==0`-th dimension: shape [1, 1], [2, 1] -> [3, 1].""" |
| |
| if _is_tpu_strategy(strategy): |
| self.skipTest('TPU does not support all_gather different shapes') |
| |
| def value_fn(ctx): |
| return constant_op.constant( |
| 1, shape=(ctx.replica_id_in_sync_group + 1, 1)) |
| |
| per_replica_value = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| |
| expect = constant_op.constant( |
| 1, shape=(sum(range(strategy.num_replicas_in_sync + 1)), 1)) |
| |
| def run(value): |
| value_identity = array_ops.identity(value) |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather(value_identity, axis=0) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| expected_result = [expect] * _get_num_replicas_per_client(strategy) |
| result = strategy.experimental_local_results( |
| strategy.run(run, args=(per_replica_value,))) |
| self.assertAllEqual(expected_result, result) |
| |
| def testAllGatherDiffShapeAtAxis1(self, strategy, pure_eager): |
| """Different `Axis`-th (not 0th) dimension: shape [1, 1], [1, 2] -> [1, 3].""" |
| if _is_tpu_strategy(strategy): |
| self.skipTest('TPU does not support all_gather different shapes') |
| |
| def value_fn(ctx): |
| return constant_op.constant( |
| 1, shape=(1, ctx.replica_id_in_sync_group + 1)) |
| |
| per_replica_value = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| |
| expect = constant_op.constant( |
| 1, shape=(1, sum(range(strategy.num_replicas_in_sync + 1)))) |
| |
| def run(value): |
| value_identity = array_ops.identity(value) |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather(value_identity, axis=1) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| expected_result = [expect] * _get_num_replicas_per_client(strategy) |
| result = strategy.experimental_local_results( |
| strategy.run(run, args=(per_replica_value,))) |
| self.assertAllEqual(expected_result, result) |
| |
| def testAllGatherNest(self, strategy, pure_eager): |
| if _is_tpu_strategy(strategy): |
| self.skipTest('TPU does not support all_gather different shapes') |
| |
| axis = 1 |
| |
| def value_fn(ctx): |
| value = constant_op.constant( |
| 1, shape=(1, ctx.replica_id_in_sync_group + 1)) |
| return value |
| per_replica_value = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| |
| expect_1 = constant_op.constant( |
| 1, shape=(1, sum(range(strategy.num_replicas_in_sync + 1)))) |
| |
| expected_per_replica_1 = [expect_1] * _get_num_replicas_per_client(strategy) |
| |
| value_2 = constant_op.constant([[[1, 2], [1, 2]]]) |
| |
| expect_2 = array_ops.concat( |
| [value_2 for _ in range(strategy.num_replicas_in_sync)], axis=axis) |
| |
| expected_per_replica_2 = [expect_2] * _get_num_replicas_per_client(strategy) |
| |
| def run(value): |
| value_1 = array_ops.identity(value) |
| value_3 = array_ops.identity(value_2) |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather([value_1, value_3], axis=axis) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| result = strategy.run(run, args=(per_replica_value,)) |
| self.assertAllEqual(expected_per_replica_1, |
| strategy.experimental_local_results(result[0])) |
| self.assertAllEqual(expected_per_replica_2, |
| strategy.experimental_local_results(result[1])) |
| |
| def testAllGatherNest1D0Axis(self, strategy, pure_eager): |
| """all_gather(..., axis=0,...) a nest of DistributedValues.""" |
| single_value = constant_op.constant([1, 2, 3]) |
| axis = 0 |
| |
| def run(): |
| value_identity = array_ops.identity(single_value) |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather([value_identity, value_identity], axis=axis) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| all_value = [single_value for _ in range(strategy.num_replicas_in_sync)] |
| expect = array_ops.concat(all_value, axis=axis) |
| expected_per_replica = [expect] * _get_num_replicas_per_client(strategy) |
| |
| result = strategy.run(run) |
| for gathered_result in result: |
| self.assertAllEqual(expected_per_replica, |
| strategy.experimental_local_results(gathered_result)) |
| |
| def testAllGatherRaiseDiffShapeAtNonAxis(self, strategy, pure_eager): |
| """Different at non-`axis`-th dimension : [2, 1], [1, 1], all_gather(...axis=1...) -> raise error.""" |
| if _is_tpu_strategy(strategy): |
| self.skipTest('TODO(b/169108777): raise a clear error message in xla.') |
| |
| if isinstance(strategy, CollectiveAllReduceStrategy |
| ) and _get_num_replicas_per_client(strategy) > 1: |
| self.skipTest('b/167331966') |
| |
| if strategy.num_replicas_in_sync <= 1: |
| self.skipTest('Test for more than 1 replica only.') |
| |
| def value_fn(ctx): |
| return constant_op.constant( |
| 1, shape=(1, ctx.replica_id_in_sync_group + 1)) |
| |
| per_replica_value = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| |
| def run(value): |
| value_identity = array_ops.identity(value) |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather(value_identity, axis=0) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| if isinstance(strategy, CollectiveAllReduceStrategy): |
| with self.assertRaisesRegex(errors.InvalidArgumentError, |
| r'Shape mismatch'): |
| strategy.run(run, args=(per_replica_value,)) |
| elif isinstance(strategy, |
| (mirrored_strategy.MirroredStrategy, |
| central_storage_strategy.CentralStorageStrategy)): |
| with self.assertRaisesRegex((errors.InvalidArgumentError, ValueError), |
| r'Dimension \d in both shapes must be equal'): |
| strategy.run(run, args=(per_replica_value,)) |
| |
| def testAllGatherRaiseSparse(self, strategy, pure_eager): |
| dense_shape = [5, 2] |
| t0 = _make_indexed_slices( |
| values=[[1., 2.]], indices=[2], dense_shape=dense_shape) |
| |
| def replica_fn(value): |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather(value, axis=0) |
| |
| with self.assertRaisesRegex( |
| NotImplementedError, |
| r'all_gather does not support IndexedSlices'): |
| if not pure_eager: |
| strategy.run(def_function.function(replica_fn), args=(t0,)) |
| else: |
| strategy.run(replica_fn, args=(t0,)) |
| |
| def testAllGatherRaiseDifferentRank(self, strategy, pure_eager): |
| """Different rank: [1,], [1, 2] -> raise error.""" |
| if _is_tpu_strategy(strategy): |
| self.skipTest('TODO(b/169108777): raise a clear error message in xla.') |
| |
| if strategy.num_replicas_in_sync <= 1: |
| self.skipTest('Test for more than 1 replicas.') |
| if isinstance(strategy, CollectiveAllReduceStrategy |
| ) and _get_num_replicas_per_client(strategy) > 1: |
| self.skipTest('b/167331966') |
| def value_fn(ctx): |
| return array_ops.ones(shape=(range(1, ctx.replica_id_in_sync_group + 2))) |
| |
| per_replica_value = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| |
| def run(value): |
| value_identity = array_ops.identity(value) |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather(value_identity, axis=0) |
| |
| if not pure_eager: |
| run = def_function.function(run) |
| |
| if isinstance(strategy, CollectiveAllReduceStrategy): |
| with self.assertRaisesRegex(errors.InvalidArgumentError, |
| r'Shape mismatch'): |
| strategy.run(run, args=(per_replica_value,)) |
| elif isinstance(strategy, |
| (mirrored_strategy.MirroredStrategy, |
| central_storage_strategy.CentralStorageStrategy)): |
| if pure_eager: |
| with self.assertRaises(errors.InvalidArgumentError) as e: |
| strategy.run(run, args=(per_replica_value,)) |
| # Different error message depending on whether collective ops is used. |
| self.assertRegexMatch( |
| str(e.exception), |
| ['Ranks of all input tensors should match', 'Shape mismatch']) |
| else: |
| with self.assertRaises((errors.InvalidArgumentError, ValueError)) as e: |
| strategy.run(run, args=(per_replica_value,)) |
| self.assertRegexMatch( |
| str(e.exception), |
| [r'Shape must be rank \d but is rank \d', 'Shape mismatch']) |
| else: |
| with self.assertRaisesRegex(ValueError, |
| r'Dimension \d in both shapes must be equal'): |
| strategy.run(run, args=(per_replica_value,)) |
| |
| def testAllGatherGradient(self, strategy, pure_eager): |
| if pure_eager: |
| self.skipTest('`tf.gradients` is not supported with eager execution ' |
| 'without using tf.functions.') |
| |
| def all_gather_fn(value): |
| axis = 1 |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather(array_ops.identity(value), axis) |
| |
| gradient_comp = sum(range(1, strategy.num_replicas_in_sync + 1)) |
| gradient = [[gradient_comp], [gradient_comp]] |
| grads_for_all_replicas = [gradient] * _get_num_replicas_per_client(strategy) |
| |
| @def_function.function |
| def step(c): |
| x = constant_op.constant([[3.], [5.]]) |
| mid = all_gather_fn(x) |
| y = mid * c |
| return gradients_impl.gradients_v2(y, [x])[0] |
| |
| def value_fn(ctx): |
| x = [1., 2., 3., 4., 5., 6., 7., 8.] |
| return array_ops.constant([x[ctx.replica_id_in_sync_group]]) |
| |
| per_replica_value = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| result = strategy.experimental_local_results( |
| strategy.run(step, args=(per_replica_value,))) |
| |
| self.assertAllEqual(grads_for_all_replicas, result) |
| |
| def testAllGatherGradientNest(self, strategy, pure_eager): |
| if pure_eager: |
| self.skipTest('`tf.gradients` is not supported with eager execution ' |
| 'without using tf.functions.') |
| |
| def all_gather_fn(value): |
| axis = 1 |
| ctx = distribute_lib.get_replica_context() |
| return ctx.all_gather(array_ops.identity(value), axis) |
| |
| gradient_comp = sum(range(1, strategy.num_replicas_in_sync + 1)) |
| gradient = [[gradient_comp], [gradient_comp]] |
| grads_for_all_replicas = [gradient] * _get_num_replicas_per_client(strategy) |
| |
| @def_function.function |
| def step(c): |
| x = constant_op.constant([[3.], [5.]]) |
| y = constant_op.constant([[2.], [4.]]) |
| mid = all_gather_fn([x, y]) |
| y = mid * c |
| return gradients_impl.gradients_v2(y, [x])[0] |
| |
| def value_fn(ctx): |
| x = [1., 2., 3., 4., 5., 6., 7., 8.] |
| return array_ops.constant([x[ctx.replica_id_in_sync_group]]) |
| |
| per_replica_value = strategy.experimental_distribute_values_from_function( |
| value_fn) |
| result = strategy.experimental_local_results( |
| strategy.run(step, args=(per_replica_value,))) |
| |
| self.assertAllEqual(grads_for_all_replicas, result) |
| |
| |
| def _make_indexed_slices(values, indices, dense_shape): |
| tensor = indexed_slices.IndexedSlices( |
| values=constant_op.constant(values), |
| indices=constant_op.constant(indices), |
| dense_shape=constant_op.constant(dense_shape)) |
| return tensor |
| |
| |
| def _get_num_replicas_per_client(strategy): |
| if isinstance(strategy, CollectiveAllReduceStrategy): |
| resolver = strategy.cluster_resolver |
| return max(nest.flatten(resolver.num_accelerators())[0], 1) |
| else: |
| return strategy.num_replicas_in_sync |
| |
| |
| def _is_tpu_strategy(strategy): |
| return isinstance(strategy, |
| (tpu_strategy.TPUStrategy, tpu_strategy.TPUStrategyV1, |
| tpu_strategy.TPUStrategyV2)) |
| |
| |
| if __name__ == '__main__': |
| test_util.main() |
