Searched defs:indices_ptr (Results 1 – 15 of 15) sorted by relevance
110 int64_t* indices_ptr) { in segmented_sort_large_segments()139 int64_t* const indices_ptr) { in segmented_sort_pairs_by_full_sort()190 int64_t* indices_ptr) { in segmented_sort_pairs()234 int64_t* indices_ptr = indices.mutable_data_ptr<int64_t>(); in launch_stable_sort_kernel() local
29 const int64_t* indices_ptr = indices.const_data_ptr<int64_t>(); in embedding_out() local
154 int64_t* const indices_ptr, in kernel_reduction_then_map_2d()347 int64_t* indices_ptr = nullptr; variable
42 const auto *indices_ptr = static_cast<const uint32_t *>(indices.data()); in gather() local
38 const CTYPE* indices_ptr = indices.const_data_ptr<CTYPE>(); in embedding_kernel() local
68 uint32_t *indices_ptr = static_cast<uint32_t *>(indices.data()); in generate_indices() local
152 int32* indices_ptr = indices_weights->GetPointer<int32>(); in Convert() local
44 int64_t* indices_ptr = indices_data + c * output_height * output_width; in cpu_adaptive_max_pool2d() local367 const int64_t* indices_ptr = indices_data + c * output_height * output_width; in cpu_adaptive_max_pool2d_backward() local415 const int64_t* indices_ptr = indices_data + n * output_height * output_width * channels; in cpu_adaptive_max_pool2d_backward_channels_last() local512 int64_t* indices_ptr = indices_data + c * output_depth * output_height * output_width; in cpu_adaptive_max_pool3d() local859 int64_t* indices_ptr = indices_data + c * output_depth * output_height * output_width; in cpu_adaptive_max_pool3d_backward() local911 … int64_t* indices_ptr = indices_data + n * output_depth * output_height * output_width * channels; in cpu_adaptive_max_pool3d_backward_channels_last() local
131 const int64_t* indices_ptr = indices_data + i * channels; in cpu_max_unpool_channels_last() local
298 int64_t* indices_ptr = indices_data + c * output_depth * output_height * output_width; in cpu_max_pool() local514 const int64_t * indices_ptr = indices_data + c * output_depth * output_height * output_width; in cpu_max_pool_backward() local574 …const int64_t* indices_ptr = indices_data + n * output_depth * output_height * output_width * chan… in cpu_max_pool_backward_channels_last() local
149 const Index* indices_ptr = indices.flat<Index>().data(); in operator ()() local
303 const int64_t* indices_ptr = indices.data(); in operator ()() local
880 se::DeviceMemoryBase indices_ptr( in DoScatterNdOnCpu() local
146 const int64_t* indices_ptr = indices.const_data_ptr<int64_t>(); in embedding_byte_per_channel() local
186 const int64_t* indices_ptr = indices.const_data_ptr<int64_t>(); in embedding_xbit_per_channel() local