Documentation
¶
Overview ¶
Package embeddings provides neural network embedding layers for the Zerfoo ML framework.
Index ¶
- type RotaryPositionalEmbedding
- func (rpe *RotaryPositionalEmbedding[T]) Backward(ctx context.Context, dOut *tensor.Tensor[T], _ ...*tensor.Tensor[T]) ([]*tensor.Tensor[T], error)
- func (rpe *RotaryPositionalEmbedding[T]) Forward(ctx context.Context, input *tensor.Tensor[T], _ ...*tensor.Tensor[T]) (*tensor.Tensor[T], error)
- func (rpe *RotaryPositionalEmbedding[T]) OutputShape(inputShapes ...[]int) ([]int, error)
- func (rpe *RotaryPositionalEmbedding[T]) Parameters() []graph.Parameter[T]
- type TokenEmbedding
- func (te *TokenEmbedding[T]) Backward(ctx context.Context, dOut *tensor.Tensor[T], _ ...*tensor.Tensor[T]) ([]*tensor.Tensor[T], error)
- func (te *TokenEmbedding[T]) Forward(ctx context.Context, tokenIDs *tensor.Tensor[int], _ ...*tensor.Tensor[T]) (*tensor.Tensor[T], error)
- func (te *TokenEmbedding[T]) OutputShape(inputShapes ...[]int) ([]int, error)
- func (te *TokenEmbedding[T]) Parameters() []graph.Parameter[T]
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type RotaryPositionalEmbedding ¶
type RotaryPositionalEmbedding[T tensor.Numeric] struct { // contains filtered or unexported fields }
RotaryPositionalEmbedding applies Rotary Positional Embedding to a tensor.
func NewRotaryPositionalEmbedding ¶
func NewRotaryPositionalEmbedding[T tensor.Numeric]( ctx context.Context, engine compute.Engine[T], headDim int, seqLen int, ) (*RotaryPositionalEmbedding[T], error)
NewRotaryPositionalEmbedding creates a new RotaryPositionalEmbedding layer. headDim: The dimension of the head. Must be even. seqLen: The maximum sequence length this embedding will be applied to. engine: The compute engine to use for tensor operations.
func (*RotaryPositionalEmbedding[T]) Backward ¶
func (rpe *RotaryPositionalEmbedding[T]) Backward(ctx context.Context, dOut *tensor.Tensor[T], _ ...*tensor.Tensor[T]) ([]*tensor.Tensor[T], error)
Backward computes the gradients for RoPE.
func (*RotaryPositionalEmbedding[T]) Forward ¶
func (rpe *RotaryPositionalEmbedding[T]) Forward(ctx context.Context, input *tensor.Tensor[T], _ ...*tensor.Tensor[T]) (*tensor.Tensor[T], error)
Forward applies Rotary Positional Embedding to the input tensor.
func (*RotaryPositionalEmbedding[T]) OutputShape ¶
func (rpe *RotaryPositionalEmbedding[T]) OutputShape(inputShapes ...[]int) ([]int, error)
OutputShape returns the output shape of the RoPE layer. Input shape is (batch_size, seq_len, head_dim). Output shape is the same.
func (*RotaryPositionalEmbedding[T]) Parameters ¶
func (rpe *RotaryPositionalEmbedding[T]) Parameters() []graph.Parameter[T]
Parameters returns no trainable parameters for RoPE.
type TokenEmbedding ¶
type TokenEmbedding[T float64] struct { // contains filtered or unexported fields }
TokenEmbedding converts token IDs into dense vector representations.
func NewTokenEmbedding ¶
func NewTokenEmbedding[T float64](engine compute.Engine[T], vocabSize, embeddingDim int) (*TokenEmbedding[T], error)
NewTokenEmbedding creates a new TokenEmbedding layer. vocabSize: The size of the vocabulary (number of unique tokens). embeddingDim: The dimension of the embedding vectors.
func (*TokenEmbedding[T]) Backward ¶
func (te *TokenEmbedding[T]) Backward(ctx context.Context, dOut *tensor.Tensor[T], _ ...*tensor.Tensor[T]) ([]*tensor.Tensor[T], error)
Backward computes the gradients for the embedding table.
func (*TokenEmbedding[T]) Forward ¶
func (te *TokenEmbedding[T]) Forward(ctx context.Context, tokenIDs *tensor.Tensor[int], _ ...*tensor.Tensor[T]) (*tensor.Tensor[T], error)
Forward performs the embedding lookup. Input: A tensor of token IDs (int type). Output: A tensor of embedding vectors (T type).
func (*TokenEmbedding[T]) OutputShape ¶
func (te *TokenEmbedding[T]) OutputShape(inputShapes ...[]int) ([]int, error)
OutputShape returns the output shape of the embedding layer. Input shape is (batch_size, seq_len). Output shape is (batch_size, seq_len, embedding_dim).
func (*TokenEmbedding[T]) Parameters ¶
func (te *TokenEmbedding[T]) Parameters() []graph.Parameter[T]
Parameters returns the trainable embedding table.
Source Files