Home
last modified time | relevance | path

Searched refs:PartitionComputation (Results 1 – 3 of 3) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_test.cc38 StatusOr<std::unique_ptr<HloModule>> PartitionComputation( in PartitionComputation() function in xla::spmd::__anon4eee67ad0111::SpmdPartitioningTest
80 auto module_status = PartitionComputation(hlo_string, /*num_devices=*/4); in TEST_F()
97 PartitionComputation(hlo_string, /*num_devices=*/2)); in TEST_F()
116 PartitionComputation(hlo_string, /*num_devices=*/2)); in TEST_F()
136 PartitionComputation(hlo_string, /*num_devices=*/2)); in TEST_F()
161 PartitionComputation(hlo_string, /*num_devices=*/2)); in TEST_F()
183 PartitionComputation(hlo_string, /*num_devices=*/2)); in TEST_F()
204 PartitionComputation(hlo_string, /*num_devices=*/2)); in TEST_F()
224 PartitionComputation(hlo_string, /*num_devices=*/2)); in TEST_F()
249 PartitionComputation(hlo_string, /*num_devices=*/2)); in TEST_F()
[all …]
Dspmd_partitioner.h198 StatusOr<bool> PartitionComputation(HloComputation* computation,
Dspmd_partitioner.cc2846 ->PartitionComputation(hlo->while_condition(), in HandleWhile()
2851 ->PartitionComputation(hlo->while_body(), sharding, in HandleWhile()
2880 ->PartitionComputation(computation, hlo->sharding(), in HandleConditional()
3632 StatusOr<bool> SpmdPartitioner::PartitionComputation( in PartitionComputation() function in xla::spmd::SpmdPartitioner
3680 PartitionComputation(module->entry_computation(), root_sharding, in Run()