From ab61e1eccc55f269d0e9c8f17aea46800b939304 Mon Sep 17 00:00:00 2001
From: David Majnemer <majnemer@google.com>
Date: Thu, 21 Jan 2021 10:26:03 -0800
Subject: [PATCH] [XLA] Fix test yet again

PiperOrigin-RevId: 353048445
Change-Id: Id1a35b23c65c4f08665eda9c6af6025a6b4500f3
---
 tensorflow/compiler/xla/tests/BUILD           | 22 ++++++++++---------
 .../compiler/xla/tests/collective_ops_test.cc |  6 ++---
 2 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/tensorflow/compiler/xla/tests/BUILD b/tensorflow/compiler/xla/tests/BUILD
index 73425353334..6ed4d4b7fb1 100644
--- a/tensorflow/compiler/xla/tests/BUILD
+++ b/tensorflow/compiler/xla/tests/BUILD
@@ -1989,19 +1989,21 @@ xla_test(
     name = "collective_ops_test",
     srcs = ["collective_ops_test.cc"],
     args = ["--xla_force_host_platform_device_count=4"],
-    backends = [
-        "gpu",
-        "cpu",
-    ],
-    tags = [
+    backend_tags = {
         # This test is tagged "manual" because it requires multiple GPUs, and
         # Forge only supports single-GPU tests.  Guitar skips "manual" tests
         # unless they're also tagged "guitar".
-        "guitar",
-        "manual",
-        "multi_gpu",
-        "no_oss",
-        "notap",
+        "gpu": [
+            "guitar",
+            "manual",
+            "multi_gpu",
+            "no_oss",
+            "notap",
+        ],
+    },
+    backends = [
+        "gpu",
+        "cpu",
     ],
     deps = [
         ":client_library_test_base",
diff --git a/tensorflow/compiler/xla/tests/collective_ops_test.cc b/tensorflow/compiler/xla/tests/collective_ops_test.cc
index a88b98a5c7d..ceabd3c62fb 100644
--- a/tensorflow/compiler/xla/tests/collective_ops_test.cc
+++ b/tensorflow/compiler/xla/tests/collective_ops_test.cc
@@ -765,8 +765,8 @@ XLA_TEST_F(CollectiveOpsTest, AllGather_Dim0) {
                                            result);
   }
 }
-// TODO(b/178047150): Fails on GPU with wrong answers.
-XLA_TEST_F(CollectiveOpsTest, DISABLED_ON_GPU(AllGather_Dim1)) {
+
+XLA_TEST_F(CollectiveOpsTest, AllGather_Dim1) {
   const char* const kModuleStr = R"(
   HloModule test
   ENTRY test_computation {
@@ -789,7 +789,7 @@ XLA_TEST_F(CollectiveOpsTest, DISABLED_ON_GPU(AllGather_Dim1)) {
                         /*use_threads=*/true, /*run_hlo_passes=*/true));
   ASSERT_EQ(results.size(), kNumReplicas);
   for (const Literal& result : results) {
-    LiteralTestUtil::ExpectR1Equal<uint32>({10, 15, 11, 16, 12, 17, 13, 18},
+    LiteralTestUtil::ExpectR1Equal<uint32>({10, 11, 12, 13, 15, 16, 17, 18},
                                            result);
   }
 }