Model Implementations¶

spectrans.models ¶

Spectral transformer model implementations.

This module provides transformer model implementations that replace traditional attention mechanisms with various spectral mixing approaches. Each model maintains the core transformer architecture with residual connections, layer normalization, and feedforward networks.

The models implement spectral transformers including FNet, Global Filter Networks, AFNO, spectral attention variants, and hybrid architectures that combine spectral and spatial processing.

Modules:

Name	Description
`afno`	Adaptive Fourier Neural Operator models.
`base`	Base classes for models and components.
`fnet`	FNet models using Fourier mixing.
`fno_transformer`	Fourier Neural Operator transformer models.
`gfnet`	Global Filter Network models.
`hybrid`	Hybrid models combining spectral and attention.
`lst`	Linear Spectral Transform models.
`spectral_attention`	Models using spectral attention mechanisms.
`wavenet_transformer`	Wavelet-based transformer models.

Classes:

Name	Description
`BaseModel`	Abstract base class for all spectral transformer models.
`PositionalEncoding`	Sinusoidal positional encoding for sequence models.
`LearnedPositionalEncoding`	Learnable positional embedding layer.
`RotaryPositionalEncoding`	Rotary Position Embedding (RoPE) for improved length generalization.
`ALiBiPositionalBias`	Attention with Linear Biases (ALiBi) positional encoding.
`ClassificationHead`	Classification head for sequence classification tasks.
`RegressionHead`	Regression head for continuous prediction tasks.
`SequenceHead`	Generic sequence-to-sequence head for various tasks.
`FNet`	FNet model with Fourier mixing layers.
`FNetEncoder`	FNet encoder stack for encoder-only architectures.
`GFNet`	Global Filter Network model with learnable spectral filters.
`GFNetEncoder`	GFNet encoder stack implementation.
`AFNOEncoder`	Adaptive Fourier Neural Operator encoder.
`AFNOModel`	AFNO model for various tasks.
`SpectralAttentionEncoder`	Encoder using spectral attention with random Fourier features.
`SpectralAttentionTransformer`	Spectral attention transformer model.
`PerformerTransformer`	Performer-style transformer with linear attention approximation.
`LSTEncoder`	Linear Spectral Transform encoder using DCT/DST.
`LSTDecoder`	Linear Spectral Transform decoder implementation.
`LSTTransformer`	LST transformer with encoder-decoder architecture.
`FNOEncoder`	Fourier Neural Operator encoder for function space learning.
`FNODecoder`	FNO decoder for continuous function approximation.
`FNOTransformer`	FNO transformer for operator learning.
`WaveletEncoder`	Wavelet transform encoder with multiresolution analysis.
`WaveletDecoder`	Wavelet decoder for signal reconstruction.
`WaveletTransformer`	Wavelet transformer model.
`HybridEncoder`	Encoder combining spectral and spatial attention layers.
`HybridTransformer`	Hybrid model alternating between spectral and attention mechanisms.
`AlternatingTransformer`	Transformer with alternating spectral and attention layers.
`StandardAttention`	Standard multi-head self-attention wrapper for hybrid models.

Examples:

Basic FNet usage:

>>> import torch
>>> from spectrans.models import FNet
>>>
>>> # Create FNet model
>>> model = FNet(
...     hidden_dim=512,
...     num_layers=12,
...     vocab_size=32000,
...     max_seq_len=512
... )
>>>
>>> # Forward pass
>>> input_ids = torch.randint(0, 32000, (2, 128))
>>> outputs = model(input_ids)
>>> print(outputs.shape)  # torch.Size([2, 128, 512])

Global Filter Network example:

>>> from spectrans.models import GFNet
>>>
>>> # Create GFNet for sequence classification
>>> model = GFNet(
...     hidden_dim=768,
...     num_layers=12,
...     num_classes=10,
...     sequence_length=256
... )
>>>
>>> # Classification forward pass
>>> x = torch.randn(4, 256, 768)
>>> logits = model(x)
>>> print(logits.shape)  # torch.Size([4, 10])

AFNO for continuous functions:

>>> from spectrans.models import AFNOModel
>>>
>>> # Create AFNO model
>>> model = AFNOModel(
...     hidden_dim=512,
...     num_layers=8,
...     n_modes=32,
...     input_dim=2,
...     output_dim=1
... )
>>>
>>> # Function approximation
>>> x = torch.randn(8, 64, 64, 2)  # Batch of 2D functions
>>> output = model(x)
>>> print(output.shape)  # torch.Size([8, 64, 64, 1])

Hybrid spectral-attention model:

>>> from spectrans.models import HybridTransformer
>>>
>>> # Create hybrid model alternating spectral and attention
>>> model = HybridTransformer(
...     hidden_dim=512,
...     num_layers=12,
...     num_heads=8,
...     spectral_type="fourier",
...     vocab_size=50000
... )
>>>
>>> input_ids = torch.randint(0, 50000, (2, 256))
>>> outputs = model(input_ids)
>>> print(outputs.shape)  # torch.Size([2, 256, 512])

Wavelet transformer for multiresolution analysis:

>>> from spectrans.models import WaveletTransformer
>>>
>>> # Create wavelet transformer
>>> model = WaveletTransformer(
...     hidden_dim=512,
...     num_layers=8,
...     wavelet="db4",
...     levels=3,
...     vocab_size=32000
... )
>>>
>>> input_ids = torch.randint(0, 32000, (2, 512))
>>> outputs = model(input_ids)
>>> print(outputs.shape)  # torch.Size([2, 512, 512])

Positional encodings with RoPE and ALiBi:

>>> from spectrans.models import FNet, RotaryPositionalEncoding, ALiBiPositionalBias
>>>
>>> # FNet with RoPE
>>> model = FNet(
...     hidden_dim=512,
...     num_layers=12,
...     vocab_size=50000,
...     pos_encoding=RotaryPositionalEncoding(dim=512)
... )
>>>
>>> # Or with ALiBi (no positional embeddings needed)
>>> model_alibi = FNet(
...     hidden_dim=512,
...     num_layers=12,
...     vocab_size=50000,
...     pos_encoding=ALiBiPositionalBias(num_heads=8)
... )

Notes

All models in this module follow the same architectural principles. Spectral processing replaces quadratic attention with spectral transforms that scale as \(O(n \log n)\) or \(O(n)\) in time complexity. Residual connections maintain gradient flow around each spectral layer and feedforward network. Layer normalization is applied before spectral mixing and feedforward operations for training stability.

The models support multiple positional encoding methods including sinusoidal, learned embeddings, RoPE, and ALiBi for various sequence modeling needs. Specialized output heads are provided for classification, regression, and sequence-to-sequence tasks.

The mathematical foundation for spectral mixing is based on the convolution theorem, which states that convolution in the spatial domain is equivalent to element-wise multiplication in the frequency domain:

\[ \mathcal{F}[f * g] = \mathcal{F}[f] \odot \mathcal{F}[g] \]

This enables efficient global mixing of sequence elements through spectral transforms like FFT, DCT, or DWT, avoiding the quadratic complexity of traditional attention mechanisms.

Model Complexity Comparison:

Standard Transformer: \(O(n^2 d + nd^2)\) time, \(O(n^2 + nd)\) space
FNet: \(O(nd \log n + nd^2)\) time, \(O(nd)\) space
GFNet: \(O(nd \log n + nd^2)\) time, \(O(nd)\) space
AFNO: \(O(k_n k_d d + nd \log n)\) time, \(O(k_n k_d d)\) space
LST: \(O(nd \log n + nd^2)\) time, \(O(nd)\) space
Wavelet: \(O(nd + nd^2)\) time, \(O(nd)\) space

Where \(n\) is sequence length, \(d\) is hidden dimension, and \(k_n, k_d\) are retained spectral modes.

References

James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. 2022. FNet: Mixing tokens with Fourier transforms. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (NAACL-HLT), pages 4296-4313, Seattle.

Yongming Rao, Wenliang Zhao, Zheng Zhu, Jiwen Lu, and Jie Zhou. 2021. Global filter networks for image classification. In Advances in Neural Information Processing Systems 34 (NeurIPS 2021), pages 980-993.

John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. 2022. Adaptive Fourier neural operators: Efficient token mixers for transformers. In Proceedings of the International Conference on Learning Representations (ICLR).

Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. 2021. Fourier neural operator for parametric partial differential equations. In Proceedings of the International Conference on Learning Representations (ICLR).

Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. 2021. Rethinking attention with performers. In Proceedings of the International Conference on Learning Representations (ICLR).

Classes¶

AFNOEncoder ¶

AFNOEncoder(hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 1024, modes_seq: int | None = None, modes_hidden: int | None = None, mlp_ratio: float = 2.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, gradient_checkpointing: bool = False)

Bases: AFNOModel

Encoder-only AFNO model for representation learning.

This variant of AFNO is designed for tasks that require extracting representations rather than making predictions. It's particularly efficient for processing very long sequences due to the mode truncation.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden dimension size. Default is 768.	`768`
`num_layers`	`int`	Number of AFNO layers. Default is 12.	`12`
`max_sequence_length`	`int`	Maximum sequence length. Default is 1024.	`1024`
`modes_seq`	`int \| None`	Number of Fourier modes in sequence dimension.	`None`
`modes_hidden`	`int \| None`	Number of Fourier modes in hidden dimension.	`None`
`mlp_ratio`	`float`	MLP expansion ratio. Default is 2.0.	`2.0`
`use_positional_encoding`	`bool`	Whether to use positional encoding. Default is True.	`True`
`positional_encoding_type`	`str`	Type of positional encoding. Default is 'sinusoidal'.	`'sinusoidal'`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN. If None, defaults to 4 * hidden_dim.	`None`
`norm_eps`	`float`	Epsilon for layer normalization. Default is 1e-12.	`1e-12`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing. Default is False.	`False`

Source code in spectrans/models/afno.py

def __init__(
    self,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 1024,
    modes_seq: int | None = None,
    modes_hidden: int | None = None,
    mlp_ratio: float = 2.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    gradient_checkpointing: bool = False,
):
    super().__init__(
        vocab_size=None,  # No token embeddings for encoder
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        modes_seq=modes_seq,
        modes_hidden=modes_hidden,
        mlp_ratio=mlp_ratio,
        num_classes=None,  # No classification head
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type="none",  # Return hidden states
        gradient_checkpointing=gradient_checkpointing,
    )

AFNOModel ¶

AFNOModel(vocab_size: int | None = None, hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 1024, modes_seq: int | None = None, modes_hidden: int | None = None, mlp_ratio: float = 2.0, num_classes: int | None = None, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, output_type: OutputHeadType = 'classification', gradient_checkpointing: bool = False)

Bases: BaseModel

Adaptive Fourier Neural Operator transformer model.

AFNO performs token mixing using truncated Fourier modes and learnable MLPs in the frequency domain, processing long sequences with \(O(n \log n)\) time complexity.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings. If None, expects pre-embedded inputs.	`None`
`hidden_dim`	`int`	Hidden dimension size. Default is 768.	`768`
`num_layers`	`int`	Number of AFNO layers. Default is 12.	`12`
`max_sequence_length`	`int`	Maximum sequence length. Default is 1024.	`1024`
`modes_seq`	`int \| None`	Number of Fourier modes to keep in sequence dimension. If None, defaults to max_sequence_length // 2.	`None`
`modes_hidden`	`int \| None`	Number of Fourier modes to keep in hidden dimension. If None, defaults to hidden_dim // 2.	`None`
`mlp_ratio`	`float`	Expansion ratio for MLP in Fourier domain. Default is 2.0.	`2.0`
`num_classes`	`int \| None`	Number of output classes for classification. Default is None.	`None`
`use_positional_encoding`	`bool`	Whether to use positional encoding. Default is True.	`True`
`positional_encoding_type`	`str`	Type of positional encoding: 'sinusoidal' or 'learned'. Default is 'sinusoidal'.	`'sinusoidal'`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN. If None, defaults to 4 * hidden_dim.	`None`
`norm_eps`	`float`	Epsilon for layer normalization. Default is 1e-12.	`1e-12`
`output_type`	`str`	Type of output head: 'classification', 'regression', 'sequence', or 'none'. Default is 'classification'.	`'classification'`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing. Default is False.	`False`

Attributes:

Name	Type	Description
`modes_seq`	`int`	Number of Fourier modes in sequence dimension.
`modes_hidden`	`int`	Number of Fourier modes in hidden dimension.
`mlp_ratio`	`float`	MLP expansion ratio in frequency domain.
`blocks`	`ModuleList`	List of AFNO transformer blocks.

Methods:

Name	Description
`build_blocks`	Build AFNO transformer blocks with adaptive Fourier mixing.
`from_config`	Create AFNO model from configuration.

