torch.nn.functional.embedding_bag
function embedding_bag<S extends readonly [number, number]>(input: Tensor<S>, weight: Tensor, options?: EmbeddingBagFunctionalOptions): Tensorfunction embedding_bag(input: Tensor, weight: Tensor, offsets: Tensor | undefined, max_norm: number | undefined, norm_type: number, scale_grad_by_freq: boolean, mode: 'sum' | 'mean' | 'max', sparse: boolean, per_sample_weights: Tensor | undefined, include_last_offset: boolean, padding_idx: number | undefined, options?: EmbeddingBagFunctionalOptions): TensorCompute sums, means, or maxes of bags of embeddings.
NOTE: Currently only 2D input is supported, where each row is a bag. 1D input with offsets is not yet implemented.
Parameters
inputTensor<S>- Tensor of indices (must be 2D: [num_bags, bag_size])
weightTensor- Embedding weight matrix of shape (num_embeddings, embedding_dim)
optionsEmbeddingBagFunctionalOptionsoptional