fn
embedding_bag_op
→TensorImplPtrint embedding_bag_op(const int & weight, const int & indices, const int & offsets, int mode, int padding_idx, bool include_last_offset)Python wrappers
Public Python APIs implemented by this engine symbol.Pooled embedding lookup with per-bag reduction.
Fuses a row gather and a reduction over bags into a single backend
call, avoiding the intermediate tensor that
embedding followed by sum/mean/max would materialize.
Variable-length bags are encoded as a flat indices tensor plus
an offsets tensor that marks the start of each bag — bag
consists of indices[offsets[b] : offsets[b+1]] (the last bag
runs to the end of indices when include_last_offset=false).
No autograd node is currently wired for this op (sparse gradients
are not yet supported on the Lucid backends), so the returned
tensor is not differentiable with respect to weight.
Math
\begin{cases} \displaystyle\sum_{j\in\mathcal{B}_b} W[j] & \text{mode}=0\\[6pt] \displaystyle\frac{1}{|\mathcal{B}_b|}\sum_{j\in\mathcal{B}_b} W[j] & \text{mode}=1\\[6pt] \displaystyle\max_{j\in\mathcal{B}_b} W[j] & \text{mode}=2 \end{cases}$$ **See Also** embedding_op : Per-token lookup without reduction.Parameters
weightTensorImplPtrEmbedding matrix of shape
(V, D).indicesTensorImplPtrFlat 1-D integer tensor of all bag contents concatenated.
offsetsTensorImplPtr1-D integer tensor of shape
(B,) marking the start of each bag.modeintReduction selector —
0 = sum, 1 = mean, 2 = max.padding_idxintIndex that contributes zero to the reduction; negative disables.
include_last_offsetboolIf
true, offsets is treated as (B+1,) so the last bag also has an explicit upper bound.Returns
TensorImplPtrOutput of shape (B, D) — one pooled embedding per bag.