Module neuralops.tests.test_aegis_hybrid

Hardening test suite for aegis_hybrid.py (AegisHybridAttention).

Functions

def test_config_none()

Expand source code

def test_config_none():
    attn = AegisHybridAttention(config=None, hidden_size=64, num_heads=4)
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)
    assert torch.isfinite(out).all()

def test_decoder()

Expand source code

def test_decoder():
    cfg = MockConfig(is_decoder=True, model_type="gpt2")
    attn = AegisHybridAttention(config=cfg)
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)
    assert torch.isfinite(out).all()

def test_distilbert_copy()

Expand source code

def test_distilbert_copy():
    attn = AegisHybridAttention(config=MockConfig(), pretrained_attn=DummyDistilBertAttn(64))
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_fail_loud()

Expand source code

def test_fail_loud():
    attn = AegisHybridAttention(config=MockConfig())
    with pytest.raises(RuntimeError, match="mandatory projections"):
        attn.copy_weights(DummyEmptyAttn())

def test_forward_stable()

Expand source code

def test_forward_stable():
    attn = AegisHybridAttention(config=MockConfig())
    out, _ = attn(torch.randn(2, 12, 64))
    assert not torch.isnan(out).any()
    assert not torch.isinf(out).any()

def test_gpt2_copy()

Expand source code

def test_gpt2_copy():
    cfg = MockConfig(is_decoder=True, model_type="gpt2")
    attn = AegisHybridAttention(config=cfg, pretrained_attn=DummyGPT2Attn(64))
    assert attn.is_decoder
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_inject_distilbert()

Expand source code

def test_inject_distilbert():
    class DummyLayer(nn.Module):
        def __init__(self):
            super().__init__()
            self.attention = DummyDistilBertAttn(64)
    layer = DummyLayer()
    res = inject_aegis_hybrid(layer=layer, model_config=MockConfig())
    assert res is not None
    assert isinstance(layer.attention, AegisHybridAttention)
    out, _ = layer.attention(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_inject_none()

Expand source code

def test_inject_none():
    assert inject_aegis_hybrid(layer=None, model=None) is None

def test_inject_none_on_failure()

Expand source code

def test_inject_none_on_failure():
    class DummyLayer(nn.Module):
        def __init__(self):
            super().__init__()
            self.attention = DummyEmptyAttn()
    assert inject_aegis_hybrid(layer=DummyLayer(), model_config=MockConfig()) is None

def test_qwen_gqa()

Expand source code

def test_qwen_gqa():
    cfg = MockConfig(hidden_size=64, num_attention_heads=4, num_key_value_heads=2,
                     head_dim=16, is_decoder=True, model_type="qwen2")
    attn = AegisHybridAttention(config=cfg, pretrained_attn=DummyQwenAttn(64, 2, 16))
    # kv_repeat may be named differently in hybrid — just verify output shape
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_registry()

Expand source code

def test_registry():
    assert ATTENTION_REGISTRY.get("aegis_hybrid") is not None

def test_roberta_copy()

Expand source code

def test_roberta_copy():
    attn = AegisHybridAttention(config=MockConfig(), pretrained_attn=DummyRobertaAttn(64))
    assert attn.is_nested_encoder
    out, _ = attn(torch.randn(1, 8, 64))
    assert out.shape == (1, 8, 64)

def test_rope_even_head_dim()

Expand source code

def test_rope_even_head_dim():
    attn = AegisHybridAttention(config=MockConfig())
    cos = torch.randn(1, 8, 16)
    sin = torch.randn(1, 8, 16)
    out, _ = attn(torch.randn(1, 8, 64), position_embeddings=(cos, sin))
    assert out.shape == (1, 8, 64)
    assert torch.isfinite(out).all()

def test_rope_odd_head_dim()

Expand source code

def test_rope_odd_head_dim():
    attn = AegisHybridAttention(config=MockConfig(hidden_size=68, num_attention_heads=4, head_dim=17))
    cos = torch.randn(1, 4, 17)
    sin = torch.randn(1, 4, 17)
    out, _ = attn(torch.randn(1, 4, 68), position_embeddings=(cos, sin))
    assert out.shape == (1, 4, 68)
    assert torch.isfinite(out).all()

Classes

class DummyDistilBertAttn (dim=64)

Expand source code

class DummyDistilBertAttn(nn.Module):
    def __init__(self, dim=64):
        super().__init__()
        self.q_lin = nn.Linear(dim, dim)
        self.k_lin = nn.Linear(dim, dim)
        self.v_lin = nn.Linear(dim, dim)
        self.out_lin = 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 DummyEmptyAttn

Expand source code

class DummyEmptyAttn(nn.Module):
    def __init__(self):
        super().__init__()
        self.foo = nn.Linear(64, 64)

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 DummyGPT2Attn (dim=64)

Expand source code

class DummyGPT2Attn(nn.Module):
    def __init__(self, dim=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 DummyQwenAttn (dim=64, num_kv_heads=2, head_dim=16)

Expand source code

class DummyQwenAttn(nn.Module):
    def __init__(self, dim=64, num_kv_heads=2, head_dim=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 DummyRobertaAttn (dim=64)

Expand source code

class DummyRobertaAttn(nn.Module):
    def __init__(self, dim=64):
        super().__init__()
        self.self = DummyRobertaSelfAttn(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 DummyRobertaSelfAttn (dim=64)

Expand source code

class DummyRobertaSelfAttn(nn.Module):
    def __init__(self, dim=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