Module neuralops.tests.test_attention_injection

Tests for the attention injection utilities.

Classes

class TestAttentionInjection

Expand source code

class TestAttentionInjection:
    """Verify that inject_attention works correctly on foundation transformer models."""

    @pytest.mark.integration
    def test_inject_attention_on_distilbert(self):
        """Inject known attention types into DistilBERT and verify forward pass execution."""
        config = DistilBertConfig(
            vocab_size=100,
            dim=768,
            n_layers=2,
            n_heads=8,
            hidden_dim=3072,
        )
        model = DistilBertModel(config)
        inject_attention(model, "full", copy_weights=False)

        # Verify forward pass still works
        dummy = (
            model.distilbert.transformer.layer[0].attention
            if hasattr(model, "distilbert")
            else model.transformer.layer[0].attention
        )

        x = torch.randn(2, 8, 768)
        out = dummy(x)
        if isinstance(out, tuple):
            out = out[0]
        assert out.shape == (2, 8, 768)

    @pytest.mark.integration
    def test_inject_attention_with_copy_weights(self):
        """Verify injection with copy_weights enabled."""
        config = DistilBertConfig(
            vocab_size=100,
            dim=768,
            n_layers=2,
            n_heads=8,
            hidden_dim=3072,
        )
        model = DistilBertModel(config)
        inject_attention(model, "full", copy_weights=True)

        x = torch.randint(0, 100, (2, 16))
        outputs = model(input_ids=x)
        assert outputs.last_hidden_state.shape == (2, 16, 768)

    def test_available_attention_types(self):
        """Verify registry queries return non-empty list of attention variants."""
        attn_types = available_attention_types()
        assert len(attn_types) > 0
        assert "full" in attn_types

Verify that inject_attention works correctly on foundation transformer models.

Methods

def test_available_attention_types(self)

Expand source code

def test_available_attention_types(self):
    """Verify registry queries return non-empty list of attention variants."""
    attn_types = available_attention_types()
    assert len(attn_types) > 0
    assert "full" in attn_types

Verify registry queries return non-empty list of attention variants.

def test_inject_attention_on_distilbert(self)

Expand source code

@pytest.mark.integration
def test_inject_attention_on_distilbert(self):
    """Inject known attention types into DistilBERT and verify forward pass execution."""
    config = DistilBertConfig(
        vocab_size=100,
        dim=768,
        n_layers=2,
        n_heads=8,
        hidden_dim=3072,
    )
    model = DistilBertModel(config)
    inject_attention(model, "full", copy_weights=False)

    # Verify forward pass still works
    dummy = (
        model.distilbert.transformer.layer[0].attention
        if hasattr(model, "distilbert")
        else model.transformer.layer[0].attention
    )

    x = torch.randn(2, 8, 768)
    out = dummy(x)
    if isinstance(out, tuple):
        out = out[0]
    assert out.shape == (2, 8, 768)

Inject known attention types into DistilBERT and verify forward pass execution.

def test_inject_attention_with_copy_weights(self)

Expand source code

@pytest.mark.integration
def test_inject_attention_with_copy_weights(self):
    """Verify injection with copy_weights enabled."""
    config = DistilBertConfig(
        vocab_size=100,
        dim=768,
        n_layers=2,
        n_heads=8,
        hidden_dim=3072,
    )
    model = DistilBertModel(config)
    inject_attention(model, "full", copy_weights=True)

    x = torch.randint(0, 100, (2, 16))
    outputs = model(input_ids=x)
    assert outputs.last_hidden_state.shape == (2, 16, 768)

Verify injection with copy_weights enabled.