EmbeddingBag
ModuleEmbeddingBag(num_embeddings: int, embedding_dim: int, mode: str = 'mean', padding_idx: int | None = None)Quantized EmbeddingBag — per-row int8 table, pooled dequant lookup.
The inference-time replacement for a calibrated float
lucid.nn.EmbeddingBag, installed by
lucid.quantization.convert / from_float. Like
lucid.nn.quantized.Embedding it stores the table as int8 with a
per-row (per-token) scale — the same dominant memory win for
large-vocabulary models — but instead of returning one row per index it
pools each bag of indices into a single vector, the fused primitive behind
recommendation and bag-of-words text models.
Representation. The table (num_embeddings, embedding_dim) is
quantized per-row into int8 codes plus a per-row scale /
zero_point at from_float time; the float table is dropped. Each
forward dequantizes the table, gathers the rows in each bag, and reduces them
by mode — fused so the per-token embedding matrix is never materialised.
Given a flat index sequence partitioned into bags by offsets, the
-th pooled output is
where are the row- scale / zero-point (symmetric
qint8) and reduce is one of sum, mean, or max selected by
mode. The reduction runs on the dequantized (float) rows.
Parameters
num_embeddingsintembedding_dimintmode(sum, mean, max)= "sum""mean".padding_idxint or None= NoneNone.Attributes
Notes
- Instances are normally produced by
lucid.quantization.convert(orfrom_float), not constructed directly: a bare instance holds a zeroed table with identity qparams and pools to zeros untilfrom_float/load_state_dictpopulates the buffers. - The table is quantized per-row on axis 0 (one
scaleper token) with a symmetricqint8grid. When the source module carries noqconfig, a default per-channelqint8observer calibrates the rows. - Memory: the int8 codes plus per-row scale shrink the table payload
~
3.55xversusfloat32— as withEmbedding, the layer wins on memory, not compute (pooling runs in float post-dequantize). offsetsmarks where each bag begins in a flat index tensor; pass a 2-D index tensor instead to treat every row as its own fixed-length bag.
Examples
>>> import lucid, lucid.nn as nn
>>> bag = nn.EmbeddingBag(1000, 32, mode="mean")
>>> qbag = nn.quantized.EmbeddingBag.from_float(bag)
>>> idx = lucid.tensor([1, 2, 4, 5, 4])
>>> offsets = lucid.tensor([0, 3]) # two bags: [1,2,4] and [5,4]
>>> qbag(idx, offsets).shape # one pooled vector per bag
(2, 32)See Also
- lucid.nn.quantized.Embedding—The unpooled (one-row-per-index) analogue.
- lucid.nn.functional.embedding_bag—The pooled lookup run each forward.
- lucid.quantization.convert—Installs this layer from a calibrated model.
Used by 1
Constructors
1Class methods
1from_float(mod: nn.Module)Quantize a float lucid.nn.EmbeddingBag's table (per-row int8).