| /aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
| H A D | EmbeddingBag.cpp | 117 Tensor& bag_size, in index_select_add() argument
126 if (bag_size.defined()) { in index_select_add()
127 bag_size_data = bag_size.data_ptr<index_t>(); in index_select_add()
150 // Decrement bag_size to reflect that the index is padded in index_select_add()
198 Tensor& bag_size, in index_select_add() argument
313 if (bag_size.defined()) { in index_select_add()
314 bag_size_data = bag_size.data_ptr<index_t>(); in index_select_add()
354 // Decrement bag_size to reflect that the index is padded in index_select_add()
376 Tensor &bag_size, in index_select_add() argument
458 if (bag_size.defined()) { in index_select_add()
[all …]
|
| H A D | EmbeddingBag.h | 46 const Tensor& bag_size,
129 Tensor& bag_size, Tensor* max_indices,
140 at::Tensor& bag_size,
|
| H A D | NonSymbolicBC.h | 19 …at::Tensor & offsets, const at::Tensor & offset2bag, const at::Tensor & bag_size, const at::Tensor…
20 …at::Tensor & offsets, const at::Tensor & offset2bag, const at::Tensor & bag_size, int64_t num_weig…
|
| H A D | native_functions.yaml | 2339 …(Tensor grad, Tensor indices, Tensor offsets, Tensor offset2bag, Tensor bag_size, Tensor maximum_i…
2343 …(Tensor grad, Tensor indices, Tensor offsets, Tensor offset2bag, Tensor bag_size, SymInt num_weigh…
2347 …g_dense_backward(Tensor grad, Tensor indices, Tensor offset2bag, Tensor bag_size, Tensor maximum_i…
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/ |
| H A D | EmbeddingBag.cu | 66 index_t *bag_size, index_t *max_indices, in EmbeddingBag_updateOutputKernel_max() argument
102 bag_size[bag] = bag_size_; in EmbeddingBag_updateOutputKernel_max()
116 int mode, index_t *bag_size, in EmbeddingBag_updateOutputKernel_sum_mean() argument
162 bag_size[bag] = bag_size_; in EmbeddingBag_updateOutputKernel_sum_mean()
172 const Tensor &bag_size, in embedding_bag_backward_cuda_sum_avg() argument
236 bag_size, per_sample_weights); in embedding_bag_backward_cuda_sum_avg()
376 auto bag_size = at::empty(offsets.sizes(), indices.options()); in _embedding_bag_cuda() local
406 weight.stride(0), weight.stride(1), bag_size.mutable_data_ptr<index_t>(), in _embedding_bag_cuda()
415 weight.stride(0), weight.stride(1), mode, bag_size.mutable_data_ptr<index_t>(), in _embedding_bag_cuda()
424 return std::tuple<Tensor, Tensor, Tensor, Tensor>(output, offset2bag, bag_size, max_indices); in _embedding_bag_cuda()
|
| H A D | EmbeddingBackwardKernel.cu | 84 int64_t stride, int mode_mean, const index_t *bag_size, in compute_grad_weight_bags() argument
116 gradient /= bag_size[seq_number]; in compute_grad_weight_bags()
221 const Tensor &bag_size, in embedding_backward_cuda_kernel() argument
326 mode_mean, bag_size.const_data_ptr<index_t>(), in embedding_backward_cuda_kernel()
|
| H A D | EmbeddingBackwardKernel.cuh | 19 const Tensor &bag_size = Tensor(),
|
| /aosp_15_r20/external/pytorch/test/nn/ |
| H A D | test_embedding.py | 511 bag_size = max_indices_per_bag
513 bag_size = 0
515 bag_size = random.randint(1, max_indices_per_bag - 1)
518 for _ in range(bag_size)
520 cur_offset += bag_size
|
| /aosp_15_r20/external/pytorch/torch/ |
| H A D | _meta_registrations.py | 3426 bag_size = indices.new_empty(offsets.size())
3437 bag_size = offsets.new_empty(num_bags)
3449 max_indices = offsets.new_empty(bag_size.size())
3450 return output, offset2bag, bag_size, max_indices
3455 output, offset2bag, bag_size, max_indices = meta_embedding_bag(
3459 bag_size = offsets.new_empty(offsets.size())
3460 return output, offset2bag, bag_size, max_indices
6354 bag_size, argument
6369 bag_size,
6381 bag_size,
[all …]
|
| /aosp_15_r20/external/pytorch/test/ |
| H A D | test_meta.py | 1622 output, offset2bag, bag_size, maximum_indices = torch.ops.aten._embedding_bag.default(
1629 grad, indices, offset2bag, bag_size, maximum_indices, weight.size(0),
1633 grad.to('meta'), indices.to('meta'), offset2bag.to('meta'), bag_size.to('meta'),
1650 output, offset2bag, bag_size, maximum_indices = torch.ops.aten._embedding_bag.default(
|
| /aosp_15_r20/external/pytorch/torch/csrc/inductor/aoti_torch/generated/ |
| H A D | c_shim_cpu.h | 23 …AtenTensorHandle indices, AtenTensorHandle offset2bag, AtenTensorHandle bag_size, AtenTensorHandle…
|
| H A D | c_shim_cuda.h | 26 …AtenTensorHandle indices, AtenTensorHandle offset2bag, AtenTensorHandle bag_size, AtenTensorHandle…
|
| /aosp_15_r20/external/pytorch/torch/onnx/ |
| H A D | symbolic_opset10.py | 664 … # aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices.
|
| H A D | symbolic_helper.py | 2041 # aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices.
|
| /aosp_15_r20/external/pytorch/tools/autograd/ |
| H A D | derivatives.yaml | 1996 …g_dense_backward(Tensor grad, Tensor indices, Tensor offset2bag, Tensor bag_size, Tensor maximum_i…
1999 bag_size: non_differentiable
|