Source code in spectrans/models/afno.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 1024,
    modes_seq: int | None = None,
    modes_hidden: int | None = None,
    mlp_ratio: float = 2.0,
    num_classes: int | None = None,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    output_type: OutputHeadType = "classification",
    gradient_checkpointing: bool = False,
):
    self.modes_seq = modes_seq or (max_sequence_length // 2)
    self.modes_hidden = modes_hidden or (hidden_dim // 2)
    self.mlp_ratio = mlp_ratio
    self._dropout_rate = dropout  # Store dropout rate for build_blocks

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type=output_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build AFNO transformer blocks with adaptive Fourier mixing.

Returns:

Type	Description
`ModuleList`	List of AFNO transformer blocks.

Source code in spectrans/models/afno.py

def build_blocks(self) -> nn.ModuleList:
    """Build AFNO transformer blocks with adaptive Fourier mixing.

    Returns
    -------
    nn.ModuleList
        List of AFNO transformer blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Create AFNO mixing layer
        mixing_layer = AFNOMixing(
            hidden_dim=self.hidden_dim,
            max_sequence_length=self.max_sequence_length,
            modes_seq=self.modes_seq,
            modes_hidden=self.modes_hidden,
            mlp_ratio=self.mlp_ratio,
            dropout=self._dropout_rate,
        )

        # Create AFNO block with pre-normalization
        block = PreNormBlock(
            mixing_layer=mixing_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            activation="gelu",
            dropout=self._dropout_rate,
            norm_eps=self.norm_eps,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

from_config `classmethod` ¶

from_config(config: AFNOModelConfig) -> AFNOModel

Create AFNO model from configuration.

Parameters:

Name	Type	Description	Default
`config`	`AFNOModelConfig`	Configuration object with model parameters.	required

Returns:

Type	Description
`AFNOModel`	Configured AFNO model.

Source code in spectrans/models/afno.py

@classmethod
def from_config(cls, config: "AFNOModelConfig") -> "AFNOModel":  # type: ignore[override]
    """Create AFNO model from configuration.

    Parameters
    ----------
    config : AFNOModelConfig
        Configuration object with model parameters.

    Returns
    -------
    AFNOModel
        Configured AFNO model.
    """
    # Handle AFNO-specific mode configuration
    n_modes = getattr(config, "n_modes", None)
    modes_seq = getattr(config, "modes_seq", None)
    modes_hidden = getattr(config, "modes_hidden", None)

    # If n_modes is provided but not modes_seq/modes_hidden, compute them
    if n_modes is not None and modes_seq is None:
        modes_seq = n_modes
    if n_modes is not None and modes_hidden is None:
        compression_ratio = getattr(config, "compression_ratio", 0.5)
        modes_hidden = int(n_modes * compression_ratio)

    return cls(
        vocab_size=getattr(config, "vocab_size", None),
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        modes_seq=modes_seq,
        modes_hidden=modes_hidden,
        mlp_ratio=getattr(config, "mlp_ratio", 2.0),
        num_classes=getattr(config, "num_classes", None),
        use_positional_encoding=getattr(config, "use_positional_encoding", True),
        positional_encoding_type=getattr(config, "positional_encoding_type", "sinusoidal"),
        dropout=config.dropout,
        ffn_hidden_dim=getattr(config, "ffn_hidden_dim", None),
        norm_eps=getattr(config, "norm_eps", 1e-12),
        output_type=getattr(config, "output_type", "classification"),
        gradient_checkpointing=getattr(config, "gradient_checkpointing", False),
    )

ALiBiPositionalBias ¶

ALiBiPositionalBias(num_heads: int, max_sequence_length: int = 5000)

Bases: Module

Attention with Linear Biases (ALiBi) positional encoding.

This module implements ALiBi, which adds a linear bias to attention scores based on the relative distance between tokens. Unlike traditional position embeddings, ALiBi enables extrapolation to longer sequences.

Parameters:

Name	Type	Description	Default
`num_heads`	`int`	Number of attention heads.	required
`max_sequence_length`	`int`	Maximum sequence length to encode.	`5000`

Attributes:

Name	Type	Description
`num_heads`	`int`	Number of attention heads.
`slopes`	`Tensor`	Head-specific slope parameters.
`alibi`	`Tensor \| None`	Cached linear bias matrix.

References

Ofir Press, Noah A. Smith, and Mike Lewis. 2022. Train short, test long: Attention with linear biases enables input length extrapolation. In Proceedings of the International Conference on Learning Representations (ICLR).

Methods:

Name	Description
`forward`	Add ALiBi bias to attention scores.
`get_bias`	Get ALiBi bias matrix for a given sequence length.

Source code in spectrans/models/base.py

def __init__(
    self,
    num_heads: int,
    max_sequence_length: int = 5000,
):
    super().__init__()
    self.num_heads = num_heads
    self.max_sequence_length = max_sequence_length

    # Compute slopes for each head
    slopes = self._get_slopes(num_heads)
    self.register_buffer("slopes", slopes)

    # Cache for bias matrix
    self.alibi: Tensor | None = None
    self._build_alibi_tensor(max_sequence_length)

Functions¶

forward ¶

forward(attention_scores: Tensor) -> Tensor

Add ALiBi bias to attention scores.

Parameters:

Name	Type	Description	Default
`attention_scores`	`Tensor`	Attention scores of shape (batch_size, num_heads, seq_len, seq_len).	required

Returns:

Type	Description
`Tensor`	Attention scores with ALiBi bias added.

Source code in spectrans/models/base.py

def forward(self, attention_scores: Tensor) -> Tensor:
    """Add ALiBi bias to attention scores.

    Parameters
    ----------
    attention_scores : Tensor
        Attention scores of shape (batch_size, num_heads, seq_len, seq_len).

    Returns
    -------
    Tensor
        Attention scores with ALiBi bias added.
    """
    _, _, seq_len, _ = attention_scores.shape

    # Rebuild cache if needed
    self._build_alibi_tensor(seq_len)

    # Add ALiBi bias
    assert self.alibi is not None
    alibi_bias = self.alibi[:, :, :seq_len, :seq_len].to(attention_scores.device)
    return attention_scores + alibi_bias

get_bias ¶

get_bias(seq_len: int, device: device | None = None) -> Tensor

Get ALiBi bias matrix for a given sequence length.

Parameters:

Name	Type	Description	Default
`seq_len`	`int`	Sequence length.	required
`device`	`device \| None`	Device to place the bias tensor.	`None`

Returns:

Type	Description
`Tensor`	ALiBi bias of shape (1, num_heads, seq_len, seq_len).

Source code in spectrans/models/base.py

def get_bias(self, seq_len: int, device: torch.device | None = None) -> Tensor:
    """Get ALiBi bias matrix for a given sequence length.

    Parameters
    ----------
    seq_len : int
        Sequence length.
    device : torch.device | None, optional
        Device to place the bias tensor.

    Returns
    -------
    Tensor
        ALiBi bias of shape (1, num_heads, seq_len, seq_len).
    """
    self._build_alibi_tensor(seq_len)
    assert self.alibi is not None
    bias = self.alibi[:, :, :seq_len, :seq_len]
    if device is not None:
        bias = bias.to(device)
    return bias

BaseModel ¶

BaseModel(vocab_size: int | None = None, hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, num_classes: int | None = None, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, output_type: OutputHeadType = 'classification', gradient_checkpointing: bool = False)

Bases: SpectralComponent, ABC

Abstract base class for spectral transformer models.

This class provides the common functionality shared by all spectral transformer models, including embeddings, positional encoding, and output heads. Subclasses must implement the build_blocks method to define their specific architecture.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings. If None, no input embedding layer is created (assumes pre-embedded inputs).	`None`
`hidden_dim`	`int`	Hidden dimension size for the model.	`768`
`num_layers`	`int`	Number of transformer blocks in the model.	`12`
`max_sequence_length`	`int`	Maximum sequence length the model can process.	`512`
`num_classes`	`int \| None`	Number of output classes for classification. If None, no classification head is added.	`None`
`use_positional_encoding`	`bool`	Whether to use positional encoding. Default is True.	`True`
`positional_encoding_type`	`PositionalEncodingType`	Type of positional encoding: 'sinusoidal', 'learned', 'rotary', 'alibi', or 'none'. Default is 'sinusoidal'.	`'sinusoidal'`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for feedforward networks. If None, defaults to 4 * hidden_dim.	`None`
`norm_eps`	`float`	Epsilon for layer normalization. Default is 1e-12.	`1e-12`
`output_type`	`OutputHeadType`	Type of output head: 'classification', 'regression', 'sequence', 'lm', or 'none'. Default is 'classification'.	`'classification'`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing for memory efficiency. Default is False.	`False`

Attributes:

Name	Type	Description
`hidden_dim`	`int`	Hidden dimension size.
`num_layers`	`int`	Number of transformer blocks.
`max_sequence_length`	`int`	Maximum sequence length.
`embedding`	`Embedding \| None`	Token embedding layer (if vocab_size is provided).
`positional_encoding`	`PositionalEncoding \| LearnedPositionalEncoding \| None`	Positional encoding module.
`blocks`	`ModuleList`	List of transformer blocks.
`output_head`	`Module \| None`	Task-specific output head.
`dropout`	`Dropout`	Dropout layer.

Methods:

Name	Description
`build_blocks`	Build the transformer blocks for the model.
`forward`	Forward pass through the model.
`from_config`	Create model instance from configuration object.

Source code in spectrans/models/base.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    num_classes: int | None = None,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    output_type: OutputHeadType = "classification",
    gradient_checkpointing: bool = False,
):
    super().__init__()
    self.hidden_dim = hidden_dim
    self.num_layers = num_layers
    self.max_sequence_length = max_sequence_length
    self.ffn_hidden_dim = ffn_hidden_dim or (4 * hidden_dim)
    self.norm_eps = norm_eps
    self.gradient_checkpointing = gradient_checkpointing
    self.output_type = output_type
    self.num_classes = num_classes

    # Input embedding
    if vocab_size is not None:
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        self._initialize_embeddings()
    else:
        self.embedding = None

    # Positional encoding
    self.positional_encoding_type = positional_encoding_type
    if use_positional_encoding and positional_encoding_type != "none":
        if positional_encoding_type == "sinusoidal":
            self.positional_encoding = PositionalEncoding(
                hidden_dim=hidden_dim,
                max_sequence_length=max_sequence_length,
                dropout=dropout,
            )
        elif positional_encoding_type == "learned":
            self.positional_encoding = LearnedPositionalEncoding(
                hidden_dim=hidden_dim,
                max_sequence_length=max_sequence_length,
                dropout=dropout,
            )
        elif positional_encoding_type == "rotary":
            self.positional_encoding = RotaryPositionalEncoding(
                hidden_dim=hidden_dim,
                max_sequence_length=max_sequence_length,
            )
        elif positional_encoding_type == "alibi":
            # ALiBi is handled differently - it's added to attention scores
            # We'll need to pass num_heads, which we'll get from the first block
            self.positional_encoding = None  # Will be initialized after blocks
        else:
            raise ValueError(f"Unknown positional encoding type: {positional_encoding_type}")
    else:
        self.positional_encoding = None

    # Build transformer blocks (implemented by subclasses)
    self.blocks = self.build_blocks()

    # Layer norm before output
    self.final_norm = nn.LayerNorm(hidden_dim, eps=norm_eps)

    # Output head
    if output_type == "classification" and num_classes is not None:
        self.output_head = ClassificationHead(hidden_dim, num_classes, dropout)
    elif output_type == "regression":
        self.output_head = RegressionHead(hidden_dim, dropout)
    elif output_type == "sequence":
        vocab_size_out = (
            num_classes
            if num_classes is not None
            else (vocab_size if vocab_size is not None else hidden_dim)
        )
        self.output_head = SequenceHead(hidden_dim, vocab_size_out, dropout)
    elif output_type == "lm":
        # Language modeling head outputs vocab_size
        if num_classes is not None:
            vocab_size_out = num_classes
        elif vocab_size is not None:
            vocab_size_out = vocab_size
        else:
            raise ValueError(
                "Either num_classes or vocab_size must be specified for language modeling output"
            )
        self.output_head = SequenceHead(hidden_dim, vocab_size_out, dropout)
    else:
        self.output_head = None

    # Dropout
    self.dropout = nn.Dropout(dropout)

Functions¶

build_blocks `abstractmethod` ¶

build_blocks() -> ModuleList

Build the transformer blocks for the model.

This method must be implemented by subclasses to define the specific architecture using appropriate mixing layers.

Returns:

Type	Description
`ModuleList`	List of transformer blocks.

Source code in spectrans/models/base.py

@abstractmethod
def build_blocks(self) -> nn.ModuleList:
    """Build the transformer blocks for the model.

    This method must be implemented by subclasses to define
    the specific architecture using appropriate mixing layers.

    Returns
    -------
    nn.ModuleList
        List of transformer blocks.
    """
    pass

forward ¶

forward(input_ids: Tensor | None = None, inputs_embeds: Tensor | None = None, attention_mask: Tensor | None = None) -> Tensor

Forward pass through the model.

Parameters:

Name	Type	Description	Default
`input_ids`	`Tensor \| None`	Input token IDs of shape (batch_size, sequence_length). Required if embedding layer exists.	`None`
`inputs_embeds`	`Tensor \| None`	Pre-embedded inputs of shape (batch_size, sequence_length, hidden_dim). Used if no embedding layer or to bypass embedding.	`None`
`attention_mask`	`Tensor \| None`	Attention mask of shape (batch_size, sequence_length). Values should be 0 or 1 (1 for tokens to attend to).	`None`

Returns:

Type	Description
`Tensor`	Output tensor. Shape depends on the output head: - Classification: (batch_size, num_classes) - Regression: (batch_size, 1) - Sequence: (batch_size, sequence_length, vocab_size) - None: (batch_size, sequence_length, hidden_dim)

Raises:

Type	Description
`ValueError`	If neither input_ids nor inputs_embeds is provided.

Source code in spectrans/models/base.py

def forward(
    self,
    input_ids: Tensor | None = None,
    inputs_embeds: Tensor | None = None,
    attention_mask: Tensor | None = None,
) -> Tensor:
    """Forward pass through the model.

    Parameters
    ----------
    input_ids : Tensor | None, optional
        Input token IDs of shape (batch_size, sequence_length).
        Required if embedding layer exists.
    inputs_embeds : Tensor | None, optional
        Pre-embedded inputs of shape (batch_size, sequence_length, hidden_dim).
        Used if no embedding layer or to bypass embedding.
    attention_mask : Tensor | None, optional
        Attention mask of shape (batch_size, sequence_length).
        Values should be 0 or 1 (1 for tokens to attend to).

    Returns
    -------
    Tensor
        Output tensor. Shape depends on the output head:
        - Classification: (batch_size, num_classes)
        - Regression: (batch_size, 1)
        - Sequence: (batch_size, sequence_length, vocab_size)
        - None: (batch_size, sequence_length, hidden_dim)

    Raises
    ------
    ValueError
        If neither input_ids nor inputs_embeds is provided.
    """
    # Get embeddings
    if inputs_embeds is not None:
        hidden_states = inputs_embeds
    elif input_ids is not None and self.embedding is not None:
        hidden_states = self.embedding(input_ids)
    elif input_ids is not None:
        raise ValueError("Model has no embedding layer but input_ids was provided")
    else:
        raise ValueError("Either input_ids or inputs_embeds must be provided")

    # Add positional encoding
    if self.positional_encoding is not None:
        hidden_states = self.positional_encoding(hidden_states)

    # Apply dropout
    hidden_states = self.dropout(hidden_states)

    # Process through transformer blocks
    for block in self.blocks:
        if self.gradient_checkpointing and self.training:
            hidden_states = checkpoint.checkpoint(block, hidden_states, use_reentrant=False)
        else:
            hidden_states = block(hidden_states)

    # Final layer norm
    hidden_states = self.final_norm(hidden_states)

    # Apply output head if present
    if self.output_head is not None:
        output: Tensor = self.output_head(hidden_states, attention_mask)
    else:
        output = hidden_states

    return output

from_config `classmethod` ¶

from_config(config: ModelConfig) -> BaseModel

Create model instance from configuration object.

Parameters:

Name	Type	Description	Default
`config`	`ModelConfig`	Configuration object with model parameters.	required

Returns:

Type	Description
`BaseModel`	Configured model instance.

Source code in spectrans/models/base.py

@classmethod
def from_config(cls, config: "ModelConfig") -> "BaseModel":
    """Create model instance from configuration object.

    Parameters
    ----------
    config : ModelConfig
        Configuration object with model parameters.

    Returns
    -------
    BaseModel
        Configured model instance.
    """
    # Build model directly from config attributes
    return cls(
        vocab_size=getattr(config, "vocab_size", None),
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        num_classes=getattr(config, "num_classes", None),
        use_positional_encoding=getattr(config, "use_positional_encoding", True),
        positional_encoding_type=getattr(config, "positional_encoding_type", "sinusoidal"),
        dropout=config.dropout,
        ffn_hidden_dim=getattr(config, "ffn_hidden_dim", None),
        norm_eps=getattr(config, "norm_eps", 1e-12),
        output_type=getattr(config, "output_type", "classification"),
        gradient_checkpointing=getattr(config, "gradient_checkpointing", False),
    )

ClassificationHead ¶

ClassificationHead(hidden_dim: int, num_classes: int, dropout: float = 0.1, pooling: PoolingType = 'cls')

Bases: Module

Classification output head.

This module pools sequence outputs and projects to class logits.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Input hidden dimension.	required
`num_classes`	`int`	Number of output classes.	required
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`pooling`	`PoolingType`	Pooling strategy: 'cls', 'mean', or 'max'. Default is 'cls'.	`'cls'`

Attributes:

Name	Type	Description
`pooling`	`PoolingType`	Pooling strategy.
`dropout`	`Dropout`	Dropout layer.
`classifier`	`Linear`	Output projection layer.

Methods:

Name	Description
`forward`	Forward pass through classification head.

Source code in spectrans/models/base.py

def __init__(
    self,
    hidden_dim: int,
    num_classes: int,
    dropout: float = 0.1,
    pooling: PoolingType = "cls",
):
    super().__init__()
    self.pooling = pooling
    self.dropout = nn.Dropout(dropout)
    self.classifier = nn.Linear(hidden_dim, num_classes)

Functions¶

forward ¶

forward(hidden_states: Tensor, attention_mask: Tensor | None = None) -> Tensor

Forward pass through classification head.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input tensor of shape (batch_size, sequence_length, hidden_dim).	required
`attention_mask`	`Tensor \| None`	Attention mask for pooling operations.	`None`

Returns:

Type	Description
`Tensor`	Classification logits of shape (batch_size, num_classes).

Source code in spectrans/models/base.py

def forward(
    self,
    hidden_states: Tensor,
    attention_mask: Tensor | None = None,
) -> Tensor:
    """Forward pass through classification head.

    Parameters
    ----------
    hidden_states : Tensor
        Input tensor of shape (batch_size, sequence_length, hidden_dim).
    attention_mask : Tensor | None, optional
        Attention mask for pooling operations.

    Returns
    -------
    Tensor
        Classification logits of shape (batch_size, num_classes).
    """
    # Pool the sequence
    if self.pooling == "cls":
        # Use first token (CLS token)
        pooled = hidden_states[:, 0, :]
    elif self.pooling == "mean":
        # Mean pooling
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
        else:
            pooled = hidden_states.mean(dim=1)
    elif self.pooling == "max":
        # Max pooling
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            hidden_states = hidden_states.masked_fill(mask == 0, -1e9)
        pooled = hidden_states.max(dim=1)[0]
    else:
        raise ValueError(f"Unknown pooling strategy: {self.pooling}")

    # Apply dropout and classifier
    pooled = self.dropout(pooled)
    logits: Tensor = self.classifier(pooled)
    return logits

LearnedPositionalEncoding ¶

LearnedPositionalEncoding(hidden_dim: int, max_sequence_length: int = 5000, dropout: float = 0.1)

Bases: Module

Learned positional embeddings.

This module uses learnable positional embeddings instead of fixed sinusoidal encodings.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Dimension of the embeddings.	required
`max_sequence_length`	`int`	Maximum sequence length to encode.	`5000`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`

Attributes:

Name	Type	Description
`position_embeddings`	`Embedding`	Learnable position embeddings.
`dropout`	`Dropout`	Dropout layer.

Methods:

Name	Description
`forward`	Add learned positional embeddings to input tensor.

Source code in spectrans/models/base.py

def __init__(
    self,
    hidden_dim: int,
    max_sequence_length: int = 5000,
    dropout: float = 0.1,
):
    super().__init__()
    self.position_embeddings = nn.Embedding(max_sequence_length, hidden_dim)
    self.dropout = nn.Dropout(p=dropout)

    # Initialize embeddings
    nn.init.xavier_uniform_(self.position_embeddings.weight)

Functions¶

forward ¶

forward(x: Tensor) -> Tensor

Add learned positional embeddings to input tensor.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor of shape (batch_size, sequence_length, hidden_dim).	required

Returns:

Type	Description
`Tensor`	Tensor with positional embeddings added.

Source code in spectrans/models/base.py

def forward(self, x: Tensor) -> Tensor:
    """Add learned positional embeddings to input tensor.

    Parameters
    ----------
    x : Tensor
        Input tensor of shape (batch_size, sequence_length, hidden_dim).

    Returns
    -------
    Tensor
        Tensor with positional embeddings added.
    """
    seq_length = x.size(1)
    position_ids = torch.arange(seq_length, dtype=torch.long, device=x.device)
    position_ids = position_ids.unsqueeze(0).expand(x.size(0), -1)

    position_embeddings = self.position_embeddings(position_ids)
    x = x + position_embeddings
    result: Tensor = self.dropout(x)
    return result

PositionalEncoding ¶

PositionalEncoding(hidden_dim: int, max_sequence_length: int = 5000, dropout: float = 0.1)

Bases: Module

Sinusoidal positional encoding.

This module adds sinusoidal positional encodings to embeddings, following the approach in "Attention is All You Need".

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Dimension of the embeddings.	required
`max_sequence_length`	`int`	Maximum sequence length to encode.	`5000`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`

Attributes:

Name	Type	Description
`dropout`	`Dropout`	Dropout layer.
`pe`	`Tensor`	Precomputed positional encodings.

Methods:

Name	Description
`forward`	Add positional encoding to input tensor.

Source code in spectrans/models/base.py

def __init__(
    self,
    hidden_dim: int,
    max_sequence_length: int = 5000,
    dropout: float = 0.1,
):
    super().__init__()
    self.dropout = nn.Dropout(p=dropout)
    self.pe: Tensor  # Type annotation for buffer

    # Create positional encodings
    pe = torch.zeros(max_sequence_length, hidden_dim)
    position = torch.arange(0, max_sequence_length).unsqueeze(1).float()

    # Create div_term for the sinusoidal pattern
    div_term = torch.exp(
        torch.arange(0, hidden_dim, 2).float() * -(math.log(10000.0) / hidden_dim)
    )

    # Apply sinusoidal functions
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)

    # Register as buffer (not a parameter)
    self.register_buffer("pe", pe.unsqueeze(0))

Functions¶

forward ¶

forward(x: Tensor) -> Tensor

Add positional encoding to input tensor.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor of shape (batch_size, sequence_length, hidden_dim).	required

Returns:

Type	Description
`Tensor`	Tensor with positional encoding added.

Source code in spectrans/models/base.py

def forward(self, x: Tensor) -> Tensor:
    """Add positional encoding to input tensor.

    Parameters
    ----------
    x : Tensor
        Input tensor of shape (batch_size, sequence_length, hidden_dim).

    Returns
    -------
    Tensor
        Tensor with positional encoding added.
    """
    # Add positional encoding
    seq_length = x.size(1)
    x = x + self.pe[:, :seq_length, :]
    result: Tensor = self.dropout(x)
    return result

RegressionHead ¶

RegressionHead(hidden_dim: int, dropout: float = 0.1, pooling: PoolingType = 'mean')

Bases: Module

Regression output head.

This module pools sequence outputs and projects to a scalar value.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Input hidden dimension.	required
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`pooling`	`PoolingType`	Pooling strategy: 'cls', 'mean', or 'max'. Default is 'mean'.	`'mean'`

Attributes:

Name	Type	Description
`pooling`	`PoolingType`	Pooling strategy.
`dropout`	`Dropout`	Dropout layer.
`regressor`	`Linear`	Output projection layer.

Methods:

Name	Description
`forward`	Forward pass through regression head.

Source code in spectrans/models/base.py

def __init__(
    self,
    hidden_dim: int,
    dropout: float = 0.1,
    pooling: PoolingType = "mean",
):
    super().__init__()
    self.pooling = pooling
    self.dropout = nn.Dropout(dropout)
    self.regressor = nn.Linear(hidden_dim, 1)

Functions¶

forward ¶

forward(hidden_states: Tensor, attention_mask: Tensor | None = None) -> Tensor

Forward pass through regression head.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input tensor of shape (batch_size, sequence_length, hidden_dim).	required
`attention_mask`	`Tensor \| None`	Attention mask for pooling operations.	`None`

Returns:

Type	Description
`Tensor`	Regression output of shape (batch_size, 1).

Source code in spectrans/models/base.py

def forward(
    self,
    hidden_states: Tensor,
    attention_mask: Tensor | None = None,
) -> Tensor:
    """Forward pass through regression head.

    Parameters
    ----------
    hidden_states : Tensor
        Input tensor of shape (batch_size, sequence_length, hidden_dim).
    attention_mask : Tensor | None, optional
        Attention mask for pooling operations.

    Returns
    -------
    Tensor
        Regression output of shape (batch_size, 1).
    """
    # Pool the sequence (same logic as classification)
    if self.pooling == "cls":
        pooled = hidden_states[:, 0, :]
    elif self.pooling == "mean":
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            pooled = (hidden_states * mask).sum(dim=1) / mask.sum(dim=1)
        else:
            pooled = hidden_states.mean(dim=1)
    elif self.pooling == "max":
        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            hidden_states = hidden_states.masked_fill(mask == 0, -1e9)
        pooled = hidden_states.max(dim=1)[0]
    else:
        raise ValueError(f"Unknown pooling strategy: {self.pooling}")

    # Apply dropout and regressor
    pooled = self.dropout(pooled)
    output: Tensor = self.regressor(pooled)
    return output

RotaryPositionalEncoding ¶

RotaryPositionalEncoding(hidden_dim: int, max_sequence_length: int = 5000, base: float = 10000.0)

Bases: Module

Rotary Position Embedding (RoPE).

This module implements Rotary Position Embeddings as described in the RoFormer paper. RoPE encodes absolute position with rotation matrix and naturally incorporates relative position dependency in self-attention formulation.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Dimension of the embeddings. Must be even.	required
`max_sequence_length`	`int`	Maximum sequence length to encode.	`5000`
`base`	`float`	Base for the frequency calculation. Default is 10000.	`10000.0`

Attributes:

Name	Type	Description
`inv_freq`	`Tensor`	Inverse frequencies for computing rotary embeddings.
`cos_cached`	`Tensor \| None`	Cached cosine values for positions.
`sin_cached`	`Tensor \| None`	Cached sine values for positions.

References

Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. 2024. RoFormer: Enhanced transformer with rotary position embedding. Neurocomputing, 568:127063.

Methods:

Name	Description
`forward`	Apply rotary position embedding to input tensor.

Source code in spectrans/models/base.py

def __init__(
    self,
    hidden_dim: int,
    max_sequence_length: int = 5000,
    base: float = 10000.0,
):
    super().__init__()
    if hidden_dim % 2 != 0:
        raise ValueError(f"hidden_dim must be even for RoPE, got {hidden_dim}")

    self.hidden_dim = hidden_dim
    self.max_sequence_length = max_sequence_length
    self.base = base

    # Compute inverse frequencies
    inv_freq = 1.0 / (base ** (torch.arange(0, hidden_dim, 2).float() / hidden_dim))
    self.register_buffer("inv_freq", inv_freq)

    # Cache for precomputed cos/sin
    self.cos_cached: Tensor | None = None
    self.sin_cached: Tensor | None = None
    self._build_cache(max_sequence_length)

Functions¶

forward ¶

forward(x: Tensor, offset: int = 0) -> Tensor

Apply rotary position embedding to input tensor.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor of shape (batch_size, num_heads, sequence_length, head_dim) or (batch_size, sequence_length, hidden_dim).	required
`offset`	`int`	Position offset for incremental decoding. Default is 0.	`0`

Returns:

Type	Description
`Tensor`	Tensor with rotary position embeddings applied.

Source code in spectrans/models/base.py

def forward(self, x: Tensor, offset: int = 0) -> Tensor:
    """Apply rotary position embedding to input tensor.

    Parameters
    ----------
    x : Tensor
        Input tensor of shape (batch_size, num_heads, sequence_length, head_dim)
        or (batch_size, sequence_length, hidden_dim).
    offset : int, optional
        Position offset for incremental decoding. Default is 0.

    Returns
    -------
    Tensor
        Tensor with rotary position embeddings applied.
    """
    # Handle both 3D and 4D inputs
    if x.ndim == 3:
        _, seq_len, _ = x.shape
        # For simplicity, we'll apply RoPE directly to the hidden dimension
        # In practice, this would be applied separately to Q and K in attention
        was_3d = True
    else:
        x.shape[0]
        seq_len = x.shape[2] if x.ndim == 4 else x.shape[1]
        was_3d = False

    # Rebuild cache if needed
    self._build_cache(seq_len + offset)

    if was_3d:
        # For 3D tensor, apply rotation directly
        assert self.cos_cached is not None
        assert self.sin_cached is not None
        cos = self.cos_cached[:, :, offset : offset + seq_len, :].squeeze(1)
        sin = self.sin_cached[:, :, offset : offset + seq_len, :].squeeze(1)

        # Split x into two halves for rotation
        x1, x2 = x.chunk(2, dim=-1)

        # Apply rotation
        rotated = torch.cat(
            [
                x1 * cos[:, :, : x1.shape[-1]] - x2 * sin[:, :, : x2.shape[-1]],
                x1 * sin[:, :, : x1.shape[-1]] + x2 * cos[:, :, : x2.shape[-1]],
            ],
            dim=-1,
        )
    else:
        # For 4D tensor (batch, heads, seq, head_dim)
        assert self.cos_cached is not None
        assert self.sin_cached is not None
        cos = self.cos_cached[:, :, offset : offset + seq_len, :]
        sin = self.sin_cached[:, :, offset : offset + seq_len, :]

        # Split x into two halves for rotation
        x1, x2 = x.chunk(2, dim=-1)

        # Apply rotation
        rotated = torch.cat(
            [
                x1 * cos[:, :, :, : x1.shape[-1]] - x2 * sin[:, :, :, : x2.shape[-1]],
                x1 * sin[:, :, :, : x1.shape[-1]] + x2 * cos[:, :, :, : x2.shape[-1]],
            ],
            dim=-1,
        )

    return rotated

SequenceHead ¶

SequenceHead(hidden_dim: int, vocab_size: int, dropout: float = 0.1)

Bases: Module

Sequence-to-sequence output head.

This module projects hidden states to vocabulary logits for sequence generation tasks.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Input hidden dimension.	required
`vocab_size`	`int`	Output vocabulary size.	required
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`

Attributes:

Name	Type	Description
`dropout`	`Dropout`	Dropout layer.
`lm_head`	`Linear`	Language modeling head.

Methods:

Name	Description
`forward`	Forward pass through sequence head.

Source code in spectrans/models/base.py

def __init__(
    self,
    hidden_dim: int,
    vocab_size: int,
    dropout: float = 0.1,
):
    super().__init__()
    self.dropout = nn.Dropout(dropout)
    self.lm_head = nn.Linear(hidden_dim, vocab_size)

Functions¶

forward ¶

forward(hidden_states: Tensor, attention_mask: Tensor | None = None) -> Tensor

Forward pass through sequence head.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Input tensor of shape (batch_size, sequence_length, hidden_dim).	required
`attention_mask`	`Tensor \| None`	Not used, kept for interface consistency.	`None`

Returns:

Type	Description
`Tensor`	Vocabulary logits of shape (batch_size, sequence_length, vocab_size).

Source code in spectrans/models/base.py

def forward(
    self,
    hidden_states: Tensor,
    attention_mask: Tensor | None = None,  # noqa: ARG002
) -> Tensor:
    """Forward pass through sequence head.

    Parameters
    ----------
    hidden_states : Tensor
        Input tensor of shape (batch_size, sequence_length, hidden_dim).
    attention_mask : Tensor | None, optional
        Not used, kept for interface consistency.

    Returns
    -------
    Tensor
        Vocabulary logits of shape (batch_size, sequence_length, vocab_size).
    """
    hidden_states = self.dropout(hidden_states)
    logits: Tensor = self.lm_head(hidden_states)
    return logits

FNet ¶

FNet(vocab_size: int | None = None, hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, num_classes: int | None = None, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, output_type: OutputHeadType = 'classification', use_real_fft: bool = True, gradient_checkpointing: bool = False)

Bases: BaseModel

FNet model with Fourier transform-based token mixing.

FNet replaces the self-attention mechanism with Fourier transforms, achieving \(O(n \log n)\) complexity.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings. If None, expects pre-embedded inputs.	`None`
`hidden_dim`	`int`	Hidden dimension size. Default is 768.	`768`
`num_layers`	`int`	Number of FNet layers. Default is 12.	`12`
`max_sequence_length`	`int`	Maximum sequence length. Default is 512.	`512`
`num_classes`	`int \| None`	Number of output classes for classification. Default is None.	`None`
`use_positional_encoding`	`bool`	Whether to use positional encoding. Default is True.	`True`
`positional_encoding_type`	`str`	Type of positional encoding: 'sinusoidal' or 'learned'. Default is 'sinusoidal'.	`'sinusoidal'`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN. If None, defaults to 4 * hidden_dim.	`None`
`norm_eps`	`float`	Epsilon for layer normalization. Default is 1e-12.	`1e-12`
`output_type`	`str`	Type of output head: 'classification', 'regression', 'sequence', or 'none'. Default is 'classification'.	`'classification'`
`use_real_fft`	`bool`	Whether to use real FFT for efficiency. Default is True.	`True`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing. Default is False.	`False`

Attributes:

Name	Type	Description
`use_real_fft`	`bool`	Whether real FFT is used for efficiency.
`blocks`	`ModuleList`	List of FNet transformer blocks.

Methods:

Name	Description
`build_blocks`	Build FNet transformer blocks with Fourier mixing.
`from_config`	Create FNet model from configuration.

Source code in spectrans/models/fnet.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    num_classes: int | None = None,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    output_type: OutputHeadType = "classification",
    use_real_fft: bool = True,
    gradient_checkpointing: bool = False,
):
    self.use_real_fft = use_real_fft
    self._dropout_rate = dropout  # Store dropout rate for build_blocks

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type=output_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build FNet transformer blocks with Fourier mixing.

Returns:

Type	Description
`ModuleList`	List of FNet transformer blocks.

Source code in spectrans/models/fnet.py

def build_blocks(self) -> nn.ModuleList:
    """Build FNet transformer blocks with Fourier mixing.

    Returns
    -------
    nn.ModuleList
        List of FNet transformer blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Choose mixing layer based on use_real_fft flag
        mixing_layer: FourierMixing | RealFourierMixing
        if self.use_real_fft:
            mixing_layer = RealFourierMixing(
                hidden_dim=self.hidden_dim,
                dropout=self._dropout_rate,
            )
        else:
            mixing_layer = FourierMixing(
                hidden_dim=self.hidden_dim,
                dropout=self._dropout_rate,
            )

        # Create FNet block with pre-normalization
        block = PreNormBlock(
            mixing_layer=mixing_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            activation="gelu",
            dropout=self._dropout_rate,
            norm_eps=self.norm_eps,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

from_config `classmethod` ¶

from_config(config: FNetModelConfig) -> FNet

Create FNet model from configuration.

Parameters:

Name	Type	Description	Default
`config`	`FNetModelConfig`	Configuration object with model parameters.	required

Returns:

Type	Description
`FNet`	Configured FNet model.

Source code in spectrans/models/fnet.py

@classmethod
def from_config(cls, config: "FNetModelConfig") -> "FNet":  # type: ignore[override]
    """Create FNet model from configuration.

    Parameters
    ----------
    config : FNetModelConfig
        Configuration object with model parameters.

    Returns
    -------
    FNet
        Configured FNet model.
    """
    return cls(
        vocab_size=getattr(config, "vocab_size", None),
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        num_classes=getattr(config, "num_classes", None),
        use_positional_encoding=getattr(config, "use_positional_encoding", True),
        positional_encoding_type=getattr(config, "positional_encoding_type", "sinusoidal"),
        dropout=config.dropout,
        ffn_hidden_dim=getattr(config, "ffn_hidden_dim", None),
        norm_eps=getattr(config, "norm_eps", 1e-12),
        output_type=getattr(config, "output_type", "classification"),
        use_real_fft=getattr(config, "use_real_fft", True),
        gradient_checkpointing=getattr(config, "gradient_checkpointing", False),
    )

FNetEncoder ¶

FNetEncoder(hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, use_real_fft: bool = True, gradient_checkpointing: bool = False)

Bases: FNet

Encoder-only FNet model for representation learning.

This variant of FNet is designed for tasks that require extracting representations rather than making predictions. It returns the hidden states from the final layer without any task-specific head.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden dimension size. Default is 768.	`768`
`num_layers`	`int`	Number of FNet layers. Default is 12.	`12`
`max_sequence_length`	`int`	Maximum sequence length. Default is 512.	`512`
`use_positional_encoding`	`bool`	Whether to use positional encoding. Default is True.	`True`
`positional_encoding_type`	`str`	Type of positional encoding. Default is 'sinusoidal'.	`'sinusoidal'`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN. If None, defaults to 4 * hidden_dim.	`None`
`norm_eps`	`float`	Epsilon for layer normalization. Default is 1e-12.	`1e-12`
`use_real_fft`	`bool`	Whether to use real FFT. Default is True.	`True`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing. Default is False.	`False`

Source code in spectrans/models/fnet.py

def __init__(
    self,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    use_real_fft: bool = True,
    gradient_checkpointing: bool = False,
):
    super().__init__(
        vocab_size=None,  # No token embeddings for encoder
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=None,  # No classification head
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type="none",  # Return hidden states
        use_real_fft=use_real_fft,
        gradient_checkpointing=gradient_checkpointing,
    )

FNODecoder ¶

FNODecoder(vocab_size: int, hidden_dim: int = 512, num_layers: int = 12, max_sequence_length: int = 2048, modes: int = 32, mlp_ratio: float = 2.0, causal: bool = True, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', gradient_checkpointing: bool = False)

Bases: BaseModel

Decoder FNO model for generation tasks.

This model uses causal FNO blocks suitable for autoregressive generation tasks. The spectral operations are modified to respect causality.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int`	Size of the vocabulary for generation.	required
`hidden_dim`	`int`	Hidden dimension size.	`512`
`num_layers`	`int`	Number of decoder blocks.	`12`
`max_sequence_length`	`int`	Maximum sequence length.	`2048`
`modes`	`int`	Number of Fourier modes (adjusted for causality).	`32`
`mlp_ratio`	`float`	MLP expansion ratio.	`2.0`
`causal`	`bool`	Whether to use causal masking.	`True`
`ffn_hidden_dim`	`int \| None`	Hidden dimension of the feedforward network.	`None`
`dropout`	`float`	Dropout probability.	`0.0`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`str`	Type of positional encoding.	`"sinusoidal"`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing.	`False`

Examples:

>>> decoder = FNODecoder(
...     vocab_size=10000,
...     hidden_dim=512,
...     num_layers=12,
...     modes=32,
...     causal=True,
...     max_sequence_length=2048
... )
>>> input_ids = torch.randint(0, 10000, (32, 100))
>>> logits = decoder(input_ids)
>>> assert logits.shape == (32, 100, 10000)

Methods:

Name	Description
`build_blocks`	Build decoder blocks with causal FNO layers.
`forward`	Forward pass through the decoder.

Source code in spectrans/models/fno_transformer.py

def __init__(
    self,
    vocab_size: int,
    hidden_dim: int = 512,
    num_layers: int = 12,
    max_sequence_length: int = 2048,
    modes: int = 32,
    mlp_ratio: float = 2.0,
    causal: bool = True,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    gradient_checkpointing: bool = False,
):
    # Store FNO-specific parameters
    self.modes = modes
    self.mlp_ratio = mlp_ratio
    self.causal = causal
    self.dropout_rate = dropout

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=None,  # Decoder uses LM head instead
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=gradient_checkpointing,
    )

    # Add language modeling head
    self.lm_head = nn.Linear(hidden_dim, vocab_size)
    self.output_type = "lm"

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build decoder blocks with causal FNO layers.

Returns:

Type	Description
`ModuleList`	List of causal FNO decoder blocks.

Source code in spectrans/models/fno_transformer.py

def build_blocks(self) -> nn.ModuleList:
    """Build decoder blocks with causal FNO layers.

    Returns
    -------
    nn.ModuleList
        List of causal FNO decoder blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Create FNO block
        # Note: Causality in spectral domain requires special handling
        # This is a simplified version - full causality would need custom implementation
        fno_block = FNOBlock(
            hidden_dim=self.hidden_dim,
            modes=self.modes,
            mlp_ratio=self.mlp_ratio,
            dropout=self.dropout_rate,
        )
        blocks.append(fno_block)

    return nn.ModuleList(blocks)

forward ¶

forward(input_ids: Tensor | None = None, inputs_embeds: Tensor | None = None, attention_mask: Tensor | None = None) -> Tensor

Forward pass through the decoder.

Parameters:

Name	Type	Description	Default
`input_ids`	`Tensor \| None`	Input token IDs of shape (batch_size, sequence_length).	`None`
`inputs_embeds`	`Tensor \| None`	Pre-embedded inputs of shape (batch_size, sequence_length, hidden_dim).	`None`
`attention_mask`	`Tensor \| None`	Attention mask for padding.	`None`

Returns:

Type	Description
`Tensor`	Logits of shape (batch_size, sequence_length, vocab_size).

Source code in spectrans/models/fno_transformer.py

def forward(
    self,
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Forward pass through the decoder.

    Parameters
    ----------
    input_ids : torch.Tensor | None, optional
        Input token IDs of shape (batch_size, sequence_length).
    inputs_embeds : torch.Tensor | None, optional
        Pre-embedded inputs of shape (batch_size, sequence_length, hidden_dim).
    attention_mask : torch.Tensor | None, optional
        Attention mask for padding.

    Returns
    -------
    torch.Tensor
        Logits of shape (batch_size, sequence_length, vocab_size).
    """
    # Use parent class forward for processing
    hidden_states = super().forward(
        input_ids=input_ids,
        inputs_embeds=inputs_embeds,
        attention_mask=attention_mask,
    )

    # Apply LM head
    logits = self.lm_head(hidden_states)
    return logits  # type: ignore[no-any-return]

FNOEncoder ¶

FNOEncoder(hidden_dim: int = 512, num_layers: int = 6, max_sequence_length: int = 1024, modes: int = 32, mlp_ratio: float = 2.0, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', gradient_checkpointing: bool = False)

Bases: BaseModel

Encoder-only FNO model for representation learning.

This model uses stacked FNO blocks without causal masking, suitable for bidirectional encoding tasks like feature extraction and representation learning.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden dimension size for the model.	`512`
`num_layers`	`int`	Number of encoder blocks.	`6`
`max_sequence_length`	`int`	Maximum sequence length.	`1024`
`modes`	`int`	Number of Fourier modes to retain.	`32`
`mlp_ratio`	`float`	MLP expansion ratio in FNO blocks.	`2.0`
`ffn_hidden_dim`	`int \| None`	Hidden dimension of the feedforward network.	`None`
`dropout`	`float`	Dropout probability.	`0.0`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`str`	Type of positional encoding.	`"sinusoidal"`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing.	`False`

Examples:

>>> encoder = FNOEncoder(
...     hidden_dim=512,
...     num_layers=6,
...     modes=32,
...     max_sequence_length=1024
... )
>>> x = torch.randn(32, 100, 512)
>>> encoded = encoder(inputs_embeds=x)
>>> assert encoded.shape == x.shape

Methods:

Name	Description
`build_blocks`	Build encoder blocks with FNO layers.

Source code in spectrans/models/fno_transformer.py

def __init__(
    self,
    hidden_dim: int = 512,
    num_layers: int = 6,
    max_sequence_length: int = 1024,
    modes: int = 32,
    mlp_ratio: float = 2.0,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    gradient_checkpointing: bool = False,
):
    # Store FNO-specific parameters
    self.modes = modes
    self.mlp_ratio = mlp_ratio
    self.dropout_rate = dropout

    super().__init__(
        vocab_size=None,  # Encoder doesn't need vocab
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=None,  # No classification head
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=gradient_checkpointing,
    )

    # Set output type to none for encoder
    self.output_type = "none"

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build encoder blocks with FNO layers.

Returns:

Type	Description
`ModuleList`	List of FNO encoder blocks.

Source code in spectrans/models/fno_transformer.py

def build_blocks(self) -> nn.ModuleList:
    """Build encoder blocks with FNO layers.

    Returns
    -------
    nn.ModuleList
        List of FNO encoder blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        fno_block = FNOBlock(
            hidden_dim=self.hidden_dim,
            modes=self.modes,
            mlp_ratio=self.mlp_ratio,
            dropout=self.dropout_rate,
        )
        blocks.append(fno_block)

    return nn.ModuleList(blocks)

FNOTransformer ¶

FNOTransformer(vocab_size: int | None = None, hidden_dim: int = 512, num_layers: int = 6, max_sequence_length: int = 1024, modes: int = 32, mlp_ratio: float = 2.0, use_2d: bool = False, spatial_dim: int | None = None, num_classes: int | None = None, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', gradient_checkpointing: bool = False)

Bases: BaseModel

Fourier Neural Operator transformer model.

This model uses Fourier Neural Operators for sequence mixing, achieving O(n log n) complexity through FFT operations. The model learns mappings between function spaces by parameterizing kernels in the Fourier domain.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings. If None, expects pre-embedded inputs.	`None`
`hidden_dim`	`int`	Hidden dimension size for the model.	`512`
`num_layers`	`int`	Number of transformer blocks.	`6`
`max_sequence_length`	`int`	Maximum sequence length the model can process.	`1024`
`modes`	`int`	Number of Fourier modes to retain (frequency truncation).	`32`
`mlp_ratio`	`float`	Expansion ratio for the MLP in FNO blocks.	`2.0`
`use_2d`	`bool`	Whether to use 2D spectral convolutions for spatial data.	`False`
`spatial_dim`	`int \| None`	Spatial dimension when using 2D convolutions (sequence = spatial_dim²).	`None`
`num_classes`	`int \| None`	Number of output classes for classification.	`None`
`ffn_hidden_dim`	`int \| None`	Hidden dimension of the feedforward network. Default is 4 * hidden_dim.	`None`
`dropout`	`float`	Dropout probability.	`0.0`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`str`	Type of positional encoding ("sinusoidal" or "learned").	`"sinusoidal"`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing to save memory.	`False`

Attributes:

Name	Type	Description
`blocks`	`ModuleList`	Stack of FNO transformer blocks.

Examples:

>>> model = FNOTransformer(
...     hidden_dim=512,
...     num_layers=6,
...     modes=32,
...     max_sequence_length=1024
... )
>>> x = torch.randn(32, 100, 512)
>>> output = model(inputs_embeds=x)
>>> assert output.shape == x.shape

Methods:

Name	Description
`build_blocks`	Build transformer blocks with FNO layers.
`from_config`	Create model from configuration.

Source code in spectrans/models/fno_transformer.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 512,
    num_layers: int = 6,
    max_sequence_length: int = 1024,
    modes: int = 32,
    mlp_ratio: float = 2.0,
    use_2d: bool = False,
    spatial_dim: int | None = None,
    num_classes: int | None = None,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    gradient_checkpointing: bool = False,
):
    # Store FNO-specific parameters
    self.modes = modes
    self.mlp_ratio = mlp_ratio
    self.use_2d = use_2d
    self.spatial_dim = spatial_dim
    self.dropout_rate = dropout

    # Validate 2D configuration
    if use_2d and spatial_dim is None:
        raise ValueError("spatial_dim must be specified when use_2d=True")
    if use_2d and spatial_dim is not None and spatial_dim * spatial_dim != max_sequence_length:
        raise ValueError(
            f"For 2D FNO, max_sequence_length ({max_sequence_length}) "
            f"must equal spatial_dim² ({spatial_dim}² = {spatial_dim * spatial_dim})"
        )

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build transformer blocks with FNO layers.

Returns:

Type	Description
`ModuleList`	List of FNO transformer blocks.

Source code in spectrans/models/fno_transformer.py

def build_blocks(self) -> nn.ModuleList:
    """Build transformer blocks with FNO layers.

    Returns
    -------
    nn.ModuleList
        List of FNO transformer blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Create FNO block with appropriate configuration
        fno_block = FNOBlock(
            hidden_dim=self.hidden_dim,
            modes=self.modes,
            mlp_ratio=self.mlp_ratio,
            dropout=self.dropout_rate,
        )
        blocks.append(fno_block)

    return nn.ModuleList(blocks)

from_config `classmethod` ¶

from_config(config: FNOTransformerConfig) -> FNOTransformer

Create model from configuration.

Parameters:

Name	Type	Description	Default
`config`	`FNOTransformerConfig`	Model configuration object.	required

Returns:

Type	Description
`FNOTransformer`	Instantiated model.

Source code in spectrans/models/fno_transformer.py

@classmethod
def from_config(cls, config: "FNOTransformerConfig") -> "FNOTransformer":  # type: ignore[override]
    """Create model from configuration.

    Parameters
    ----------
    config : FNOTransformerConfig
        Model configuration object.

    Returns
    -------
    FNOTransformer
        Instantiated model.
    """
    return cls(
        vocab_size=config.vocab_size,
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        modes=config.modes,
        mlp_ratio=config.mlp_ratio,
        use_2d=config.use_2d,
        spatial_dim=config.spatial_dim,
        num_classes=config.num_classes,
        ffn_hidden_dim=config.ffn_hidden_dim,
        dropout=config.dropout,
        use_positional_encoding=config.use_positional_encoding,
        positional_encoding_type=config.positional_encoding_type,
        gradient_checkpointing=config.gradient_checkpointing,
    )

GFNet ¶

GFNet(vocab_size: int | None = None, hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, num_classes: int | None = None, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, output_type: OutputHeadType = 'classification', filter_activation: FilterActivationType = 'sigmoid', gradient_checkpointing: bool = False)

Bases: BaseModel

Global Filter Network model with learnable frequency domain filters.

GFNet uses learnable complex filters in the Fourier domain for token mixing, maintaining \(O(n \log n)\) complexity.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings. If None, expects pre-embedded inputs.	`None`
`hidden_dim`	`int`	Hidden dimension size. Default is 768.	`768`
`num_layers`	`int`	Number of GFNet layers. Default is 12.	`12`
`max_sequence_length`	`int`	Maximum sequence length. Default is 512.	`512`
`num_classes`	`int \| None`	Number of output classes for classification. Default is None.	`None`
`use_positional_encoding`	`bool`	Whether to use positional encoding. Default is True.	`True`
`positional_encoding_type`	`str`	Type of positional encoding: 'sinusoidal' or 'learned'. Default is 'sinusoidal'.	`'sinusoidal'`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN. If None, defaults to 4 * hidden_dim.	`None`
`norm_eps`	`float`	Epsilon for layer normalization. Default is 1e-12.	`1e-12`
`output_type`	`str`	Type of output head: 'classification', 'regression', 'sequence', or 'none'. Default is 'classification'.	`'classification'`
`filter_activation`	`str`	Activation function for filters: 'sigmoid' or 'tanh'. Default is 'sigmoid'.	`'sigmoid'`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing. Default is False.	`False`

Attributes:

Name	Type	Description
`filter_activation`	`str`	Activation function used for filters.
`blocks`	`ModuleList`	List of GFNet transformer blocks.

Methods:

Name	Description
`build_blocks`	Build GFNet transformer blocks with global filter mixing.
`from_config`	Create GFNet model from configuration.

Source code in spectrans/models/gfnet.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    num_classes: int | None = None,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    output_type: OutputHeadType = "classification",
    filter_activation: FilterActivationType = "sigmoid",
    gradient_checkpointing: bool = False,
):
    self.filter_activation = filter_activation
    self._dropout_rate = dropout  # Store dropout rate for build_blocks

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type=output_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build GFNet transformer blocks with global filter mixing.

Returns:

Type	Description
`ModuleList`	List of GFNet transformer blocks.

Source code in spectrans/models/gfnet.py

def build_blocks(self) -> nn.ModuleList:
    """Build GFNet transformer blocks with global filter mixing.

    Returns
    -------
    nn.ModuleList
        List of GFNet transformer blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Create global filter mixing layer
        mixing_layer = GlobalFilterMixing(
            hidden_dim=self.hidden_dim,
            sequence_length=self.max_sequence_length,
            activation=self.filter_activation,  # Note: parameter is 'activation' not 'filter_activation'
            dropout=self._dropout_rate,
        )

        # Create GFNet block with pre-normalization
        block = PreNormBlock(
            mixing_layer=mixing_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            activation="gelu",
            dropout=self._dropout_rate,
            norm_eps=self.norm_eps,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

from_config `classmethod` ¶

from_config(config: GFNetModelConfig) -> GFNet

Create GFNet model from configuration.

Parameters:

Name	Type	Description	Default
`config`	`GFNetModelConfig`	Configuration object with model parameters.	required

Returns:

Type	Description
`GFNet`	Configured GFNet model.

Source code in spectrans/models/gfnet.py

@classmethod
def from_config(cls, config: "GFNetModelConfig") -> "GFNet":  # type: ignore[override]
    """Create GFNet model from configuration.

    Parameters
    ----------
    config : GFNetModelConfig
        Configuration object with model parameters.

    Returns
    -------
    GFNet
        Configured GFNet model.
    """
    return cls(
        vocab_size=getattr(config, "vocab_size", None),
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        num_classes=getattr(config, "num_classes", None),
        use_positional_encoding=getattr(config, "use_positional_encoding", True),
        positional_encoding_type=getattr(config, "positional_encoding_type", "sinusoidal"),
        dropout=config.dropout,
        ffn_hidden_dim=getattr(config, "ffn_hidden_dim", None),
        norm_eps=getattr(config, "norm_eps", 1e-12),
        output_type=getattr(config, "output_type", "classification"),
        filter_activation=getattr(config, "filter_activation", "sigmoid"),
        gradient_checkpointing=getattr(config, "gradient_checkpointing", False),
    )

GFNetEncoder ¶

GFNetEncoder(hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, filter_activation: FilterActivationType = 'sigmoid', gradient_checkpointing: bool = False)

Bases: GFNet

Encoder-only GFNet model for representation learning.

This variant of GFNet is designed for tasks that require extracting representations rather than making predictions. It returns the hidden states from the final layer without any task-specific head.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden dimension size. Default is 768.	`768`
`num_layers`	`int`	Number of GFNet layers. Default is 12.	`12`
`max_sequence_length`	`int`	Maximum sequence length. Default is 512.	`512`
`use_positional_encoding`	`bool`	Whether to use positional encoding. Default is True.	`True`
`positional_encoding_type`	`str`	Type of positional encoding. Default is 'sinusoidal'.	`'sinusoidal'`
`dropout`	`float`	Dropout probability. Default is 0.1.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN. If None, defaults to 4 * hidden_dim.	`None`
`norm_eps`	`float`	Epsilon for layer normalization. Default is 1e-12.	`1e-12`
`filter_activation`	`str`	Activation function for filters. Default is 'sigmoid'.	`'sigmoid'`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing. Default is False.	`False`

Source code in spectrans/models/gfnet.py

def __init__(
    self,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    filter_activation: FilterActivationType = "sigmoid",
    gradient_checkpointing: bool = False,
):
    super().__init__(
        vocab_size=None,  # No token embeddings for encoder
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=None,  # No classification head
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type="none",  # Return hidden states
        filter_activation=filter_activation,
        gradient_checkpointing=gradient_checkpointing,
    )

AlternatingTransformer ¶

AlternatingTransformer(vocab_size: int | None = None, hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, layer1_type: str = 'fourier', layer2_type: str = 'attention', layer1_config: dict | None = None, layer2_config: dict | None = None, num_classes: int | None = None, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, output_type: OutputHeadType = 'classification', gradient_checkpointing: bool = False)

Bases: BaseModel

Transformer that strictly alternates between two mixing strategies.

A simplified hybrid model that alternates between exactly two types of mixing layers following a strict pattern: layer1_type for even-indexed layers, layer2_type for odd-indexed layers. This design enables controlled comparisons between different mixing strategies.

For \(L\) layers, the alternation follows:

\[ \text{Layer}(\ell) = \begin{cases} \text{layer1_type} & \text{if } \ell \bmod 2 = 0 \\ \text{layer2_type} & \text{if } \ell \bmod 2 = 1 \end{cases} \]

Each layer applies the mixing operation with residual connection:

\[ \mathbf{X}_{\ell+1} = \mathbf{X}_\ell + \text{MixingLayer}_\ell(\text{LayerNorm}(\mathbf{X}_\ell)) \]

followed by the standard feedforward block with another residual connection.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings.	`None`
`hidden_dim`	`int`	Hidden dimension size.	`768`
`num_layers`	`int`	Number of transformer blocks.	`12`
`max_sequence_length`	`int`	Maximum sequence length.	`512`
`layer1_type`	`str`	Type of first mixing layer.	`'fourier'`
`layer2_type`	`str`	Type of second mixing layer.	`'attention'`
`layer1_config`	`dict \| None`	Configuration for first layer type.	`None`
`layer2_config`	`dict \| None`	Configuration for second layer type.	`None`
`num_classes`	`int \| None`	Number of output classes.	`None`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`PositionalEncodingType`	Type of positional encoding.	`'sinusoidal'`
`dropout`	`float`	Dropout probability.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN.	`None`
`norm_eps`	`float`	Layer normalization epsilon.	`1e-12`
`output_type`	`OutputHeadType`	Type of output head.	`'classification'`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing.	`False`

Methods:

Name	Description
`build_blocks`	Build alternating transformer blocks.

Source code in spectrans/models/hybrid.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    layer1_type: str = "fourier",
    layer2_type: str = "attention",
    layer1_config: dict | None = None,
    layer2_config: dict | None = None,
    num_classes: int | None = None,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    output_type: OutputHeadType = "classification",
    gradient_checkpointing: bool = False,
):
    # Store configuration for alternating layers before super().__init__
    # These need to be available when build_blocks is called
    self.layer1_type = layer1_type
    self.layer2_type = layer2_type
    self.layer1_config = layer1_config or {}
    self.layer2_config = layer2_config or {}
    self._dropout_rate = dropout
    self._layer1_is_spectral = layer1_type in ["fourier", "wavelet", "afno", "gfnet"]
    self._layer2_is_spectral = layer2_type in ["fourier", "wavelet", "afno", "gfnet"]

    # Initialize BaseModel
    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type=output_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build alternating transformer blocks.

Returns:

Type	Description
`ModuleList`	List of alternating transformer blocks.

Source code in spectrans/models/hybrid.py

def build_blocks(self) -> nn.ModuleList:
    """Build alternating transformer blocks.

    Returns
    -------
    nn.ModuleList
        List of alternating transformer blocks.
    """
    blocks = []

    for layer_idx in range(self.num_layers):
        # Alternate between layer types
        use_layer1 = layer_idx % 2 == 0
        mixing_layer: RealFourierMixing | FourierMixing | StandardAttention

        if use_layer1:
            # Create layer1 type
            if self.layer1_type == "fourier":
                use_real_fft = self.layer1_config.get("use_real_fft", True)
                if use_real_fft:
                    mixing_layer = RealFourierMixing(
                        hidden_dim=self.hidden_dim, dropout=self._dropout_rate
                    )
                else:
                    mixing_layer = FourierMixing(
                        hidden_dim=self.hidden_dim, dropout=self._dropout_rate
                    )
            elif self.layer1_type == "attention":
                mixing_layer = StandardAttention(
                    hidden_dim=self.hidden_dim,
                    num_heads=self.layer1_config.get("num_heads", 8),
                    dropout=self._dropout_rate,
                )
            else:
                raise ValueError(
                    f"Invalid layer1_type '{self.layer1_type}'. "
                    f"Supported types: ['fourier', 'attention']"
                )
        else:
            # Create layer2 type
            if self.layer2_type == "attention":
                mixing_layer = StandardAttention(
                    hidden_dim=self.hidden_dim,
                    num_heads=self.layer2_config.get("num_heads", 8),
                    dropout=self._dropout_rate,
                )
            elif self.layer2_type == "fourier":
                use_real_fft = self.layer2_config.get("use_real_fft", True)
                if use_real_fft:
                    mixing_layer = RealFourierMixing(
                        hidden_dim=self.hidden_dim, dropout=self._dropout_rate
                    )
                else:
                    mixing_layer = FourierMixing(
                        hidden_dim=self.hidden_dim, dropout=self._dropout_rate
                    )
            else:
                raise ValueError(
                    f"Invalid layer2_type '{self.layer2_type}'. "
                    f"Supported types: ['attention', 'fourier']"
                )

        # Create block with pre-normalization
        block = PreNormBlock(
            mixing_layer=mixing_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            activation="gelu",
            dropout=self._dropout_rate,
            norm_eps=self.norm_eps,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

HybridEncoder ¶

HybridEncoder(hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, spectral_type: str = 'fourier', spatial_type: str = 'attention', alternation_pattern: str = 'even_spectral', num_heads: int = 8, spectral_config: dict | None = None, spatial_config: dict | None = None, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, gradient_checkpointing: bool = False)

Bases: HybridTransformer

Encoder-only hybrid transformer for representation learning.

This variant returns hidden states without any task-specific head, suitable for feature extraction and representation learning.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden dimension size.	`768`
`num_layers`	`int`	Number of transformer blocks.	`12`
`max_sequence_length`	`int`	Maximum sequence length.	`512`
`spectral_type`	`str`	Type of spectral mixing.	`'fourier'`
`spatial_type`	`str`	Type of spatial mixing.	`'attention'`
`alternation_pattern`	`str`	Layer alternation pattern.	`'even_spectral'`
`num_heads`	`int`	Number of attention heads.	`8`
`spectral_config`	`dict \| None`	Spectral layer configuration.	`None`
`spatial_config`	`dict \| None`	Spatial layer configuration.	`None`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`PositionalEncodingType`	Type of positional encoding.	`'sinusoidal'`
`dropout`	`float`	Dropout probability.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN.	`None`
`norm_eps`	`float`	Layer normalization epsilon.	`1e-12`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing.	`False`

Source code in spectrans/models/hybrid.py

def __init__(
    self,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    spectral_type: str = "fourier",
    spatial_type: str = "attention",
    alternation_pattern: str = "even_spectral",
    num_heads: int = 8,
    spectral_config: dict | None = None,
    spatial_config: dict | None = None,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    gradient_checkpointing: bool = False,
):
    super().__init__(
        vocab_size=None,  # No token embeddings
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        spectral_type=spectral_type,
        spatial_type=spatial_type,
        alternation_pattern=alternation_pattern,
        num_heads=num_heads,
        spectral_config=spectral_config,
        spatial_config=spatial_config,
        num_classes=None,  # No classification head
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type="none",  # Return hidden states
        gradient_checkpointing=gradient_checkpointing,
    )

HybridTransformer ¶

HybridTransformer(vocab_size: int | None = None, hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, spectral_type: str = 'fourier', spatial_type: str = 'attention', alternation_pattern: str = 'even_spectral', num_heads: int = 8, spectral_config: dict | None = None, spatial_config: dict | None = None, num_classes: int | None = None, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, output_type: OutputHeadType = 'classification', gradient_checkpointing: bool = False)

Bases: BaseModel

Hybrid Spectral-Spatial Transformer model.

Combines spectral and spatial mixing strategies across layers to balance computational efficiency with modeling expressiveness. The model alternates between spectral layers (efficient global mixing) and spatial layers (expressive local modeling) according to configurable patterns.

For a sequence \(\mathbf{X}_0 \in \mathbb{R}^{n \times d}\), the hybrid transformer applies alternating transformations:

Spectral layers (\(\ell\) even for "even_spectral" pattern):

\[ \mathbf{Z}_\ell = \mathbf{X}_\ell + \text{SpectralMix}(\text{LN}(\mathbf{X}_\ell)) \]

Spatial layers (\(\ell\) odd for "even_spectral" pattern):

\[ \mathbf{Z}_\ell = \mathbf{X}_\ell + \text{SpatialMix}(\text{LN}(\mathbf{X}_\ell)) \]

where \(\text{LN}(\cdot)\) denotes LayerNorm and each block concludes with:

\[ \mathbf{X}_{\ell+1} = \mathbf{Z}_\ell + \text{FFN}(\text{LN}(\mathbf{Z}_\ell)) \]

The spectral mixing operations provide different complexity-accuracy tradeoffs: - Fourier: \(O(n \log n)\) via FFT/IFFT - Wavelet: \(O(n)\) via fast DWT algorithms - AFNO: \(O(k_n k_d d)\) with mode truncation parameters \(k_n, k_d\) - GFNet: \(O(n \log n)\) with learnable spectral filters

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings.	`None`
`hidden_dim`	`int`	Hidden dimension size.	`768`
`num_layers`	`int`	Number of transformer blocks.	`12`
`max_sequence_length`	`int`	Maximum sequence length.	`512`
`spectral_type`	`str`	Type of spectral mixing: 'fourier', 'wavelet', 'afno', 'gfnet'.	`'fourier'`
`spatial_type`	`str`	Type of spatial mixing: 'attention', 'spectral_attention', 'lst'.	`'attention'`
`alternation_pattern`	`str`	How to alternate: 'even_spectral', 'alternate', 'custom'.	`'even_spectral'`
`num_heads`	`int`	Number of attention heads for spatial layers.	`8`
`spectral_config`	`dict \| None`	Additional configuration for spectral layers.	`None`
`spatial_config`	`dict \| None`	Additional configuration for spatial layers.	`None`
`num_classes`	`int \| None`	Number of output classes for classification.	`None`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`PositionalEncodingType`	Type of positional encoding.	`'sinusoidal'`
`dropout`	`float`	Dropout probability.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN.	`None`
`norm_eps`	`float`	Layer normalization epsilon.	`1e-12`
`output_type`	`OutputHeadType`	Type of output head.	`'classification'`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing.	`False`

Attributes:

Name	Type	Description
`spectral_type`	`str`	Type of spectral mixing being used.
`spatial_type`	`str`	Type of spatial mixing being used.
`alternation_pattern`	`str`	The alternation pattern.
`blocks`	`ModuleList`	List of hybrid transformer blocks.

Methods:

Name	Description
`build_blocks`	Build hybrid transformer blocks.
`from_config`	Create hybrid transformer from configuration.

Source code in spectrans/models/hybrid.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    spectral_type: str = "fourier",
    spatial_type: str = "attention",
    alternation_pattern: str = "even_spectral",
    num_heads: int = 8,
    spectral_config: dict | None = None,
    spatial_config: dict | None = None,
    num_classes: int | None = None,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    output_type: OutputHeadType = "classification",
    gradient_checkpointing: bool = False,
):
    self.spectral_type = spectral_type
    self.spatial_type = spatial_type
    self.alternation_pattern = alternation_pattern
    self.num_heads = num_heads
    self.spectral_config = spectral_config or {}
    self.spatial_config = spatial_config or {}
    self._dropout_rate = dropout

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type=output_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build hybrid transformer blocks.

Returns:

Type	Description
`ModuleList`	List of transformer blocks with alternating mixing strategies.

Source code in spectrans/models/hybrid.py

def build_blocks(self) -> nn.ModuleList:
    """Build hybrid transformer blocks.

    Returns
    -------
    nn.ModuleList
        List of transformer blocks with alternating mixing strategies.
    """
    blocks = []

    for layer_idx in range(self.num_layers):
        # Determine which mixing layer to use based on pattern
        if self.alternation_pattern == "even_spectral":
            # Even layers use spectral, odd use spatial
            use_spectral = layer_idx % 2 == 0
        elif self.alternation_pattern == "alternate":
            # Strictly alternate starting with spectral
            use_spectral = layer_idx % 2 == 0
        else:  # custom or other patterns
            # Default to alternating
            use_spectral = layer_idx % 2 == 0

        # Create appropriate mixing layer
        if use_spectral:
            mixing_layer = self._create_spectral_layer()
        else:
            mixing_layer = self._create_spatial_layer()

        # Create block with pre-normalization
        block = PreNormBlock(
            mixing_layer=mixing_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            activation="gelu",
            dropout=self._dropout_rate,
            norm_eps=self.norm_eps,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

from_config `classmethod` ¶

from_config(config: HybridModelConfig) -> HybridTransformer

Create hybrid transformer from configuration.

Parameters:

Name	Type	Description	Default
`config`	`HybridModelConfig`	Configuration object with model parameters.	required

Returns:

Type	Description
`HybridTransformer`	Configured hybrid transformer model.

Source code in spectrans/models/hybrid.py

@classmethod
def from_config(cls, config: "HybridModelConfig") -> "HybridTransformer":  # type: ignore[override]
    """Create hybrid transformer from configuration.

    Parameters
    ----------
    config : HybridModelConfig
        Configuration object with model parameters.

    Returns
    -------
    HybridTransformer
        Configured hybrid transformer model.
    """
    return cls(
        vocab_size=getattr(config, "vocab_size", None),
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        spectral_type=getattr(config, "spectral_type", "fourier"),
        spatial_type=getattr(config, "spatial_type", "attention"),
        alternation_pattern=getattr(config, "alternation_pattern", "even_spectral"),
        num_heads=getattr(config, "num_heads", 8),
        spectral_config=getattr(config, "spectral_config", None),
        spatial_config=getattr(config, "spatial_config", None),
        num_classes=getattr(config, "num_classes", None),
        use_positional_encoding=getattr(config, "use_positional_encoding", True),
        positional_encoding_type=getattr(config, "positional_encoding_type", "sinusoidal"),
        dropout=config.dropout,
        ffn_hidden_dim=getattr(config, "ffn_hidden_dim", None),
        norm_eps=getattr(config, "norm_eps", 1e-12),
        output_type=getattr(config, "output_type", "classification"),
        gradient_checkpointing=getattr(config, "gradient_checkpointing", False),
    )

StandardAttention ¶

StandardAttention(hidden_dim: int, num_heads: int = 8, dropout: float = 0.0)

Bases: Module

Standard multi-head self-attention wrapper.

Wraps PyTorch's MultiheadAttention for use as a mixing layer.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden dimension size.	required
`num_heads`	`int`	Number of attention heads.	`8`
`dropout`	`float`	Dropout probability.	`0.0`

Methods:

Name	Description
`forward`	Apply self-attention.

Source code in spectrans/models/hybrid.py

def __init__(
    self,
    hidden_dim: int,
    num_heads: int = 8,
    dropout: float = 0.0,
):
    super().__init__()
    self.attention = nn.MultiheadAttention(
        hidden_dim,
        num_heads=num_heads,
        dropout=dropout,
        batch_first=True,
    )

Functions¶

forward ¶

forward(x: Tensor) -> Tensor

Apply self-attention.

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Input tensor of shape (batch_size, seq_len, hidden_dim).	required

Returns:

Type	Description
`Tensor`	Output tensor of same shape.

Source code in spectrans/models/hybrid.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply self-attention.

    Parameters
    ----------
    x : torch.Tensor
        Input tensor of shape (batch_size, seq_len, hidden_dim).

    Returns
    -------
    torch.Tensor
        Output tensor of same shape.
    """
    # Self-attention: queries, keys, and values are all x
    attn_output: torch.Tensor
    attn_output, _ = self.attention(x, x, x, need_weights=False)
    return attn_output

LSTDecoder ¶

LSTDecoder(vocab_size: int, hidden_dim: int = 512, num_layers: int = 12, max_sequence_length: int = 2048, transform_type: TransformLSTType = 'dst', causal: bool = True, use_conv_bias: bool = True, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', gradient_checkpointing: bool = False)

Bases: BaseModel

Decoder LST model with optional causal masking.

This model uses linear spectral transforms with support for causal masking, suitable for autoregressive generation tasks.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int`	Size of the vocabulary.	required
`hidden_dim`	`int`	Hidden dimension size.	`512`
`num_layers`	`int`	Number of transformer blocks.	`12`
`max_sequence_length`	`int`	Maximum sequence length.	`2048`
`transform_type`	`TransformLSTType`	Type of spectral transform (DST is preferred for causal).	`"dst"`
`causal`	`bool`	Whether to use causal masking.	`True`
`use_conv_bias`	`bool`	Use bias in spectral convolution.	`True`
`ffn_hidden_dim`	`int \| None`	FFN hidden dimension.	`None`
`dropout`	`float`	Dropout probability.	`0.0`
`use_positional_encoding`	`bool`	Use positional encoding.	`True`
`positional_encoding_type`	`str`	Positional encoding type.	`"sinusoidal"`
`gradient_checkpointing`	`bool`	Use gradient checkpointing.	`False`

Examples:

>>> decoder = LSTDecoder(
...     vocab_size=10000,
...     hidden_dim=512,
...     num_layers=12,
...     transform_type="dst",
...     causal=True,
...     max_sequence_length=2048
... )
>>> input_ids = torch.randint(0, 10000, (32, 100))
>>> logits = decoder(input_ids)
>>> assert logits.shape == (32, 100, 10000)

Methods:

Name	Description
`build_blocks`	Build decoder blocks with causal LST layers.
`forward`	Forward pass through the decoder.

Source code in spectrans/models/lst.py

def __init__(
    self,
    vocab_size: int,
    hidden_dim: int = 512,
    num_layers: int = 12,
    max_sequence_length: int = 2048,
    transform_type: TransformLSTType = "dst",
    causal: bool = True,
    use_conv_bias: bool = True,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    gradient_checkpointing: bool = False,
):
    # Store decoder-specific parameters
    self.transform_type = transform_type
    self.causal = causal
    self.use_conv_bias = use_conv_bias
    self.dropout_rate = dropout
    self.vocab_size_decoder = vocab_size

    # Initialize with language modeling head
    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=None,  # Use LM head instead
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=gradient_checkpointing,
    )

    # Add language modeling head
    self.lm_head = nn.Linear(hidden_dim, vocab_size)
    self.output_type = "lm"

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build decoder blocks with causal LST layers.

Returns:

Type	Description
`ModuleList`	List of causal LST decoder blocks.

Source code in spectrans/models/lst.py

def build_blocks(self) -> nn.ModuleList:
    """Build decoder blocks with causal LST layers.

    Returns
    -------
    nn.ModuleList
        List of causal LST decoder blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Use appropriate LST attention based on transform type
        # Note: For causal decoder, DST is preferred as it naturally handles causality
        attention_layer: DCTAttention | HadamardAttention | LSTAttention
        if self.transform_type == "dct":
            attention_layer = DCTAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )
        elif self.transform_type == "hadamard":
            attention_layer = HadamardAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )
        else:  # dst - DST is preferred for causal
            attention_layer = LSTAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,
                transform_type=self.transform_type,
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )

        block = PreNormBlock(
            mixing_layer=attention_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            dropout=self.dropout_rate,
            norm_eps=1e-12,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

forward ¶

forward(input_ids: Tensor | None = None, inputs_embeds: Tensor | None = None, attention_mask: Tensor | None = None) -> Tensor

Forward pass through the decoder.

Parameters:

Name	Type	Description	Default
`input_ids`	`Tensor \| None`	Input token IDs of shape (batch_size, sequence_length).	`None`
`inputs_embeds`	`Tensor \| None`	Pre-embedded inputs of shape (batch_size, sequence_length, hidden_dim).	`None`

Returns:

Type	Description
`Tensor`	Language modeling logits of shape (batch_size, sequence_length, vocab_size).

Source code in spectrans/models/lst.py

def forward(
    self,
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Forward pass through the decoder.

    Parameters
    ----------
    input_ids : torch.Tensor | None
        Input token IDs of shape (batch_size, sequence_length).
    inputs_embeds : torch.Tensor | None
        Pre-embedded inputs of shape (batch_size, sequence_length, hidden_dim).

    Returns
    -------
    torch.Tensor
        Language modeling logits of shape (batch_size, sequence_length, vocab_size).
    """
    # Get hidden states from base forward
    hidden_states = super().forward(
        input_ids=input_ids,
        inputs_embeds=inputs_embeds,
        attention_mask=attention_mask,
    )

    # Apply LM head
    logits = self.lm_head(hidden_states)
    return logits  # type: ignore[no-any-return]

LSTEncoder ¶

LSTEncoder(vocab_size: int | None = None, hidden_dim: int = 512, num_layers: int = 6, max_sequence_length: int = 1024, transform_type: TransformLSTType = 'dct', use_conv_bias: bool = True, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal')

Bases: BaseModel

Encoder-only LST model for representation learning.

This model uses linear spectral transforms without a classification head, suitable for generating embeddings or as a component in larger architectures.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings.	`None`
`hidden_dim`	`int`	Hidden dimension size.	`512`
`num_layers`	`int`	Number of transformer blocks.	`6`
`max_sequence_length`	`int`	Maximum sequence length.	`1024`
`transform_type`	`TransformLSTType`	Type of spectral transform.	`"dct"`
`use_conv_bias`	`bool`	Use bias in spectral convolution.	`True`
`ffn_hidden_dim`	`int \| None`	FFN hidden dimension.	`None`
`dropout`	`float`	Dropout probability.	`0.0`
`use_positional_encoding`	`bool`	Use positional encoding.	`True`
`positional_encoding_type`	`str`	Positional encoding type.	`"sinusoidal"`

Methods:

Name	Description
`build_blocks`	Build encoder blocks with LST layers.

Source code in spectrans/models/lst.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 512,
    num_layers: int = 6,
    max_sequence_length: int = 1024,
    transform_type: TransformLSTType = "dct",
    use_conv_bias: bool = True,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
):
    # Store parameters
    self.transform_type = transform_type
    self.use_conv_bias = use_conv_bias
    self.dropout_rate = dropout

    # Initialize without classification head
    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=None,  # No classification head
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=False,
    )

    # Set output type to none for encoder
    self.output_type = "none"

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build encoder blocks with LST layers.

Returns:

Type	Description
`ModuleList`	List of LST encoder blocks.

Source code in spectrans/models/lst.py

def build_blocks(self) -> nn.ModuleList:
    """Build encoder blocks with LST layers.

    Returns
    -------
    nn.ModuleList
        List of LST encoder blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Use appropriate LST attention based on transform type
        attention_layer: DCTAttention | HadamardAttention | LSTAttention
        if self.transform_type == "dct":
            attention_layer = DCTAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )
        elif self.transform_type == "hadamard":
            attention_layer = HadamardAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )
        else:  # dst
            attention_layer = LSTAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,
                transform_type=self.transform_type,
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )

        block = PreNormBlock(
            mixing_layer=attention_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            dropout=self.dropout_rate,
            norm_eps=1e-12,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

LSTTransformer ¶

LSTTransformer(vocab_size: int | None = None, hidden_dim: int = 512, num_layers: int = 6, max_sequence_length: int = 1024, transform_type: TransformLSTType = 'dct', use_conv_bias: bool = True, num_classes: int | None = None, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', gradient_checkpointing: bool = False)

Bases: BaseModel

Linear Spectral Transform transformer model.

This model uses linear spectral transforms (DCT/DST/Hadamard) for sequence mixing, achieving O(n log n) complexity through fast transform algorithms. The model applies learned transformations in the spectral domain for efficient global token interactions.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings. If None, expects pre-embedded inputs.	`None`
`hidden_dim`	`int`	Hidden dimension size for the model.	`512`
`num_layers`	`int`	Number of transformer blocks.	`6`
`max_sequence_length`	`int`	Maximum sequence length the model can process.	`1024`
`transform_type`	`TransformLSTType`	Type of spectral transform to use.	`"dct"`
`use_conv_bias`	`bool`	Whether to use bias in spectral convolution.	`True`
`num_classes`	`int \| None`	Number of output classes for classification.	`None`
`ffn_hidden_dim`	`int \| None`	Hidden dimension of the feedforward network. Default is 4 * hidden_dim.	`None`
`dropout`	`float`	Dropout probability.	`0.0`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`PositionalEncodingType`	Type of positional encoding ("sinusoidal", "learned", "rotary", "alibi", or "none").	`"sinusoidal"`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing to save memory.	`False`

Attributes:

Name	Type	Description
`blocks`	`ModuleList`	Stack of LST transformer blocks.

Examples:

>>> model = LSTTransformer(
...     hidden_dim=512,
...     num_layers=6,
...     transform_type="dct",
...     max_sequence_length=1024
... )
>>> x = torch.randn(32, 100, 512)
>>> output = model(inputs_embeds=x)
>>> assert output.shape == x.shape

Methods:

Name	Description
`build_blocks`	Build transformer blocks with LST layers.
`from_config`	Create model from configuration.

Source code in spectrans/models/lst.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 512,
    num_layers: int = 6,
    max_sequence_length: int = 1024,
    transform_type: TransformLSTType = "dct",
    use_conv_bias: bool = True,
    num_classes: int | None = None,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    gradient_checkpointing: bool = False,
):
    # Store LST-specific parameters
    self.transform_type = transform_type
    self.use_conv_bias = use_conv_bias
    self.dropout_rate = dropout

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build transformer blocks with LST layers.

Returns:

Type	Description
`ModuleList`	List of LST transformer blocks.

Source code in spectrans/models/lst.py

def build_blocks(self) -> nn.ModuleList:
    """Build transformer blocks with LST layers.

    Returns
    -------
    nn.ModuleList
        List of LST transformer blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Use appropriate LST attention based on transform type
        attention_layer: DCTAttention | HadamardAttention | LSTAttention
        if self.transform_type == "dct":
            attention_layer = DCTAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,  # Default num_heads
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )
        elif self.transform_type == "hadamard":
            attention_layer = HadamardAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )
        else:  # dst - use general LST attention
            attention_layer = LSTAttention(
                hidden_dim=self.hidden_dim,
                num_heads=8,
                transform_type=self.transform_type,
                learnable_scale=self.use_conv_bias,
                dropout=self.dropout_rate,
            )

        block = PreNormBlock(
            mixing_layer=attention_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            dropout=self.dropout_rate,
            norm_eps=1e-12,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

from_config `classmethod` ¶

from_config(config: LSTModelConfig) -> LSTTransformer

Create model from configuration.

Parameters:

Name	Type	Description	Default
`config`	`LSTModelConfig`	Model configuration object.	required

Returns:

Type	Description
`LSTTransformer`	Configured model instance.

Source code in spectrans/models/lst.py

@classmethod
def from_config(cls, config: "LSTModelConfig") -> "LSTTransformer":  # type: ignore[override]
    """Create model from configuration.

    Parameters
    ----------
    config : LSTModelConfig
        Model configuration object.

    Returns
    -------
    LSTTransformer
        Configured model instance.
    """
    return cls(
        vocab_size=config.vocab_size,
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        transform_type=config.transform_type,
        use_conv_bias=config.use_conv_bias,
        num_classes=config.num_classes,
        ffn_hidden_dim=config.ffn_hidden_dim,
        dropout=config.dropout,
        use_positional_encoding=config.use_positional_encoding,
        positional_encoding_type=config.positional_encoding_type,
        gradient_checkpointing=config.gradient_checkpointing,
    )

PerformerTransformer ¶

PerformerTransformer(vocab_size: int | None = None, hidden_dim: int = 512, num_layers: int = 6, max_sequence_length: int = 1024, num_heads: int = 8, num_features: int | None = None, num_classes: int | None = None, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', gradient_checkpointing: bool = False)

Bases: BaseModel

Performer transformer with positive orthogonal random features.

This model implements the Performer architecture which uses positive orthogonal random features (PORF) to approximate the softmax kernel with improved variance reduction compared to standard RFF.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Vocabulary size.	`None`
`hidden_dim`	`int`	Hidden dimension.	`512`
`num_layers`	`int`	Number of layers.	`6`
`max_sequence_length`	`int`	Maximum sequence length.	`1024`
`num_heads`	`int`	Number of heads.	`8`
`num_features`	`int \| None`	Number of random features.	`None`
`num_classes`	`int \| None`	Number of classes.	`None`
`ffn_hidden_dim`	`int \| None`	FFN dimension.	`None`
`dropout`	`float`	Dropout rate.	`0.0`
`use_positional_encoding`	`bool`	Use positional encoding.	`True`
`positional_encoding_type`	`str`	Positional encoding type.	`"sinusoidal"`
`gradient_checkpointing`	`bool`	Use gradient checkpointing.	`False`

Examples:

>>> performer = PerformerTransformer(
...     hidden_dim=512,
...     num_layers=6,
...     num_heads=8,
...     num_features=256,
...     max_sequence_length=1024
... )
>>> x = torch.randn(32, 100, 512)
>>> output = performer(inputs_embeds=x)

Methods:

Name	Description
`build_blocks`	Build Performer blocks with orthogonal features.

Source code in spectrans/models/spectral_attention.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 512,
    num_layers: int = 6,
    max_sequence_length: int = 1024,
    num_heads: int = 8,
    num_features: int | None = None,
    num_classes: int | None = None,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    gradient_checkpointing: bool = False,
):
    # Store parameters before calling super().__init__ since build_blocks needs them
    self.num_heads = num_heads
    self.num_features = num_features or hidden_dim
    self.dropout_rate = dropout

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build Performer blocks with orthogonal features.

Returns:

Type	Description
`ModuleList`	List of Performer blocks.

Source code in spectrans/models/spectral_attention.py

def build_blocks(self) -> nn.ModuleList:
    """Build Performer blocks with orthogonal features.

    Returns
    -------
    nn.ModuleList
        List of Performer blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        attention_layer = PerformerAttention(
            hidden_dim=self.hidden_dim,
            num_heads=self.num_heads,
            num_features=self.num_features,
            dropout=self.dropout_rate,
        )

        block = PreNormBlock(
            mixing_layer=attention_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            dropout=self.dropout_rate,
            norm_eps=1e-12,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

SpectralAttentionEncoder ¶

SpectralAttentionEncoder(vocab_size: int | None = None, hidden_dim: int = 512, num_layers: int = 6, max_sequence_length: int = 1024, num_heads: int = 8, num_features: int | None = None, kernel_type: KernelType = 'softmax', use_orthogonal: bool = False, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal')

Bases: BaseModel

Encoder-only spectral attention model for representation learning.

This model uses spectral attention layers without a classification head, suitable for generating embeddings or as a component in larger architectures.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings.	`None`
`hidden_dim`	`int`	Hidden dimension size.	`512`
`num_layers`	`int`	Number of transformer blocks.	`6`
`max_sequence_length`	`int`	Maximum sequence length.	`1024`
`num_heads`	`int`	Number of attention heads.	`8`
`num_features`	`int \| None`	Number of random features.	`None`
`kernel_type`	`KernelType`	Kernel type.	`"softmax"`
`use_orthogonal`	`bool`	Use orthogonal features.	`False`
`ffn_hidden_dim`	`int \| None`	FFN hidden dimension.	`None`
`dropout`	`float`	Dropout probability.	`0.0`
`use_positional_encoding`	`bool`	Use positional encoding.	`True`
`positional_encoding_type`	`str`	Positional encoding type.	`"sinusoidal"`

Methods:

Name	Description
`build_blocks`	Build encoder blocks with spectral attention.

Source code in spectrans/models/spectral_attention.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 512,
    num_layers: int = 6,
    max_sequence_length: int = 1024,
    num_heads: int = 8,
    num_features: int | None = None,
    kernel_type: KernelType = "softmax",
    use_orthogonal: bool = False,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
):
    # Store parameters before calling super().__init__ since build_blocks needs them
    self.num_heads = num_heads
    self.num_features = num_features or hidden_dim
    self.kernel_type = kernel_type
    self.use_orthogonal = use_orthogonal
    self.dropout_rate = dropout

    # Initialize without classification head
    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=None,  # No classification head
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=False,
    )

    # Set output type to none for encoder
    self.output_type = "none"

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build encoder blocks with spectral attention.

Returns:

Type	Description
`ModuleList`	List of spectral attention blocks.

Source code in spectrans/models/spectral_attention.py

def build_blocks(self) -> nn.ModuleList:
    """Build encoder blocks with spectral attention.

    Returns
    -------
    nn.ModuleList
        List of spectral attention blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        attention_layer = SpectralAttention(
            hidden_dim=self.hidden_dim,
            num_heads=self.num_heads,
            num_features=self.num_features,
            kernel_type=self.kernel_type,
            use_orthogonal=self.use_orthogonal,
            dropout=self.dropout_rate,
        )

        block = PreNormBlock(
            mixing_layer=attention_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            dropout=self.dropout_rate,
            norm_eps=1e-12,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

SpectralAttentionTransformer ¶

SpectralAttentionTransformer(vocab_size: int | None = None, hidden_dim: int = 512, num_layers: int = 6, max_sequence_length: int = 1024, num_heads: int = 8, num_features: int | None = None, kernel_type: KernelType = 'softmax', use_orthogonal: bool = False, num_classes: int | None = None, ffn_hidden_dim: int | None = None, dropout: float = 0.0, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', gradient_checkpointing: bool = False)

Bases: BaseModel

Spectral Attention transformer using Random Fourier Features.

This model uses spectral attention layers with RFF approximation to achieve linear complexity attention computation. The model maintains the expressive power of standard transformers while being efficient for long sequences.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings. If None, expects pre-embedded inputs.	`None`
`hidden_dim`	`int`	Hidden dimension size for the model.	`512`
`num_layers`	`int`	Number of transformer blocks.	`6`
`max_sequence_length`	`int`	Maximum sequence length the model can process.	`1024`
`num_heads`	`int`	Number of attention heads.	`8`
`num_features`	`int \| None`	Number of random features for RFF approximation. If None, uses hidden_dim.	`None`
`kernel_type`	`KernelType`	Type of kernel to approximate.	`"softmax"`
`use_orthogonal`	`bool`	Whether to use orthogonal random features.	`False`
`num_classes`	`int \| None`	Number of output classes for classification.	`None`
`ffn_hidden_dim`	`int \| None`	Hidden dimension of the feedforward network. Default is 4 * hidden_dim.	`None`
`dropout`	`float`	Dropout probability.	`0.0`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`str`	Type of positional encoding ("sinusoidal" or "learned").	`"sinusoidal"`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing to save memory.	`False`

Attributes:

Name	Type	Description
`blocks`	`ModuleList`	Stack of spectral attention transformer blocks.

Examples:

>>> model = SpectralAttentionTransformer(
...     hidden_dim=512,
...     num_layers=6,
...     num_heads=8,
...     num_features=256,
...     max_sequence_length=1024
... )
>>> x = torch.randn(32, 100, 512)
>>> output = model(inputs_embeds=x)
>>> assert output.shape == x.shape

Methods:

Name	Description
`build_blocks`	Build transformer blocks with spectral attention layers.
`from_config`	Create model from configuration.

Source code in spectrans/models/spectral_attention.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 512,
    num_layers: int = 6,
    max_sequence_length: int = 1024,
    num_heads: int = 8,
    num_features: int | None = None,
    kernel_type: KernelType = "softmax",
    use_orthogonal: bool = False,
    num_classes: int | None = None,
    ffn_hidden_dim: int | None = None,
    dropout: float = 0.0,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    gradient_checkpointing: bool = False,
):
    # Store all parameters before calling super().__init__ since build_blocks needs them
    self.num_heads = num_heads
    self.num_features = num_features or hidden_dim
    self.kernel_type = kernel_type
    self.use_orthogonal = use_orthogonal
    self.dropout_rate = dropout  # Store as different name to avoid conflict

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        ffn_hidden_dim=ffn_hidden_dim,
        dropout=dropout,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build transformer blocks with spectral attention layers.

Returns:

Type	Description
`ModuleList`	List of spectral attention transformer blocks.

Source code in spectrans/models/spectral_attention.py

def build_blocks(self) -> nn.ModuleList:
    """Build transformer blocks with spectral attention layers.

    Returns
    -------
    nn.ModuleList
        List of spectral attention transformer blocks.
    """
    blocks = []
    for _ in range(self.num_layers):
        attention_layer = SpectralAttention(
            hidden_dim=self.hidden_dim,
            num_heads=self.num_heads,
            num_features=self.num_features,
            kernel_type=self.kernel_type,
            use_orthogonal=self.use_orthogonal,
            dropout=self.dropout_rate,
        )

        block = PreNormBlock(
            mixing_layer=attention_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            dropout=self.dropout_rate,
            norm_eps=1e-12,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

from_config `classmethod` ¶

from_config(config: SpectralAttentionModelConfig) -> SpectralAttentionTransformer

Create model from configuration.

Parameters:

Name	Type	Description	Default
`config`	`SpectralAttentionModelConfig`	Model configuration object.	required

Returns:

Type	Description
`SpectralAttentionTransformer`	Configured model instance.

Source code in spectrans/models/spectral_attention.py

@classmethod
def from_config(cls, config: "SpectralAttentionModelConfig") -> "SpectralAttentionTransformer":  # type: ignore[override]
    """Create model from configuration.

    Parameters
    ----------
    config : SpectralAttentionModelConfig
        Model configuration object.

    Returns
    -------
    SpectralAttentionTransformer
        Configured model instance.
    """
    # Extract spectral attention specific config
    return cls(
        vocab_size=config.vocab_size,
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        num_heads=config.num_heads,
        num_features=config.num_features,
        kernel_type=config.kernel_type,
        use_orthogonal=config.use_orthogonal,
        num_classes=config.num_classes,
        ffn_hidden_dim=config.ffn_hidden_dim,
        dropout=config.dropout,
        use_positional_encoding=config.use_positional_encoding,
        positional_encoding_type=config.positional_encoding_type,
        gradient_checkpointing=config.gradient_checkpointing,
    )

WaveletDecoder ¶

WaveletDecoder(vocab_size: int, hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, wavelet: WaveletType = 'db4', levels: int = 2, mixing_mode: str = 'pointwise', use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, gradient_checkpointing: bool = False)

Bases: WaveletTransformer

Decoder wavelet transformer for sequence generation.

This variant uses causal wavelet processing suitable for autoregressive generation tasks. The wavelet decomposition is modified to respect causality constraints.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int`	Size of the vocabulary for token generation.	required
`hidden_dim`	`int`	Hidden dimension size.	`768`
`num_layers`	`int`	Number of wavelet transformer blocks.	`12`
`max_sequence_length`	`int`	Maximum sequence length.	`512`
`wavelet`	`WaveletType`	Type of wavelet to use.	`'db4'`
`levels`	`int`	Number of decomposition levels (typically lower for causality).	`2`
`mixing_mode`	`str`	Coefficient mixing strategy.	`'pointwise'`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`PositionalEncodingType`	Type of positional encoding.	`'sinusoidal'`
`dropout`	`float`	Dropout probability.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN.	`None`
`norm_eps`	`float`	Layer normalization epsilon.	`1e-12`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing.	`False`

Source code in spectrans/models/wavenet_transformer.py

def __init__(
    self,
    vocab_size: int,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    wavelet: WaveletType = "db4",
    levels: int = 2,  # Lower default for causality
    mixing_mode: str = "pointwise",
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    gradient_checkpointing: bool = False,
):
    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        wavelet=wavelet,
        levels=levels,
        mixing_mode=mixing_mode,
        num_classes=vocab_size,  # Output vocabulary size
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type="lm",  # Language modeling head
        gradient_checkpointing=gradient_checkpointing,
    )

WaveletEncoder ¶

WaveletEncoder(hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, wavelet: WaveletType = 'db4', levels: int = 3, mixing_mode: str = 'pointwise', use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, gradient_checkpointing: bool = False)

Bases: WaveletTransformer

Encoder-only wavelet transformer for representation learning.

This variant is designed for extracting representations from sequences using wavelet-based mixing, without any task-specific output head.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Hidden dimension size.	`768`
`num_layers`	`int`	Number of wavelet transformer blocks.	`12`
`max_sequence_length`	`int`	Maximum sequence length.	`512`
`wavelet`	`WaveletType`	Type of wavelet to use.	`'db4'`
`levels`	`int`	Number of decomposition levels.	`3`
`mixing_mode`	`str`	Coefficient mixing strategy.	`'pointwise'`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`PositionalEncodingType`	Type of positional encoding.	`'sinusoidal'`
`dropout`	`float`	Dropout probability.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN.	`None`
`norm_eps`	`float`	Layer normalization epsilon.	`1e-12`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing.	`False`

Source code in spectrans/models/wavenet_transformer.py

def __init__(
    self,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    wavelet: WaveletType = "db4",
    levels: int = 3,
    mixing_mode: str = "pointwise",
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    gradient_checkpointing: bool = False,
):
    super().__init__(
        vocab_size=None,  # No token embeddings for encoder
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        wavelet=wavelet,
        levels=levels,
        mixing_mode=mixing_mode,
        num_classes=None,  # No classification head
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type="none",  # Return hidden states
        gradient_checkpointing=gradient_checkpointing,
    )

WaveletTransformer ¶

WaveletTransformer(vocab_size: int | None = None, hidden_dim: int = 768, num_layers: int = 12, max_sequence_length: int = 512, wavelet: WaveletType = 'db4', levels: int = 3, mixing_mode: str = 'pointwise', num_classes: int | None = None, use_positional_encoding: bool = True, positional_encoding_type: PositionalEncodingType = 'sinusoidal', dropout: float = 0.1, ffn_hidden_dim: int | None = None, norm_eps: float = 1e-12, output_type: OutputHeadType = 'classification', gradient_checkpointing: bool = False)

Bases: BaseModel

Wavelet transformer with DWT-based sequence mixing.

This model replaces attention mechanisms with discrete wavelet transforms, providing multi-resolution analysis of sequences with \(O(n)\) complexity per channel. The DWT decomposes input sequences into approximation and detail coefficients at multiple scales, representing both local transients and global structure.

The wavelet mixing operation applies the DWT along the sequence dimension for each channel independently, processes the coefficients through learnable transformations, and reconstructs the sequence via the inverse DWT (IDWT). Perfect reconstruction is maintained when no coefficient modification occurs.

For input :math:\mathbf{X} \in \mathbb{R}^{n \times d}, each channel undergoes:

.. math:: \mathbf{c} = \text{DWT}J(\mathbf{X} i \in [1,d]}) \quad \text{for

.. math:: \tilde{\mathbf{c}} = f_{\theta}(\mathbf{c})

.. math:: \mathbf{Y}_{:,i} = \text{IDWT}_J(\tilde{\mathbf{c}})

where :math:f_{\theta} represents learnable coefficient transformations and :math:J is the number of decomposition levels.

Parameters:

Name	Type	Description	Default
`vocab_size`	`int \| None`	Size of the vocabulary for token embeddings. If None, expects pre-embedded inputs.	`None`
`hidden_dim`	`int`	Hidden dimension size.	`768`
`num_layers`	`int`	Number of wavelet transformer blocks.	`12`
`max_sequence_length`	`int`	Maximum sequence length the model can process.	`512`
`wavelet`	`WaveletType`	Type of wavelet to use (e.g., 'db4', 'sym6', 'coif3').	`'db4'`
`levels`	`int`	Number of wavelet decomposition levels.	`3`
`mixing_mode`	`str`	How to mix wavelet coefficients: 'pointwise', 'channel', or 'level'.	`'pointwise'`
`num_classes`	`int \| None`	Number of output classes for classification.	`None`
`use_positional_encoding`	`bool`	Whether to use positional encoding.	`True`
`positional_encoding_type`	`PositionalEncodingType`	Type of positional encoding.	`'sinusoidal'`
`dropout`	`float`	Dropout probability.	`0.1`
`ffn_hidden_dim`	`int \| None`	Hidden dimension for FFN. If None, defaults to 4 * hidden_dim.	`None`
`norm_eps`	`float`	Epsilon for layer normalization.	`1e-12`
`output_type`	`OutputHeadType`	Type of output head.	`'classification'`
`gradient_checkpointing`	`bool`	Whether to use gradient checkpointing for memory efficiency.	`False`

Attributes:

Name	Type	Description
`wavelet`	`WaveletType`	The wavelet family being used.
`levels`	`int`	Number of decomposition levels.
`mixing_mode`	`str`	Coefficient mixing strategy.
`blocks`	`ModuleList`	List of wavelet transformer blocks.

Methods:

Name	Description
`build_blocks`	Build wavelet transformer blocks.
`from_config`	Create wavelet transformer from configuration.

Source code in spectrans/models/wavenet_transformer.py

def __init__(
    self,
    vocab_size: int | None = None,
    hidden_dim: int = 768,
    num_layers: int = 12,
    max_sequence_length: int = 512,
    wavelet: WaveletType = "db4",
    levels: int = 3,
    mixing_mode: str = "pointwise",
    num_classes: int | None = None,
    use_positional_encoding: bool = True,
    positional_encoding_type: PositionalEncodingType = "sinusoidal",
    dropout: float = 0.1,
    ffn_hidden_dim: int | None = None,
    norm_eps: float = 1e-12,
    output_type: OutputHeadType = "classification",
    gradient_checkpointing: bool = False,
):
    self.wavelet = wavelet
    self.levels = levels
    self.mixing_mode = mixing_mode
    self._dropout_rate = dropout  # Store for build_blocks

    super().__init__(
        vocab_size=vocab_size,
        hidden_dim=hidden_dim,
        num_layers=num_layers,
        max_sequence_length=max_sequence_length,
        num_classes=num_classes,
        use_positional_encoding=use_positional_encoding,
        positional_encoding_type=positional_encoding_type,
        dropout=dropout,
        ffn_hidden_dim=ffn_hidden_dim,
        norm_eps=norm_eps,
        output_type=output_type,
        gradient_checkpointing=gradient_checkpointing,
    )

Functions¶

build_blocks ¶

build_blocks() -> ModuleList

Build wavelet transformer blocks.

Returns:

Type	Description
`ModuleList`	List of wavelet transformer blocks with DWT mixing layers.

Source code in spectrans/models/wavenet_transformer.py

def build_blocks(self) -> nn.ModuleList:
    """Build wavelet transformer blocks.

    Returns
    -------
    nn.ModuleList
        List of wavelet transformer blocks with DWT mixing layers.
    """
    blocks = []
    for _ in range(self.num_layers):
        # Create wavelet mixing layer
        mixing_layer = WaveletMixing(
            hidden_dim=self.hidden_dim,
            wavelet=self.wavelet,
            levels=self.levels,
            mixing_mode=self.mixing_mode,
            dropout=self._dropout_rate,
        )

        # Create block with pre-normalization
        block = PreNormBlock(
            mixing_layer=mixing_layer,
            hidden_dim=self.hidden_dim,
            ffn_hidden_dim=self.ffn_hidden_dim,
            activation="gelu",
            dropout=self._dropout_rate,
            norm_eps=self.norm_eps,
        )
        blocks.append(block)

    return nn.ModuleList(blocks)

from_config `classmethod` ¶

from_config(config: WaveletTransformerConfig) -> WaveletTransformer

Create wavelet transformer from configuration.

Parameters:

Name	Type	Description	Default
`config`	`WaveletTransformerConfig`	Configuration object with model parameters.	required

Returns:

Type	Description
`WaveletTransformer`	Configured wavelet transformer model.

Source code in spectrans/models/wavenet_transformer.py

@classmethod
def from_config(cls, config: "WaveletTransformerConfig") -> "WaveletTransformer":  # type: ignore[override]
    """Create wavelet transformer from configuration.

    Parameters
    ----------
    config : WaveletTransformerConfig
        Configuration object with model parameters.

    Returns
    -------
    WaveletTransformer
        Configured wavelet transformer model.
    """
    return cls(
        vocab_size=getattr(config, "vocab_size", None),
        hidden_dim=config.hidden_dim,
        num_layers=config.num_layers,
        max_sequence_length=config.sequence_length,
        wavelet=getattr(config, "wavelet", "db4"),
        levels=getattr(config, "levels", 3),
        mixing_mode=getattr(config, "mixing_mode", "pointwise"),
        num_classes=getattr(config, "num_classes", None),
        use_positional_encoding=getattr(config, "use_positional_encoding", True),
        positional_encoding_type=getattr(config, "positional_encoding_type", "sinusoidal"),
        dropout=config.dropout,
        ffn_hidden_dim=getattr(config, "ffn_hidden_dim", None),
        norm_eps=getattr(config, "norm_eps", 1e-12),
        output_type=getattr(config, "output_type", "classification"),
        gradient_checkpointing=getattr(config, "gradient_checkpointing", False),
    )

Model Implementations¶

spectrans.models ¶

Classes¶

AFNOEncoder ¶

AFNOModel ¶

Functions¶

build_blocks ¶

from_config classmethod ¶

ALiBiPositionalBias ¶

Functions¶

forward ¶

get_bias ¶

BaseModel ¶

Functions¶

build_blocks abstractmethod ¶

forward ¶

from_config classmethod ¶

ClassificationHead ¶

Functions¶

forward ¶

LearnedPositionalEncoding ¶

Functions¶

forward ¶

PositionalEncoding ¶

Functions¶

forward ¶

RegressionHead ¶

Functions¶

forward ¶

RotaryPositionalEncoding ¶

Functions¶

forward ¶

SequenceHead ¶

Functions¶

forward ¶

FNet ¶

Functions¶

build_blocks ¶

from_config classmethod ¶

FNetEncoder ¶

FNODecoder ¶

Functions¶

build_blocks ¶

forward ¶

FNOEncoder ¶

Functions¶

build_blocks ¶

FNOTransformer ¶

Functions¶

build_blocks ¶

from_config classmethod ¶

GFNet ¶

Functions¶

build_blocks ¶

from_config classmethod ¶

GFNetEncoder ¶

AlternatingTransformer ¶

Functions¶

build_blocks ¶

HybridEncoder ¶

HybridTransformer ¶

Functions¶

build_blocks ¶

from_config classmethod ¶

StandardAttention ¶

Functions¶

forward ¶

LSTDecoder ¶

Functions¶

build_blocks ¶

forward ¶

LSTEncoder ¶

Functions¶

build_blocks ¶

LSTTransformer ¶

Functions¶

build_blocks ¶

from_config classmethod ¶

PerformerTransformer ¶

Functions¶

from_config `classmethod` ¶

build_blocks `abstractmethod` ¶

from_config `classmethod` ¶

from_config `classmethod` ¶

from_config `classmethod` ¶

from_config `classmethod` ¶

from_config `classmethod` ¶

from_config `classmethod` ¶

from_config `classmethod` ¶

from_config `classmethod` ¶