Home
last modified time | relevance | path

Searched refs:SetPartitionedHlo (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.cc1401 SetPartitionedHlo(hlo, in DefaultAction()
1478 SetPartitionedHlo(hlo, [&] { in HandleElementwise()
1499 SetPartitionedHlo(hlo, [&] { in HandleConcatenate()
1566 SetPartitionedHlo(hlo, [&] { in HandleConcatenate()
1625 SetPartitionedHlo(hlo, [&] { in HandleSlice()
1738 SetPartitionedHlo(hlo, replicated_sort.Reshard(hlo->sharding())); in HandleSort()
1768 SetPartitionedHlo(hlo, [&] { in HandleSort()
1788 SetPartitionedHlo(hlo, [&] { return copy; }); in HandleCustomCall()
1799 SetPartitionedHlo(hlo, [&] { return copy; }); in HandleCustomCall()
1939 SetPartitionedHlo(hlo, PartitionedHlo(create_tuple, create_tuple->shape(), in HandleCustomCall()
[all …]
Dgather_scatter_handler.cc609 SetPartitionedHlo(hlo, [&]() { in HandleScatter()
628 SetPartitionedHlo(hlo, [&]() { in HandleScatter()
660 SetPartitionedHlo(hlo, [&]() { in HandleScatter()
689 SetPartitionedHlo(gather, [pgather] { return pgather; }); in HandleGather()
699 SetPartitionedHlo(gather, [pgather] { return pgather; }); in HandleGather()
Dspmd_partitioner.h461 void SetPartitionedHlo(const HloInstruction* hlo, in SetPartitionedHlo() function
470 void SetPartitionedHlo(const HloInstruction* hlo, in SetPartitionedHlo() function
475 SetPartitionedHlo( in SetPartitionedHlo()
Dfft_handler.cc431 SetPartitionedHlo(hlo, partitioned_fft); in HandleFft()
Ddot_handler.cc2852 SetPartitionedHlo(hlo, [&] { return partitioned_dot; }); in HandleDotHelper()