Transformer Config Defaults to GPT2

Sinusoidal Positional Encoding Module. The purpose of positional encoding is to add information about the relative positions of tokens within

Encoder Only Transformer.

Multi-Head Attention Module. This module takes an input residual stream and applies multiple attention heads that operate in parallel to proc

Test Dataset.

Embedding Module. Embedding involves taking each token, and embedding it from the high-dimensional vocabulary space (which can be of size e.g

Full layer (attention + feed forward). The layer receives the residual stream as an input. It then applies the multi-head attention sub-layer

MLP The MLP module takes an input of the residual stream and applies a standard two-layer feed forward network. The resulting output will the

Simple model for testing.

Unembedding Module. Unembedding involves taking each token at the end of the residual stream, and "unembedding" it from the low-dimensional r

Flagged Token Dataset Dataset designed to test that the model can learn to attend to a token, and that it can learn a what to do after attend

Ordered integers dataset. The samples are contiguous sets of numbers (e.g. {3,1,2,4,0}). The targets are simply the next numbers after this (

Zero Layer Model. Contains just embedding and unembedding layers.

Bit Reversal Dataset. The model should learn to reverse the order of the input bits. For example, if the input is [0, 1, 1, 0, 1, 0, 0, 1, 0,

Identity Dataset. The model should learn to output the same input values. Each input neuron is connected to the corresponding output neuron.

Odd Even Dataset. The model should learn to separate the input bits based on their index (odd or even). The output should have the odd-indexe

Regression task dataset. Extend this to add datasets that the model should be able to learn.

Backwards compatible string enum.

Tensor Shape Labels Labels (aka dimension names) used in describing tensors.

Attention Calculation Tests.

Default device tests.

Evaluate tests.

Attention Mask Tests.

Multi-Head Attention Module Tests.

TensorShapeLabels tests.

Train loop tests.