| /aosp_15_r20/external/pytorch/aten/src/ATen/native/ |
| H A D | EmbeddingBag.cpp | 57 static void make_offset2bag(const Tensor &offsets, Tensor& offset2bag) { in make_offset2bag() argument
58 offset2bag.index_add_( in make_offset2bag()
59 … 0, offsets, at::ones_like(offsets, LEGACY_CONTIGUOUS_MEMORY_FORMAT)); // offset2bag = [1 0 1 0 1] in make_offset2bag()
60 offset2bag[0] -= 1; // offset2bag = [0 0 1 0 1] in make_offset2bag()
61 offset2bag = offset2bag.cumsum(0, offset2bag.scalar_type()); // offset2bag = [0 0 1 1 2] in make_offset2bag()
966 Tensor& offset2bag, in make_offset2bag_out() argument
975 // mode, we skip calculating offset2bag, since it is not going to be used. in make_offset2bag_out()
980 at::native::resize_(offset2bag, {indices.size(0) + 1}, std::nullopt); in make_offset2bag_out()
981 at::native::zero_(offset2bag); in make_offset2bag_out()
992 make_offset2bag(_offsets, offset2bag); in make_offset2bag_out()
[all …]
|
| H A D | EmbeddingBag.h | 51 Tensor& offset2bag,
128 void _embedding_bag_cpu_impl_out(Tensor& output, Tensor& offset2bag,
139 at::Tensor& offset2bag,
|
| H A D | NonSymbolicBC.h | 19 …t at::Tensor & indices, const at::Tensor & offsets, const at::Tensor & offset2bag, const at::Tenso…
20 …t at::Tensor & indices, const at::Tensor & offsets, const at::Tensor & offset2bag, const at::Tenso…
|
| H A D | native_functions.yaml | 2303 # e.g. `indices` and `offsets` (and `offset2bag`), are contiguous.
2339 - func: _embedding_bag_backward(Tensor grad, Tensor indices, Tensor offsets, Tensor offset2bag, Ten…
2343 …ag_sparse_backward(Tensor grad, Tensor indices, Tensor offsets, Tensor offset2bag, Tensor bag_size…
2347 - func: _embedding_bag_dense_backward(Tensor grad, Tensor indices, Tensor offset2bag, Tensor bag_si…
2353 …ard(Tensor grad, Tensor weight, Tensor indices, Tensor offsets, Tensor offset2bag, int mode, int p…
|
| /aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/ |
| H A D | EmbeddingBag.cu | 64 index_t *offset2bag, int64_t numIndices, int64_t numBags, in EmbeddingBag_updateOutputKernel_max() argument
99 offset2bag[emb] = bag; in EmbeddingBag_updateOutputKernel_max()
114 index_t *offset2bag, int64_t numIndices, int64_t numBags, in EmbeddingBag_updateOutputKernel_sum_mean() argument
154 offset2bag[emb] = bag; in EmbeddingBag_updateOutputKernel_sum_mean()
171 const Tensor &offset2bag, in embedding_bag_backward_cuda_sum_avg() argument
235 count, num_weights, padding_idx, mode == EmbeddingBagMode::MEAN, offset2bag, in embedding_bag_backward_cuda_sum_avg()
377 auto offset2bag = in _embedding_bag_cuda() local
378 at::empty({indices.size(0)}, indices.options()); // offset2bag = [0 0 0 0 0] in _embedding_bag_cuda()
405 offset2bag.mutable_data_ptr<index_t>(), numIndices, numBags, featureSize, in _embedding_bag_cuda()
414 offset2bag.mutable_data_ptr<index_t>(), numIndices, numBags, featureSize, in _embedding_bag_cuda()
[all …]
|
| H A D | EmbeddingBackwardKernel.cu | 83 const index_t *offset2bag, const index_t *count, ptrdiff_t numel, in compute_grad_weight_bags() argument
106 const int seq_number = offset2bag[origRow]; in compute_grad_weight_bags()
220 const Tensor &offset2bag, in embedding_backward_cuda_kernel() argument
320 if (offset2bag.defined()) { in embedding_backward_cuda_kernel()
324 offset2bag.const_data_ptr<index_t>(), in embedding_backward_cuda_kernel()
|
| H A D | EmbeddingBackwardKernel.cuh | 18 const Tensor &offset2bag = Tensor(),
|
| /aosp_15_r20/external/pytorch/test/ |
| H A D | test_meta.py | 1607 # causes offset2bag to not get allocated... but the backward function
1608 # needs it, and the offset2bag computation lives inside the
1610 # To test sum, need to manually compute offset2bag
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(
1657 grad, weight, indices, offsets, offset2bag, mode, padding_idx
1661 offsets.to('meta'), offset2bag.to('meta'), mode, padding_idx
|
| /aosp_15_r20/external/pytorch/torch/csrc/inductor/aoti_torch/generated/ |
| H A D | c_shim_cpu.h | 23 …ward(AtenTensorHandle grad, AtenTensorHandle indices, AtenTensorHandle offset2bag, AtenTensorHandl…
25 …, AtenTensorHandle indices, AtenTensorHandle offsets, AtenTensorHandle offset2bag, int64_t mode, i…
|
| H A D | c_shim_cuda.h | 26 …ward(AtenTensorHandle grad, AtenTensorHandle indices, AtenTensorHandle offset2bag, AtenTensorHandl…
28 …, AtenTensorHandle indices, AtenTensorHandle offsets, AtenTensorHandle offset2bag, int64_t mode, i…
|
| /aosp_15_r20/external/pytorch/torch/ |
| H A D | _meta_registrations.py | 3425 offset2bag = indices.new_empty(indices.size(0))
3434 offset2bag = offsets.new_empty(indices.size(0))
3436 offset2bag = offsets.new_empty(0)
3450 return output, offset2bag, bag_size, max_indices
3455 output, offset2bag, bag_size, max_indices = meta_embedding_bag(
3460 return output, offset2bag, bag_size, max_indices
6353 offset2bag, argument
6368 offset2bag,
6380 offset2bag,
6395 offset2bag, argument
[all …]
|
| /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 - name: _embedding_bag_dense_backward(Tensor grad, Tensor indices, Tensor offset2bag, Tensor bag_si…
1998 offset2bag: non_differentiable
|