Home
last modified time | relevance | path

Searched defs:offset2bag (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/
H A DEmbeddingBag.cpp57 static void make_offset2bag(const Tensor &offsets, Tensor& offset2bag) { in make_offset2bag()
966 Tensor& offset2bag, in make_offset2bag_out()
1030 Tensor offset2bag = at::empty({0}, offsets.options()); in make_offset2bag() local
1052 const Tensor &offset2bag, in apply_bag_size_backward()
1068 const Tensor& offset2bag, in embedding_bag_cpu_max_out()
1128 void _embedding_bag_cpu_impl_out(Tensor& output, Tensor& offset2bag, in _embedding_bag_cpu_impl_out()
1208 …Tensor offset2bag = make_offset2bag(output, weight, indices, offsets, mode, per_sample_weights, pa… in _embedding_bag_cpu_impl() local
1316 at::Tensor& offset2bag, in _embedding_bag_cpu_out()
1376 const Tensor &offset2bag, in _embedding_bag_backward()
1391 const Tensor &offset2bag, in _embedding_bag_backward_symint()
[all …]
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DEmbeddingBag.cu64 index_t *offset2bag, int64_t numIndices, int64_t numBags, in EmbeddingBag_updateOutputKernel_max()
114 index_t *offset2bag, int64_t numIndices, int64_t numBags, in EmbeddingBag_updateOutputKernel_sum_mean()
171 const Tensor &offset2bag, in embedding_bag_backward_cuda_sum_avg()
377 auto offset2bag = in _embedding_bag_cuda() local
428 const Tensor &offset2bag, in _embedding_bag_dense_backward_cuda()
475 const index_t* offset2bag, // contiguous in _embedding_bag_per_sample_weights_backward_kernel()
512 const Tensor& offset2bag, in _embedding_bag_per_sample_weights_backward_cuda()
H A DEmbeddingBackwardKernel.cu83 const index_t *offset2bag, const index_t *count, ptrdiff_t numel, in compute_grad_weight_bags()
220 const Tensor &offset2bag, in embedding_backward_cuda_kernel()
/aosp_15_r20/external/pytorch/torch/
H A D_meta_registrations.py6353 offset2bag, argument
6395 offset2bag, argument
6421 offset2bag, argument