Searched refs:sorted_input_inds (Results 1 – 1 of 1) sorted by relevance
| /aosp_15_r20/external/tensorflow/tensorflow/core/kernels/ |
| H A D | unique_op_gpu.cu.h | 61 const TIndex* __restrict__ sorted_input_inds, in ExtractFirstOccurrenceIndicesKernel() argument
67 unique_input_inds[sorted_input_unique_id] = sorted_input_inds[i]; in ExtractFirstOccurrenceIndicesKernel()
88 const TIndex* sorted_input_inds, in ExtractFirstOccurrenceIndices() argument
98 d.stream(), input_size, uniq_size, sorted_input_inds, in ExtractFirstOccurrenceIndices()
145 int64_t input_size, const TIndex* sorted_input_inds, in LookupAndScatterUniqueIdsKernel() argument
150 idx[sorted_input_inds[i]] = in LookupAndScatterUniqueIdsKernel()
159 const TIndex* sorted_input_inds, in LookupAndScatterUniqueIds() argument
169 d.stream(), input_size, sorted_input_inds, in LookupAndScatterUniqueIds()
272 Tensor sorted_input_inds; in ComputeAsync() local
274 AllocateTemp(context, input_size, &sorted_input_inds, in ComputeAsync()
[all …]
|