Documentation
¶
Overview ¶
Package transformer provides transformer-related neural network layers.
Index ¶
- type Block
- func (tb *Block[T]) Backward(ctx context.Context, dOut *tensor.Tensor[T], inputs ...*tensor.Tensor[T]) ([]*tensor.Tensor[T], error)
- func (tb *Block[T]) Forward(ctx context.Context, inputs ...*tensor.Tensor[T]) (*tensor.Tensor[T], error)
- func (tb *Block[T]) OutputShape(inputShapes ...[]int) ([]int, error)
- func (tb *Block[T]) Parameters() []graph.Parameter[T]
Constants ¶
This section is empty.
Variables ¶
This section is empty.
Functions ¶
This section is empty.
Types ¶
type Block ¶
Block implements a single Transformer encoder block.
func NewTransformerBlock ¶
func NewTransformerBlock[T tensor.Numeric](engine compute.Engine[T], ops numeric.Arithmetic[T], modelDim, numQueryHeads, numKeyValueHeads, ffnDim int, epsilon T) (*Block[T], error)
NewTransformerBlock creates a new Transformer encoder block. modelDim: The input and output dimension of the block (d_model). numQueryHeads: The number of query heads. numKeyValueHeads: The number of key/value heads. ffnDim: The inner dimension of the feed-forward network. epsilon: Epsilon for Layer Normalization.
func (*Block[T]) Backward ¶
func (tb *Block[T]) Backward(ctx context.Context, dOut *tensor.Tensor[T], inputs ...*tensor.Tensor[T]) ([]*tensor.Tensor[T], error)
Backward computes the gradients for the Transformer Block.
func (*Block[T]) Forward ¶
func (tb *Block[T]) Forward(ctx context.Context, inputs ...*tensor.Tensor[T]) (*tensor.Tensor[T], error)
Forward computes the Transformer Block's forward pass.
func (*Block[T]) OutputShape ¶
OutputShape returns the output shape, which is the same as the input shape.
func (*Block[T]) Parameters ¶
Parameters returns all trainable parameters from its sub-layers.
Source Files