Module neuralops.tests.test_fms

Comprehensive test suite for Factorised Matrix State Attention.

Tests: 1. Basic initialization and bidirectional forward pass. 2. Causal forward pass (decoder mode). 3. Chunked linear scan numerical equivalence with parallel scan on short sequence. 4. Linear memory and runtime behavior across sequence lengths (L=64, 128, 256). 5. Pretrained weight copy from DistilBERT, RoBERTa, GPT-2, and Qwen. 6. Missing weight extraction raises RuntimeError. 7. Nested encoder bypass of out_proj. 8. RoPE embedding and odd head dimension safety. 9. Head mask handling. 10. Symmetry breaking in L and R initialization (gradients flow into L while R starts at 0).

Functions

def test_basic_forward()

Expand source code

def test_basic_forward():
    cfg = MockConfig()
    fms = FactorisedMatrixState(config=cfg)
    x = torch.randn(2, 16, 64)
    out, _ = fms(x)
    assert out.shape == (2, 16, 64)
    assert not torch.isnan(out).any()

def test_causal_forward()

Expand source code

def test_causal_forward():
    cfg = MockConfig(is_decoder=True, model_type="gpt2")
    fms = FactorisedMatrixState(config=cfg)
    assert fms.is_decoder
    x = torch.randn(2, 32, 64)
    out, _ = fms(x)
    assert out.shape == (2, 32, 64)
    assert not torch.isnan(out).any()

def test_chunked_scan_equivalence()

Expand source code

def test_chunked_scan_equivalence():
    """Verify chunked linear scan produces numerically consistent results with small chunk size."""
    cfg = MockConfig()
    fms = FactorisedMatrixState(config=cfg)
    B, H, L, d = 2, 4, 128, 16
    q = torch.randn(B, H, L, d)
    k = torch.randn(B, H, L, d)
    v = torch.randn(B, H, L, d)
    decay = torch.sigmoid(fms.beta)

    # Run with small chunk size (32) and larger chunk size (64)
    y_c32 = fms._matrix_scan_linear(q, k, v, decay, chunk_size=32)
    y_c64 = fms._matrix_scan_linear(q, k, v, decay, chunk_size=64)
    diff = (y_c32 - y_c64).abs().max().item()
    assert diff < 1e-4, f"Chunked scan outputs differ across chunk sizes: diff={diff}"

Verify chunked linear scan produces numerically consistent results with small chunk size.

def test_distilbert_weight_copy()

Expand source code

def test_distilbert_weight_copy():
    cfg = MockConfig()
    old_attn = DummyDistilBertAttention(64)
    fms = FactorisedMatrixState(config=cfg, pretrained_attn=old_attn)
    assert not fms.is_nested_encoder
    x = torch.randn(2, 8, 64)
    out, _ = fms(x)
    assert out.shape == (2, 8, 64)

def test_gpt2_weight_copy()

Expand source code

def test_gpt2_weight_copy():
    cfg = MockConfig(is_decoder=True, model_type="gpt2")
    old_attn = DummyGPT2Attention(64)
    fms = FactorisedMatrixState(config=cfg, pretrained_attn=old_attn)
    assert fms.is_decoder
    x = torch.randn(2, 16, 64)
    out, _ = fms(x)
    assert out.shape == (2, 16, 64)

def test_long_sequence_linear_scaling()

Expand source code

def test_long_sequence_linear_scaling():
    """Test that chunked linear scan handles long sequences without OOM."""
    cfg = MockConfig()
    fms = FactorisedMatrixState(config=cfg)
    x = torch.randn(1, 256, 64)
    out, _ = fms(x)
    assert out.shape == (1, 256, 64)

Test that chunked linear scan handles long sequences without OOM.

def test_lr_symmetry_breaking()

Expand source code

def test_lr_symmetry_breaking():
    """Verify L is non-zero (Kaiming uniform) and R is zero initially, allowing gradients into L."""
    cfg = MockConfig()
    fms = FactorisedMatrixState(config=cfg)
    assert fms.L.norm().item() > 0.1, "L should have non-zero initial norm"
    assert fms.R.norm().item() == 0.0, "R should start at zero"

Verify L is non-zero (Kaiming uniform) and R is zero initially, allowing gradients into L.

def test_missing_weights_raises()

Expand source code

def test_missing_weights_raises():
    class BrokenAttn(nn.Module):
        def __init__(self):
            super().__init__()
            self.linear = nn.Linear(64, 64)

    cfg = MockConfig()
    with pytest.raises(RuntimeError, match="missing mandatory projections"):
        FactorisedMatrixState(config=cfg, pretrained_attn=BrokenAttn())

def test_odd_head_dimension_rotate_half()

Expand source code

def test_odd_head_dimension_rotate_half():
    """Verify _rotate_half handles odd dimension without shape mismatch."""
    x = torch.randn(2, 4, 8, 17)
    rot = _rotate_half(x)
    assert rot.shape == x.shape

Verify _rotate_half handles odd dimension without shape mismatch.

def test_qwen_gqa_and_rope()

Expand source code

def test_qwen_gqa_and_rope():
    cfg = MockConfig(hidden_size=64, num_attention_heads=4, num_key_value_heads=2, head_dim=16, is_decoder=True, model_type="qwen2")
    old_attn = DummyQwenAttention(dim=64, num_kv_heads=2, head_dim=16)
    fms = FactorisedMatrixState(config=cfg, pretrained_attn=old_attn)
    assert fms.kv_repeat == 2
    x = torch.randn(2, 8, 64)
    cos = torch.randn(2, 8, 16)
    sin = torch.randn(2, 8, 16)
    out, _ = fms(x, position_embeddings=(cos, sin))
    assert out.shape == (2, 8, 64)

def test_roberta_weight_copy_and_nested_bypass()

Expand source code

def test_roberta_weight_copy_and_nested_bypass():
    cfg = MockConfig()
    old_attn = DummyRobertaAttention(64)
    fms = FactorisedMatrixState(config=cfg, pretrained_attn=old_attn)
    assert fms.is_nested_encoder
    x = torch.randn(2, 8, 64)
    out, _ = fms(x)
    assert out.shape == (2, 8, 64)

Classes

class DummyGPT2Attention (dim: int = 64)

Expand source code

class DummyGPT2Attention(nn.Module):
    def __init__(self, dim: int = 64):
        super().__init__()
        self.c_attn = nn.Linear(dim, 3 * dim)
        self.c_proj = nn.Linear(dim, dim)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

class DummyQwenAttention (dim: int = 64, num_kv_heads: int = 2, head_dim: int = 16)

Expand source code

class DummyQwenAttention(nn.Module):
    def __init__(self, dim: int = 64, num_kv_heads: int = 2, head_dim: int = 16):
        super().__init__()
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, num_kv_heads * head_dim)
        self.v_proj = nn.Linear(dim, num_kv_heads * head_dim)
        self.o_proj = nn.Linear(dim, dim)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

class DummyRobertaAttention (dim: int = 64)

Expand source code

class DummyRobertaAttention(nn.Module):
    def __init__(self, dim: int = 64):
        super().__init__()
        self.self = DummyRobertaSelfAttention(dim)
        self.output = nn.Module()
        self.output.dense = nn.Linear(dim, dim)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

class DummyRobertaSelfAttention (dim: int = 64)

Expand source code

class DummyRobertaSelfAttention(nn.Module):
    def __init__(self, dim: int = 64):
        super().__init__()
        self.query = nn.Linear(dim, dim)
        self.key = nn.Linear(dim, dim)
        self.value = nn.Linear(dim, dim)

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:to, etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module