torch_brain.nn#

Embedding modules#

`Embedding`	A simple extension of `torch.nn.Embedding` to allow more control over the weights initializer.
`InfiniteVocabEmbedding`	Embedding layer with a vocabulary that can be extended.
`RotaryTimeEmbedding`	Rotary time/positional embedding layer.
`SinusoidalTimeEmbedding`	Sinusoidal time/position embedding layer.

`MultitaskReadout`	A module that performs multi-task linear readouts from output embeddings.
`prepare_for_multitask_readout`

`RotaryCrossAttention`	Cross-attention layer with rotary positional embeddings.
`RotarySelfAttention`	Self-attention layer with rotary positional embeddings.