commit 09984bd3a044d9f8cf9539b66c8d0dd3bd0bd19e
author Rick Chao <rchao@google.com> Thu Mar 18 11:55:31 2021 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Thu Mar 18 12:01:53 2021 -0700
tree 3f3bba7b5b82d1d3fa5f1d1816606ce3958e8c7a
parent 631f2b77dc22cf2e1b9a41d0c7518b32fe02e61b

PSv2: Make multi-GPU tests compatible with guitar tests where GPU count is > 2.

PiperOrigin-RevId: 363710173
Change-Id: Ida8e3118f5cd097fe07bed375b2fc4e8bc7804ea

tensorflow/python/distribute/coordinator/cluster_coordinator_test.py

diff --git a/tensorflow/python/distribute/coordinator/cluster_coordinator_test.py b/tensorflow/python/distribute/coordinator/cluster_coordinator_test.py
index deeec98..5f6ce2e 100644
--- a/tensorflow/python/distribute/coordinator/cluster_coordinator_test.py
+++ b/tensorflow/python/distribute/coordinator/cluster_coordinator_test.py
@@ -914,7 +914,7 @@
       # Invoking `run` without `coordinator.schedule` again should error.
       self.strategy.run(replica_fn)
 
-    all_results = [(2, 0), (2, 0)]
+    all_results = [(2, 0)] * self.strategy.num_replicas_in_sync
     expected_result = []
     for i in range(self.strategy.num_replicas_in_sync):
       expected_result.append(all_results[i])
@@ -1067,7 +1067,7 @@
   def testAsyncScheduleWithDistributedDataset(self):
 
     def input_fn():
-      dataset = dataset_ops.DatasetV2.from_tensor_slices([2.] * 24).batch(
+      dataset = dataset_ops.DatasetV2.from_tensor_slices([2.]).repeat().batch(
           self.strategy.num_replicas_in_sync)
       return self.strategy.experimental_distribute_dataset(dataset)
 
@@ -1122,7 +1122,8 @@
         self._map_fn_tracing_count += 1
         return x + 10
 
-      dataset = dataset_ops.DatasetV2.range(0, 10).batch(2).map(map_fn)
+      dataset = dataset_ops.DatasetV2.range(0, 10).batch(
+          self.strategy.num_replicas_in_sync).map(map_fn)
       return self.strategy.experimental_distribute_dataset(dataset)
 
     @def_function.function
@@ -1134,7 +1135,7 @@
         worker_fn, args=(iter(distributed_dataset),))
 
     expected_result = array_ops.split(
-        math_ops.range(10., 12.),
+        math_ops.range(10., 10. + self.strategy.num_replicas_in_sync),
         num_or_size_splits=self.strategy.num_replicas_in_sync,
         axis=0)