tensorplay.nn.functional.embedding_bag

tensorplay.nn.functional.embedding_bag(input: TensorBase, weight: TensorBase, offsets=None, max_norm=None, norm_type: float = 2, scale_grad_by_freq: bool = False, mode: str = 'mean', sparse: bool = False, per_sample_weights=None, include_last_offset: bool = False, padding_idx=None) → TensorBase[source]

Compute sums, means or maxes of bags of embeddings.

Composition of aten::embedding_bag from dispatched primitives; supports 1-D inputs with offsets (incl. include_last_offset), fixed length 2-D inputs, per_sample_weights (sum mode), padding_idx exclusion and max_norm renormalization. See tensorplay.nn.EmbeddingBag for details.