Module neuralops.tests.test_ssm_spectral_comprehensive

Functions

def main()

Expand source code

def main():
    modules_to_test = [
        "mamba", "mamba_svd_flash", "spectral_mamba", "upi_stuffed_mamba",
        "simple_ssm", "decay_linear_attention", "gated_linear_attention",
        "factorised_matrix_state", "soa_fused_attention", "soa_attn",
        "leka_v1", "leka_v2", "leka_v3"
    ]
    for m in modules_to_test:
        test_module(m)

def test_module(module_name)

Expand source code

@pytest.mark.parametrize("module_name", MODULES_TO_TEST)
def test_module(module_name):
    print(f"Testing {module_name}...")
    try:
        module_class = ATTENTION_REGISTRY.get(module_name)
        if module_class is None:
            print(f"Skipping {module_name} (not found in registry via get)")
            return
        try:
            module = module_class(d_model=D, hidden_size=D)
        except TypeError:
            try:
                module = module_class(hidden_size=D)
            except TypeError:
                module = module_class()
        
        x = torch.randn(2, 8, D)
        out, _ = module(x)
        assert out.shape == (2, 8, D)
    except Exception as e:
        print(f"{module_name} failed: {e}")