Module neuralops.tests.test_recurrent_ssm_attentions

Unit tests for all 16 Recurrent, State Space & Linear DeltaNet Attention modules.

Verifies: 1. Clean instantiation from ATTENTION_REGISTRY. 2. Forward pass output shape (2, 128, 768) with no NaNs/Infs. 3. Backward pass gradient flow with non-null gradients on inputs and parameters. 4. Pretrained Multi-Head Attention weight transfer via copy_weights(). 5. 2D and 4D attention mask handling.

Classes

class DummyConfig (hidden_size: int = 768, num_attention_heads: int = 12, rank: int = 16)

Expand source code

class DummyConfig:
    """Mock config for 768-dim Transformer (DistilBERT / BERT)."""
    def __init__(self, hidden_size: int = 768, num_attention_heads: int = 12, rank: int = 16) -> None:
        self.hidden_size = hidden_size
        self.dim = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_heads = num_attention_heads
        self.head_dim = hidden_size // num_attention_heads
        self.rank = rank
        self.dropout = 0.0
        self.hidden_dropout_prob = 0.0
        self.attention_probs_dropout_prob = 0.0

Mock config for 768-dim Transformer (DistilBERT / BERT).

class TestRecurrentSSMAttentions (methodName='runTest')

Expand source code

class TestRecurrentSSMAttentions(unittest.TestCase):
    """Rigorous verification suite for all 16 recurrent & SSM attention modules."""

    def setUp(self) -> None:
        torch.manual_seed(42)
        self.batch_size = 2
        self.seq_len = 128
        self.hidden_size = 768
        self.config = DummyConfig(hidden_size=self.hidden_size, num_attention_heads=12, rank=16)
        self.dummy_x = torch.randn(
            self.batch_size, self.seq_len, self.hidden_size, requires_grad=True
        )
        self.mask_2d = torch.ones(self.batch_size, self.seq_len)
        self.mask_2d[:, 100:] = 0.0  # Last 28 tokens padded

        # Mock pretrained MHA layer
        self.old_mha = nn.MultiheadAttention(self.hidden_size, 12, batch_first=True)

    def _test_module_lifecycle(self, key: str) -> None:
        self.assertIn(key, ATTENTION_REGISTRY, f"Key '{key}' not in ATTENTION_REGISTRY!")
        cls = ATTENTION_REGISTRY.get(key)
        
        # Instantiate
        try:
            mod = cls(config=self.config)
        except TypeError:
            try:
                mod = cls(self.config)
            except TypeError:
                mod = cls(d_model=self.hidden_size)

        # Weight transfer
        if hasattr(mod, "copy_weights") and callable(mod.copy_weights):
            mod.copy_weights(self.old_mha)

        # Forward pass (unmasked)
        x = self.dummy_x.clone().detach().requires_grad_(True)
        res = mod(x)
        out = res[0] if isinstance(res, tuple) else res

        self.assertEqual(
            out.shape,
            (self.batch_size, self.seq_len, self.hidden_size),
            f"Module '{key}' produced wrong output shape: {out.shape}",
        )
        self.assertFalse(torch.isnan(out).any(), f"Module '{key}' output contains NaNs!")
        self.assertFalse(torch.isinf(out).any(), f"Module '{key}' output contains Infs!")

        # Backward pass
        loss = out.sum()
        loss.backward()

        self.assertIsNotNone(x.grad, f"Module '{key}' did not backpropagate gradient to input x!")
        self.assertFalse(torch.isnan(x.grad).any(), f"Module '{key}' input gradients contain NaNs!")

        # Forward pass (with 2D attention mask)
        x_masked = self.dummy_x.clone().detach().requires_grad_(True)
        res_m = mod(x_masked, attention_mask=self.mask_2d)
        out_m = res_m[0] if isinstance(res_m, tuple) else res_m
        self.assertEqual(out_m.shape, (self.batch_size, self.seq_len, self.hidden_size))
        self.assertFalse(torch.isnan(out_m).any(), f"Module '{key}' masked output contains NaNs!")

    def test_01_mamba(self) -> None:
        self._test_module_lifecycle("mamba")

    def test_02_mamba_svd_flash(self) -> None:
        self._test_module_lifecycle("mamba_svd_flash")

    def test_03_gated_delta_net(self) -> None:
        self._test_module_lifecycle("gated_delta_net")

    def test_04_gated_deltanet_block(self) -> None:
        self._test_module_lifecycle("gated_deltanet_block")

    def test_05_hgrn2_layer(self) -> None:
        self._test_module_lifecycle("hgrn2_layer")

    def test_06_hgrn2_gated_deltanet(self) -> None:
        self._test_module_lifecycle("hgrn2_gated_deltanet")

    def test_07_hgrn2_outer(self) -> None:
        self._test_module_lifecycle("hgrn2_outer")

    def test_08_griffin_ssd(self) -> None:
        self._test_module_lifecycle("griffin_ssd")

    def test_09_hla_attention(self) -> None:
        self._test_module_lifecycle("hla")

    def test_10_kda_attention(self) -> None:
        self._test_module_lifecycle("kda")

    def test_11_simple_ssm(self) -> None:
        self._test_module_lifecycle("simple_ssm")

    def test_12_decay_linear_attention(self) -> None:
        self._test_module_lifecycle("decay_linear_attention")

    def test_13_gated_linear_attention(self) -> None:
        self._test_module_lifecycle("gated_linear_attention")

    def test_14_factorised_matrix_state(self) -> None:
        self._test_module_lifecycle("factorised_matrix_state")

    def test_15_upi_stuffed_mamba(self) -> None:
        self._test_module_lifecycle("upi_stuffed_mamba")

