[XLA] Fix test yet again

PiperOrigin-RevId: 353048445
Change-Id: Id1a35b23c65c4f08665eda9c6af6025a6b4500f3
This commit is contained in:
David Majnemer 2021-01-21 10:26:03 -08:00 committed by TensorFlower Gardener
parent 87e111077f
commit ab61e1eccc
2 changed files with 15 additions and 13 deletions

View File

@ -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",

View File

@ -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);
}
}