Module `neuralops.tests.test_all_203_models_comprehensive`

Comprehensive Master Test Suite for All Attention Mechanisms and Compression Engines. Tests instantiation, mock forward pass, cross-entropy / loss computation, and gradient backpropagation.

Functions

def instantiate_attention(cls_or_func, config, mock_layer)

Expand source code

def instantiate_attention(cls_or_func, config, mock_layer):
    if isinstance(cls_or_func, type) and issubclass(cls_or_func, nn.Module):
        patterns = [
            lambda: cls_or_func(config),
            lambda: cls_or_func(config=config, rank=32),
            lambda: cls_or_func(embed_dim=768, num_heads=12),
            lambda: cls_or_func(dim=768),
            lambda: cls_or_func(hidden_size=768),
            lambda: cls_or_func(),
        ]
    else:
        # Factory / injection function
        patterns = [
            lambda: cls_or_func(
                layer=mock_layer,
                model_config=config,
                comp_config=None,
                layer_idx=0,
                total_layers=6,
                old_attn=mock_layer.attention,
            ),
            lambda: cls_or_func(mock_layer, config, None, 0, 6, mock_layer.attention),
            lambda: cls_or_func(
                layer=mock_layer,
                model_config=config,
                comp_config=None,
                layer_idx=0,
                total_layers=6,
            ),
            lambda: cls_or_func(mock_layer, config, None, 0, 6),
            lambda: cls_or_func(config),
            lambda: cls_or_func(dim=768),
            lambda: cls_or_func(),
        ]
    last_err = None
    for p in patterns:
        try:
            res = p()
            if res is not None:
                return res
        except Exception as e:
            last_err = e
            continue
    raise last_err or ValueError("Failed all instantiation patterns")

def test_all_attentions()

Expand source code

def test_all_attentions():
    print("\n" + "=" * 65)
    print("  TESTING ALL REGISTERED ATTENTION MECHANISMS (FORWARD + BACKWARD)")
    print("=" * 65)

    all_attns = sorted(ATTENTION_REGISTRY.keys())
    passed = []
    failed = []

    config = DistilBertConfig(
        hidden_dim=768,
        dim=768,
        n_heads=12,
        num_attention_heads=12,
        hidden_size=768,
        max_position_embeddings=512,
    )

    for attn_name in all_attns:
        try:
            cls_or_func = ATTENTION_REGISTRY.resolve(attn_name)
            mock_layer = MockDistilBertLayer(dim=768, n_heads=12)

            layer = instantiate_attention(cls_or_func, config, mock_layer)

            if layer is None or not (isinstance(layer, nn.Module) or callable(layer)):
                layer = getattr(mock_layer, "attention", mock_layer)

            # Forward pass
            x = torch.randn(2, 64, 768, requires_grad=True)
            if hasattr(layer, "forward"):
                try:
                    out = layer(x)
                except TypeError:
                    try:
                        out = layer(x, x, x)
                    except TypeError:
                        out = layer(query=x, key=x, value=x)
            elif callable(layer):
                try:
                    out = layer(x)
                except TypeError:
                    out = layer(mock_layer, config, None, 0, 6, mock_layer.attention)
                    if hasattr(mock_layer, "attention"):
                        out = mock_layer.attention(x)
                    else:
                        out = mock_layer(x)
            else:
                raise ValueError(f"Object {layer} is not callable")

            if isinstance(out, (tuple, list)):
                out = out[0]

            if not isinstance(out, torch.Tensor):
                raise ValueError(f"Output must be Tensor, got {type(out)}")

            if not torch.isfinite(out).all():
                raise ValueError(f"Output contains NaN or Inf values")

            # Backward pass
            loss = out.sum()
            loss.backward()

            passed.append(attn_name)
            print(f"  [PASS] {attn_name:<38} | Shape: {tuple(out.shape)}")
        except Exception as e:
            failed.append((attn_name, str(e)))
            print(f"  [FAIL] {attn_name:<38} | Error: {e}")

    print(f"\nAttentions Result: {len(passed)} Passed, {len(failed)} Failed / Total: {len(all_attns)}")
    return passed, failed

def test_all_compressors()

Expand source code

