Home
last modified time | relevance | path

Searched defs:segment_offsets (Results 1 – 7 of 7) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DEmbeddingBackwardKernel.cu47 void krn_partials_per_segment(index_t *ret, const index_t *segment_offsets, in krn_partials_per_segment()
65 const index_t *segment_offsets, in krn_partial_segment_offset()
86 const index_t* segment_offsets, const int64_t *num_of_segments_ptr, in compute_grad_weight_bags()
130 const index_t* segment_offsets, in compute_grad_weight()
162 const index_t* segment_offsets, const int64_t *num_of_segments_ptr, in sum_and_scatter()
234 auto segment_offsets = at::empty({numel}, orig_indices.options()); in embedding_backward_cuda_kernel() local
H A DLegacyThrustHelpers.cu89 …bedding_backward_cuda_kernel_unique_by_key(const Tensor &sorted_indices, Tensor &segment_offsets) { in embedding_backward_cuda_kernel_unique_by_key()
/aosp_15_r20/external/tensorflow/tensorflow/core/kernels/
H A Dsegment_reduction_ops_gpu.cu.h144 const Index* __restrict__ segment_offsets, // [nsegments + 1] in SegmentMeanNormalizeKernel()
159 const Index* __restrict__ segment_offsets, // [nsegments + 1] in LaunchSegmentMeanNormalizeKernel()
197 Tindex* __restrict__ segment_offsets) { // [nsegments + 1] in SegmentOffsetsKernel()
227 Tindex* segment_offsets) { // [nsegments + 1] in LaunchSegmentOffsetsKernel()
H A Dgpu_prim_helpers.h216 OffsetIteratorT segment_offsets, // [num_segments + 1] in GpuSegmentedReduce()
H A Dgpu_prim_helpers_test.cu.cc141 const Tensor& segment_offsets = context->input(1); in Compute() local
/aosp_15_r20/external/executorch/runtime/executor/
H A Dprogram.cpp447 const auto& segment_offsets = internal_program_->mutable_data_segments()->Get( in load_mutable_subsegment_into() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/sparse/cuda/
H A DSparseCUDAApplyUtils.cuh247 int64_t *segment_offsets, int64_t *value_indices, in coalesceValuesKernel()