[XLA] Fix test yet again
PiperOrigin-RevId: 353048445 Change-Id: Id1a35b23c65c4f08665eda9c6af6025a6b4500f3
This commit is contained in:
parent
87e111077f
commit
ab61e1eccc
@ -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",
|
||||
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user