core.nn

Layers shared by more than one model, and the attention backends they run on.

Layers

Embedding	A simple extension of torch.nn.Embedding with a configurable initializer.
MultitaskReadout	Linear readout with one head per task, routed by integer index.

Initialization

tfixup_init_

Scale a Transformer encoder's weights in place for optimization without warmup.

Variable-length attention

Two interchangeable backends behind one call signature, so a model built on chained tokens can drop xformers entirely.

self_attn	forward_varlen for RotarySelfAttention, on either backend.
cross_attn	forward_varlen for RotaryCrossAttention, on either backend.
validate_attn_impl	Narrow a config string to a known backend, raising on anything else.
uses_xformers	What to pass as use_xformers when constructing the attention modules.