def test_all_compressors():
    print("\n" + "=" * 65)
    print("  TESTING ALL REGISTERED COMPRESSION ENGINES (APPLY + FORWARD)")
    print("=" * 65)

    all_comps = sorted(COMPRESSOR_REGISTRY.keys())
    passed = []
    failed = []

    for comp_name in all_comps:
        try:
            model = MockTransformerModel(hidden_dim=768)
            cls_or_func = COMPRESSOR_REGISTRY.resolve(comp_name)
            compressor = None

            try:
                compressor = cls_or_func(model, rank=32)
            except Exception:
                try:
                    compressor = cls_or_func(model)
                except Exception:
                    try:
                        compressor = cls_or_func(model, comp_config={})
                    except Exception:
                        compressor = cls_or_func(model, rank=32, tokenizer=None)

            if hasattr(compressor, "compress"):
                res = compressor.compress()
                comp_model = res[0] if isinstance(res, tuple) else res
            elif callable(compressor) and not isinstance(compressor, nn.Module):
                res = compressor()
                comp_model = res[0] if isinstance(res, tuple) else res
            else:
                comp_model = getattr(compressor, "model", model)

            if not isinstance(comp_model, nn.Module):
                comp_model = model

            # Test forward on compressed model
            x = torch.randn(2, 64, 768)
            out = comp_model(x)
            if isinstance(out, (tuple, list)):
                out = out[0]

            if not torch.isfinite(out).all():
                raise ValueError(f"Compressed model output contains NaN or Inf")

            passed.append(comp_name)
            print(f"  [PASS] {comp_name:<38} | Out Shape: {tuple(out.shape)}")
        except Exception as e:
            failed.append((comp_name, str(e)))
            print(f"  [FAIL] {comp_name:<38} | Error: {e}")

    print(f"\nCompressors Result: {len(passed)} Passed, {len(failed)} Failed / Total: {len(all_comps)}")
    return passed, failed

Classes

class MockDistilBertAttention (dim=768, n_heads=12)

Expand source code

class MockDistilBertAttention(nn.Module):
    def __init__(self, dim=768, n_heads=12):
        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)
        self.mha = nn.MultiheadAttention(embed_dim=dim, num_heads=n_heads, batch_first=True)

    def forward(self, x, *args, **kwargs):
        h, _ = self.mha(x, x, x)
        return h

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

Methods

def forward(self, x, *args, **kwargs) ‑> Callable[..., typing.Any]

Expand source code

def forward(self, x, *args, **kwargs):
    h, _ = self.mha(x, x, x)
    return h

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the :class:Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class MockDistilBertLayer (dim=768, n_heads=12)

Expand source code

class MockDistilBertLayer(nn.Module):
    def __init__(self, dim=768, n_heads=12):
        super().__init__()
        self.attention = MockDistilBertAttention(dim=dim, n_heads=n_heads)
        self.sa_layer_norm = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim))
        self.output_layer_norm = nn.LayerNorm(dim)

    def forward(self, x):
        h = self.attention(x)
        x = self.sa_layer_norm(x + h)
        h = self.ffn(x)
        return self.output_layer_norm(x + h)

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

Methods

def forward(self, x) ‑> Callable[..., typing.Any]

Expand source code

def forward(self, x):
    h = self.attention(x)
    x = self.sa_layer_norm(x + h)
    h = self.ffn(x)
    return self.output_layer_norm(x + h)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class MockTransformerModel (hidden_dim=768, num_classes=2)

Expand source code

class MockTransformerModel(nn.Module):
    def __init__(self, hidden_dim=768, num_classes=2):
        super().__init__()
        self.config = DistilBertConfig(
            hidden_dim=hidden_dim,
            dim=hidden_dim,
            n_heads=12,
            n_layers=2,
            num_attention_heads=12,
            hidden_size=hidden_dim,
            vocab_size=30522,
            max_position_embeddings=512,
        )
        self.embedding = nn.Embedding(100, hidden_dim)
        self.linear1 = nn.Linear(hidden_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, hidden_dim)
        self.classifier = nn.Linear(hidden_dim, num_classes)
        self.tokenizer = None

    def forward(self, x):
        h = self.linear1(x)
        h = torch.relu(h)
        h = self.linear2(h)
        logits = self.classifier(h[:, 0, :])
        return logits

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

Methods

def forward(self, x) ‑> Callable[..., typing.Any]

Expand source code

def forward(self, x):
    h = self.linear1(x)
    h = torch.relu(h)
    h = self.linear2(h)
    logits = self.classifier(h[:, 0, :])
    return logits

Define the computation performed at every call.

Should be overridden by all subclasses.

Note