Module neuralops.tests.test_cosformer

Comprehensive test suite for Cosformer Attention (cosformer.py).

Functions

def test_copy_weights_distilbert()

Expand source code

def test_copy_weights_distilbert():
    class MockDistilBertAttn(nn.Module):
        def __init__(self):
            super().__init__()
            self.q_lin = nn.Linear(64, 64)
            self.k_lin = nn.Linear(64, 64)
            self.v_lin = nn.Linear(64, 64)
            self.out_lin = nn.Linear(64, 64)

    mock = MockDistilBertAttn()
    cfg = MockConfig()
    layer = CosformerAttention(config=cfg, d_model=64, num_heads=4)
    layer.copy_weights(mock)
    assert torch.allclose(layer.q_proj.weight, mock.q_lin.weight)
    assert torch.allclose(layer.out_proj.weight, mock.out_lin.weight)

def test_copy_weights_gpt2()

Expand source code

def test_copy_weights_gpt2():
    class MockGPT2Attn(nn.Module):
        def __init__(self):
            super().__init__()
            self.c_attn = nn.Linear(64, 192)
            self.c_proj = nn.Linear(64, 64)

    mock = MockGPT2Attn()
    cfg = MockConfig(is_decoder=True)
    layer = CosformerAttention(config=cfg, d_model=64, num_heads=4)
    layer.copy_weights(mock)
    assert layer.q_proj.weight.shape == (64, 64)
    assert layer.out_proj.weight.shape == (64, 64)

def test_copy_weights_missing_raises()

Expand source code

def test_copy_weights_missing_raises():
    class IncompleteAttn(nn.Module):
        def __init__(self):
            super().__init__()
            self.q_lin = nn.Linear(64, 64)

    mock = IncompleteAttn()
    cfg = MockConfig()
    layer = CosformerAttention(config=cfg, d_model=64, num_heads=4)
    with pytest.raises(RuntimeError):
        layer.copy_weights(mock)

def test_copy_weights_roberta_nested()

Expand source code

def test_copy_weights_roberta_nested():
    class MockRobertaSelf(nn.Module):
        def __init__(self):
            super().__init__()
            self.query = nn.Linear(64, 64)
            self.key = nn.Linear(64, 64)
            self.value = nn.Linear(64, 64)

    class MockRobertaOutput(nn.Module):
        def __init__(self):
            super().__init__()
            self.dense = nn.Linear(64, 64)

    class MockRobertaAttention(nn.Module):
        def __init__(self):
            super().__init__()
            self.self = MockRobertaSelf()
            self.output = MockRobertaOutput()

    mock = MockRobertaAttention()
    cfg = MockConfig()
    layer = CosformerAttention(config=cfg, d_model=64, num_heads=4)
    layer.copy_weights(mock)
    assert layer.is_nested_encoder
    assert torch.allclose(layer.q_proj.weight, mock.self.query.weight)

def test_forward_bidirectional()

Expand source code

def test_forward_bidirectional():
    cfg = MockConfig(is_decoder=False)
    layer = CosformerAttention(config=cfg, d_model=64, num_heads=4)
    x = torch.randn(2, 16, 64)
    mask = torch.ones(2, 16)
    out, attn = layer(x, attention_mask=mask)
    assert out.shape == (2, 16, 64)
    assert attn is None
    assert not torch.isnan(out).any()

def test_forward_causal()

Expand source code

def test_forward_causal():
    cfg = MockConfig(is_decoder=True)
    layer = CosformerAttention(config=cfg, d_model=64, num_heads=4)
    x = torch.randn(2, 16, 64)
    out, attn = layer(x)
    assert out.shape == (2, 16, 64)
    assert not torch.isnan(out).any()

def test_inject_cosformer_single_layer()

Expand source code

def test_inject_cosformer_single_layer():
    class MockAttn(nn.Module):
        def __init__(self):
            super().__init__()
            self.q_lin = nn.Linear(64, 64)
            self.k_lin = nn.Linear(64, 64)
            self.v_lin = nn.Linear(64, 64)
            self.out_lin = nn.Linear(64, 64)

    class MockLayer(nn.Module):
        def __init__(self):
            super().__init__()
            self.attention = MockAttn()

    layer = MockLayer()
    cfg = MockConfig()
    ret = inject_cosformer(layer=layer, model_config=cfg, layer_idx=0, total_layers=6)
    assert isinstance(ret, CosformerAttention)
    assert layer.attention is ret

def test_registry()

Expand source code

def test_registry():
    assert "cosformer" in ATTENTION_REGISTRY.list_available()
    assert "cosformer_attention" in ATTENTION_REGISTRY.list_available()
    assert ATTENTION_REGISTRY.get("cosformer") is CosformerAttention

def test_rotate_half_odd_dimension()

Expand source code

def test_rotate_half_odd_dimension():
    x = torch.randn(2, 4, 15)
    rot = _rotate_half(x)
    assert rot.shape == x.shape
    assert torch.allclose(rot[..., 14], x[..., 14])