Searched refs:iota_s32 (Results 1 – 1 of 1) sorted by relevance
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | sorting.cc | 50 XlaOp iota_s32 = Iota(builder, iota_shape, last_dim); in TopK() local
54 iota_s32 = SetDimensionSize(iota_s32, GetDimensionSize(input, i), i); in TopK()
59 Sort({input, iota_s32}, in TopK()
61 iota_s32.builder()), in TopK()
92 XlaOp iota_s32 = Iota(builder, iota_shape, last_dim); in TopKWithPartitions() local
96 iota_s32 = SetDimensionSize(iota_s32, GetDimensionSize(input, i), i); in TopKWithPartitions()
106 auto iota_s32 = values_and_indices[3]; in TopKWithPartitions() local
114 DynamicSliceInMinorDims(iota_s32, {start}, {per_partition_size}); in TopKWithPartitions()
136 return std::vector<XlaOp>{values, indices, input, iota_s32}; in TopKWithPartitions()
149 Slice(iota_s32, start_indices, limit_indices, strides); in TopKWithPartitions()
[all …]
|