diff options
author | 2018-08-08 17:10:33 -0700 | |
---|---|---|
committer | 2018-08-08 17:15:13 -0700 | |
commit | 3e02edc1f33fb3bfa43b5828d8ecea0dbc7738ea (patch) | |
tree | 926591cdb6e0b62c2dc8c041b1c61cbaf3c2ab85 /tensorflow/compiler/xla/service/hlo_parser_test.cc | |
parent | aacb29a4ab88f9fa27c3301977e7f2cc289a3976 (diff) |
[XLA] Add the xla interface for AllToAll.
PiperOrigin-RevId: 207971529
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_parser_test.cc')
-rw-r--r-- | tensorflow/compiler/xla/service/hlo_parser_test.cc | 24 |
1 files changed, 24 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_parser_test.cc b/tensorflow/compiler/xla/service/hlo_parser_test.cc index 7344679bb6..4cd21841f4 100644 --- a/tensorflow/compiler/xla/service/hlo_parser_test.cc +++ b/tensorflow/compiler/xla/service/hlo_parser_test.cc @@ -1072,6 +1072,30 @@ ENTRY CrossReplicaSumWithSubgroups { )" }, +// all-to-all +{ +"AllToAll", +R"(HloModule AllToAll + +ENTRY AllToAll { + input = f32[128,32]{0,1} parameter(0) + ROOT a2a = f32[128,32]{0,1} all-to-all(input), replica_groups={} +} + +)" +}, +// all-to-all with subgroups +{ +"AllToAllWithSubgroups", +R"(HloModule AllToAllWithSubgroups + +ENTRY AllToAllWithSubgroups { + input = f32[128,32]{0,1} parameter(0) + ROOT a2a = f32[128,32]{0,1} all-to-all(input), replica_groups={{1,2},{3,0}}, barrier="abc" +} + +)" +}, // Iota { "Iota", |