Python wrappers
Public Python APIs implemented by this engine symbols.Autograd node for the embedding lookup on a weight matrix.
Only the weight matrix is a differentiable input. Integer indices
have no meaningful gradient and are not stored in the standard
saved_inputs_ slot — they live in saved_indices_
instead, alongside their shape and dtype, so the backward can
scatter-add grad_out into the correct rows of dW.
Forward dispatches to IBackend::embedding_forward which gathers
rows from the weight storage at positions given by indices,
zeroing any output row whose index equals padding_idx. Backward
uses IBackend::embedding_backward to scatter-add grad_out
into a zero-initialized dW at those same positions.
Attributes
schema_v1OpSchema"embedding", version 1, AmpPolicy::Promote).padding_idx_int-1 disables the masking entirely.weight_shape_Shape(num_embeddings, embed_dim); recovered during backward to size dW.saved_indices_Storagesaved_indices_shape_Shapesaved_indices_dtype_DtypeI64).Static methods
1forward
→TensorImplPtrint forward(const int & weight, const int & indices, int padding_idx)Look up rows of weight at positions given by indices.
Math
\qquad y[\ldots] = 0 \text{ if } \text{idx}[\ldots] = \text{padding\_idx}.$$ **Shape** - `weight` : $(V, D)$ - `indices` : $(\ast)$ integer - Output : $(\ast, D)$Parameters
weightTensorImplPtr(num_embeddings, embed_dim). Must be exactly 2-D.indicesTensorImplPtr[0, num_embeddings).padding_idxintReturns
TensorImplPtrOutput of shape — the input index shape with an extra trailing dim of embed_dim.
Raises
ShapeMismatchweight is not 2-D.DeviceMismatchweight and indices live on different devices.Methods
1Scatter-add backward.
Returns a single dW storage of shape weight_shape_,
zero-initialized except where the saved indices accumulate
contributions from grad_out. Rows at padding_idx_
accumulate nothing.
Parameters
grad_outStorageReturns
std::vector<Storage>Single-element {dW}.