Neural Network Components API

Composable PyTorch modules for language-model architectures.

Modules

Module	Public API
`olm.nn.activations.base`	`ActivationBase`
`olm.nn.activations.elu`	`ELU`
`olm.nn.activations.geglu`	`GeGLU`
`olm.nn.activations.gelu`	`GELU`
`olm.nn.activations.glu`	`GLU`
`olm.nn.activations.identity`	`Identity`
`olm.nn.activations.leaky_relu`	`LeakyReLU`
`olm.nn.activations.liglu`	`LiGLU`
`olm.nn.activations.mish`	`Mish`
`olm.nn.activations.prelu`	`PReLU`
`olm.nn.activations.reglu`	`ReGLU`
`olm.nn.activations.relu`	`ReLU`
`olm.nn.activations.selu`	`SELU`
`olm.nn.activations.sigmoid`	`Sigmoid`
`olm.nn.activations.silu`	`SiLU`, `Swish`
`olm.nn.activations.softmax`	`Softmax`
`olm.nn.activations.softplus`	`Softplus`
`olm.nn.activations.swiglu`	`SwiGLU`
`olm.nn.activations.swish`	`Swish`
`olm.nn.activations.tanh`	`Tanh`
`olm.nn.attention`	`AttentionBase`, `AttentionwithRoPEBase`, `FlashAttention`, `FlashAttentionwithRoPE`, `GroupedQueryAttention`, `MultiHeadAttention`, `MultiHeadAttentionwithALiBi`, `MultiHeadAttentionwithRoPE`
`olm.nn.attention.alibi`	`MultiHeadAttentionwithALiBi`
`olm.nn.attention.base`	`AttentionBase`, `AttentionwithRoPEBase`
`olm.nn.attention.flash`	`FlashAttention`, `FlashAttentionwithRoPE`
`olm.nn.attention.gqa`	`GroupedQueryAttention`
`olm.nn.attention.masks`	`attention_mask_to_bool`
`olm.nn.attention.mha`	`MultiHeadAttention`, `MultiHeadAttentionwithRoPE`
`olm.nn.blocks.LM`	`LM`
`olm.nn.blocks.linear_projections`	`QKVProjection`
`olm.nn.blocks.output_head`	`OutputHead`
`olm.nn.blocks.transformer_block`	`TransformerBlock`
`olm.nn.embeddings`	`ALiBiPositionalBias`, `AbsolutePositionalEmbedding`, `Embedding`, `PartialRotaryPositionalEmbedding`, `PartialScaledRotaryPositionalEmbedding`, `PositionalEmbeddingBase`, `RotaryPositionalEmbedding`, `ScaledRotaryPositionalEmbedding`, +1 more
`olm.nn.embeddings.positional`	`ALiBiPositionalBias`, `AbsolutePositionalEmbedding`, `PartialRotaryPositionalEmbedding`, `PartialScaledRotaryPositionalEmbedding`, `PositionalEmbeddingBase`, `RotaryPositionalEmbedding`, `ScaledRotaryPositionalEmbedding`, `SinusoidalPositionalEmbedding`
`olm.nn.embeddings.positional.absolute`	`AbsolutePositionalEmbedding`
`olm.nn.embeddings.positional.alibi`	`ALiBiPositionalBias`
`olm.nn.embeddings.positional.base`	`PositionalEmbeddingBase`
`olm.nn.embeddings.positional.rope`	`PartialRotaryPositionalEmbedding`, `PartialScaledRotaryPositionalEmbedding`, `RotaryPositionalEmbedding`, `ScaledRotaryPositionalEmbedding`
`olm.nn.embeddings.positional.sinusoidal`	`SinusoidalPositionalEmbedding`
`olm.nn.embeddings.token_embed`	`Embedding`
`olm.nn.feedforward`	`ClassicFFN`, `ClassicMoEFFN`, `FeedForwardBase`, `GeGLUFFN`, `GeGLUMoEFFN`, `SwiGLUFFN`, `SwiGLUMoEFFN`
`olm.nn.feedforward.base`	`FeedForwardBase`
`olm.nn.feedforward.classic_ffn`	`ClassicFFN`
`olm.nn.feedforward.classic_moe`	`ClassicMoEFFN`
`olm.nn.feedforward.geglu_ffn`	`GeGLUFFN`
`olm.nn.feedforward.geglu_moe`	`GeGLUMoEFFN`
`olm.nn.feedforward.moe_base`	`MoEFeedForwardBase`, `MoERouter`
`olm.nn.feedforward.swiglu_ffn`	`SwiGLUFFN`
`olm.nn.feedforward.swiglu_moe`	`SwiGLUMoEFFN`
`olm.nn.norms`	`LayerNorm`, `RMSNorm`
`olm.nn.norms.base`	`NormBase`
`olm.nn.norms.layer_norm`	`LayerNorm`
`olm.nn.norms.rms_norm`	`RMSNorm`
`olm.nn.structure.block`	`Block`, `load`, `load_block`, `load_model`
`olm.nn.structure.combinators`	`BaseCombinator`, `Parallel`, `Repeat`, `Residual`
`olm.nn.structure.combinators.base`	`BaseCombinator`
`olm.nn.structure.combinators.parallel`	`Parallel`
`olm.nn.structure.combinators.repeat`	`Repeat`
`olm.nn.structure.combinators.residual`	`Residual`
`olm.nn.torch_nn_wrappers`	`Linear`