bitorch.layers.qembedding.QEmbeddingBag

class bitorch.layers.qembedding.QEmbeddingBag(*args: Any, embedding_dim: int, weight_quantization: Optional[Union[Quantization, str]] = None, output_quantization: Optional[Union[Quantization, str]] = None, **kwargs: Any)[source]

Quantized version of pytorchs embedding bag. With the input indices the embedding is computed with a quantized version of the layers weight table. The output embedding will be also quantized before return.

Methods

__init__

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward

generates embeddings for received bags.

Attributes

__init__(*args: Any, embedding_dim: int, weight_quantization: Optional[Union[Quantization, str]] = None, output_quantization: Optional[Union[Quantization, str]] = None, **kwargs: Any) → None[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(input: Tensor, offsets: Optional[Tensor] = None, per_sample_weights: Optional[Tensor] = None) → Tensor[source]

generates embeddings for received bags. then quantizes these embeddings and depending on configuration forwards it through another quantized linear layer.

Parameters:

  • input (Tensor) – indices list for embeddings

  • offsets (Optional[Tensor], optional) – offsets to determine embedding sequences. Defaults to None.

  • per_sample_weights (Optional[Tensor], optional) – sample weights. Defaults to None.

Returns:

embeddings for given sequences

Return type:

Tensor