Versioned name: EmbeddingBagPackedSum-3
Category: Sparse
Short description: Computes sums of "bags" of embeddings, without instantiating the intermediate embeddings.
Detailed description: This is the first case of the PyTorch EmbeddingBag, it has indices in the tensor of format [batch, indices_per_bag]
. If 3rd input is not provided, this operation is equivalent to Gather followed by ReduceSum(axis=0). However, EmbeddingBagPackedSum is much more time and memory efficient than using a chain of these operations.
Inputs:
emb_table
tensor containing the embedding lookup table of the module of shape [num_emb, emb_dim1, emb_dim2, ...]
and of type T. Required.indices
tensor of shape [batch, indices_per_bag]
and of type T_IND. Required.per_sample_weights
tensor of the same shape as indices
and of type T. Each value in this tensor are multiplied with each value pooled from embedding table for each index. Optional, default is tensor of ones.Outputs:
[batch, emb_dim1, emb_dim2, ...]
and of type T containing embeddings for each bag.Types
int32
or int64
.Example