Embedding

Token gathering and row-wise scatter-add.

2 min readv1.0.0C++20

From the implementationinclude/nexus_model/layers/embedding.hpp

01Table and shapes#

Embedding(num_embeddings,embedding_dim,padding_idx=-1). [V,D] weights start normal(0,1). D is appended to the input shape: [B,T]→[B,T,D]. Keep input rank≤7 for the extra axis.

1Embedding tokens(1000,64,0);2auto ids=Tensor::from_values({1,4},{12,9,0,4});3auto y=tokens.forward(ids); // [1,4,64]

02Index contract#

Indices are stored in float32 Tensor and cast to int in the kernel. Use integer-valued indices in range; fractional values may not be rejected. Nonnegative padding_idx must lie in [0,V).

03Repeated-token gradients#

Repeated tokens add into the same row. padding_idx has zero output and no gradient contribution. backward returns zero index gradients. grad remains dense despite updating only touched rows.

Y_i=W_{k_i},\qquad dW_j\mathrel{+}=\sum_{i:k_i=j}dY_i

◆Read alongside the source. API behavior reflects the current implementation.

01Table and shapes#

Embedding(num_embeddings,embedding_dim,padding_idx=-1). [V,D] weights start normal(0,1). D is appended to the input shape: [B,T]→[B,T,D]. Keep input rank≤7 for the extra axis.

1Embedding tokens(1000,64,0);2auto ids=Tensor::from_values({1,4},{12,9,0,4});3auto y=tokens.forward(ids); // [1,4,64]

03Repeated-token gradients#

Repeated tokens add into the same row. padding_idx has zero output and no gradient contribution. backward returns zero index gradients. grad remains dense despite updating only touched rows.

Y_i=W_{k_i},\qquad dW_j\mathrel{+}=\sum_{i:k_i=j}dY_i