Rigorous verification suite for all 16 recurrent & SSM attention modules.

Create an instance of the class that will use the named test method when executed. Raises a ValueError if the instance does not have a method with the specified name.

Ancestors

• unittest.case.TestCase

Methods

def setUp(self) ‑> None

Expand source code

def setUp(self) -> None:
    torch.manual_seed(42)
    self.batch_size = 2
    self.seq_len = 128
    self.hidden_size = 768
    self.config = DummyConfig(hidden_size=self.hidden_size, num_attention_heads=12, rank=16)
    self.dummy_x = torch.randn(
        self.batch_size, self.seq_len, self.hidden_size, requires_grad=True
    )
    self.mask_2d = torch.ones(self.batch_size, self.seq_len)
    self.mask_2d[:, 100:] = 0.0  # Last 28 tokens padded

    # Mock pretrained MHA layer
    self.old_mha = nn.MultiheadAttention(self.hidden_size, 12, batch_first=True)

Hook method for setting up the test fixture before exercising it.

def test_01_mamba(self) ‑> None

Expand source code

def test_01_mamba(self) -> None:
    self._test_module_lifecycle("mamba")

def test_02_mamba_svd_flash(self) ‑> None

Expand source code

def test_02_mamba_svd_flash(self) -> None:
    self._test_module_lifecycle("mamba_svd_flash")

def test_03_gated_delta_net(self) ‑> None

Expand source code

def test_03_gated_delta_net(self) -> None:
    self._test_module_lifecycle("gated_delta_net")

def test_04_gated_deltanet_block(self) ‑> None

Expand source code

def test_04_gated_deltanet_block(self) -> None:
    self._test_module_lifecycle("gated_deltanet_block")

def test_05_hgrn2_layer(self) ‑> None

Expand source code

def test_05_hgrn2_layer(self) -> None:
    self._test_module_lifecycle("hgrn2_layer")

def test_06_hgrn2_gated_deltanet(self) ‑> None

Expand source code

def test_06_hgrn2_gated_deltanet(self) -> None:
    self._test_module_lifecycle("hgrn2_gated_deltanet")

def test_07_hgrn2_outer(self) ‑> None

Expand source code

def test_07_hgrn2_outer(self) -> None:
    self._test_module_lifecycle("hgrn2_outer")

def test_08_griffin_ssd(self) ‑> None

Expand source code

def test_08_griffin_ssd(self) -> None:
    self._test_module_lifecycle("griffin_ssd")

def test_09_hla_attention(self) ‑> None

Expand source code

def test_09_hla_attention(self) -> None:
    self._test_module_lifecycle("hla")

def test_10_kda_attention(self) ‑> None

Expand source code

def test_10_kda_attention(self) -> None:
    self._test_module_lifecycle("kda")

def test_11_simple_ssm(self) ‑> None

Expand source code

def test_11_simple_ssm(self) -> None:
    self._test_module_lifecycle("simple_ssm")

def test_12_decay_linear_attention(self) ‑> None

Expand source code

def test_12_decay_linear_attention(self) -> None:
    self._test_module_lifecycle("decay_linear_attention")

def test_13_gated_linear_attention(self) ‑> None

Expand source code

def test_13_gated_linear_attention(self) -> None:
    self._test_module_lifecycle("gated_linear_attention")

def test_14_factorised_matrix_state(self) ‑> None

Expand source code

def test_14_factorised_matrix_state(self) -> None:
    self._test_module_lifecycle("factorised_matrix_state")

def test_15_upi_stuffed_mamba(self) ‑> None

Expand source code

def test_15_upi_stuffed_mamba(self) -> None:
    self._test_module_lifecycle("upi_stuffed_mamba")