Module neuralops.tests.test_attentions

Classes

class DummyConfig (hidden_size=64, num_heads=4, rank=16)

Expand source code

class DummyConfig:
    def __init__(self, hidden_size=64, num_heads=4, rank=16):
        self.hidden_size = hidden_size
        self.num_attention_heads = num_heads
        self.num_heads = num_heads
        self.head_dim = hidden_size // max(1, num_heads)
        self.rank = rank
        self.dropout = 0.0
        self.attention_dropout = 0.0
        self.window_size = 32
        self.svd_rank = 16

class TestAttentionMechanisms (methodName='runTest')

Expand source code

class TestAttentionMechanisms(unittest.TestCase):
    """Test suite verifying attention module registry, forward propagation, and output shapes."""

    def setUp(self):
        self.cfg = DummyConfig(hidden_size=64, num_heads=4, rank=16)
        self.batch_size = 2
        self.seq_len = 16
        self.hidden_size = 64
        self.dummy_x = torch.randn(self.batch_size, self.seq_len, self.hidden_size)

    def _instantiate_module(self, cls):
        try:
            return cls(config=self.cfg)
        except TypeError:
            try:
                return cls(self.cfg)
            except TypeError:
                return cls(hidden_size=self.hidden_size, num_heads=self.cfg.num_heads)

    def test_registry_size(self):
        self.assertGreaterEqual(len(ATTENTION_REGISTRY), 50)

    def test_full_and_gqa_attention(self):
        for name in ["full", "gqa", "mla"]:
            if name in ATTENTION_REGISTRY:
                cls = ATTENTION_REGISTRY.get(name)
                mod = self._instantiate_module(cls)
                out = mod(self.dummy_x)
                if isinstance(out, tuple):
                    out = out[0]
                self.assertEqual(out.shape, (self.batch_size, self.seq_len, self.hidden_size))

    def test_svd_linear_attention(self):
        for name in ["svd_linear", "svd_linear_attn"]:
            if name in ATTENTION_REGISTRY:
                cls = ATTENTION_REGISTRY.get(name)
                mod = self._instantiate_module(cls)
                out = mod(self.dummy_x)
                if isinstance(out, tuple):
                    out = out[0]
                self.assertEqual(out.shape, (self.batch_size, self.seq_len, self.hidden_size))

    def test_next_gen_hybrids(self):
        hybrids = ["mamba_svd_flash", "moe_deltakv_paged", "aegis_spectron_fp8", "aegis_hybrid"]
        for name in hybrids:
            if name in ATTENTION_REGISTRY:
                cls = ATTENTION_REGISTRY.get(name)
                mod = self._instantiate_module(cls)
                out = mod(self.dummy_x)
                if isinstance(out, tuple):
                    out = out[0]
                self.assertEqual(out.shape, (self.batch_size, self.seq_len, self.hidden_size))

Test suite verifying attention module registry, forward propagation, and output shapes.

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)

Expand source code

def setUp(self):
    self.cfg = DummyConfig(hidden_size=64, num_heads=4, rank=16)
    self.batch_size = 2
    self.seq_len = 16
    self.hidden_size = 64
    self.dummy_x = torch.randn(self.batch_size, self.seq_len, self.hidden_size)

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

def test_full_and_gqa_attention(self)

Expand source code

def test_full_and_gqa_attention(self):
    for name in ["full", "gqa", "mla"]:
        if name in ATTENTION_REGISTRY:
            cls = ATTENTION_REGISTRY.get(name)
            mod = self._instantiate_module(cls)
            out = mod(self.dummy_x)
            if isinstance(out, tuple):
                out = out[0]
            self.assertEqual(out.shape, (self.batch_size, self.seq_len, self.hidden_size))

def test_next_gen_hybrids(self)

Expand source code

def test_next_gen_hybrids(self):
    hybrids = ["mamba_svd_flash", "moe_deltakv_paged", "aegis_spectron_fp8", "aegis_hybrid"]
    for name in hybrids:
        if name in ATTENTION_REGISTRY:
            cls = ATTENTION_REGISTRY.get(name)
            mod = self._instantiate_module(cls)
            out = mod(self.dummy_x)
            if isinstance(out, tuple):
                out = out[0]
            self.assertEqual(out.shape, (self.batch_size, self.seq_len, self.hidden_size))

def test_registry_size(self)

Expand source code

def test_registry_size(self):
    self.assertGreaterEqual(len(ATTENTION_REGISTRY), 50)

def test_svd_linear_attention(self)

Expand source code

def test_svd_linear_attention(self):
    for name in ["svd_linear", "svd_linear_attn"]:
        if name in ATTENTION_REGISTRY:
            cls = ATTENTION_REGISTRY.get(name)
            mod = self._instantiate_module(cls)
            out = mod(self.dummy_x)
            if isinstance(out, tuple):
                out = out[0]
            self.assertEqual(out.shape, (self.batch_size, self.seq_len, self.hidden_size))