Searched refs:SliceFirstK (Results 1 – 3 of 3) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
| D | spmd_partitioner_util.h | 277 HloInstruction* SliceFirstK(HloInstruction* hlo, SpmdBuilder* builder,
|
| D | spmd_partitioner.cc | 1705 auto slice_input = SliceFirstK(value_gte, &b_, sort_dim, k.value()); in HandleSort()
1715 auto slice_index = SliceFirstK(index_gte, &b_, sort_dim, k.value()); in HandleSort()
1931 SliceFirstK(sort_value_gte, &b_, sort_dim, k); in HandleCustomCall()
1934 SliceFirstK(sort_index_gte, &b_, sort_dim, k); in HandleCustomCall()
|
| D | spmd_partitioner_util.cc | 1261 HloInstruction* SliceFirstK(HloInstruction* hlo, SpmdBuilder* builder, in SliceFirstK() function
|