Module neuralops.tests.conftest

Shared pytest fixtures for neuralops test suite.

Functions

def attention_registry()

Expand source code

@pytest.fixture(scope="session")
def attention_registry():
    from neuralops.core.registry import ATTENTION_REGISTRY
    import neuralops.attentions  # noqa: F401
    return ATTENTION_REGISTRY

def compressor_registry()

Expand source code

@pytest.fixture(scope="session")
def compressor_registry():
    from neuralops.core.registry import COMPRESSOR_REGISTRY
    import neuralops.compression_engine  # noqa: F401
    return COMPRESSOR_REGISTRY

def cpu_device()

Expand source code

@pytest.fixture
def cpu_device():
    return torch.device("cpu")

def device()

Expand source code

@pytest.fixture
def device():
    """Return available device (cuda > mps > cpu)."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")

Return available device (cuda > mps > cpu).

def dtype(request)

Expand source code

@pytest.fixture(params=[torch.float32, torch.float16, torch.bfloat16])
def dtype(request):
    """Parametrized dtype fixture."""
    if request.param == torch.bfloat16:
        if not torch.cuda.is_available():
            pytest.skip("bfloat16 requires CUDA")
        if not torch.cuda.is_bf16_supported():
            pytest.skip("bfloat16 not supported on this GPU")
    return request.param

Parametrized dtype fixture.

def dummy_config()

Expand source code

@pytest.fixture
def dummy_config():
    """Minimal config for testing attention modules."""
    class DummyConfig:
        hidden_size = 64
        num_attention_heads = 4
        num_key_value_heads = 4
        head_dim = 16
        dropout = 0.0
        is_decoder = False
        model_type = "bert"
    return DummyConfig()

Minimal config for testing attention modules.

def dummy_hidden_states()

Expand source code

@pytest.fixture
def dummy_hidden_states():
    """Random hidden states for forward passes."""
    return torch.randn(2, 16, 64)

Random hidden states for forward passes.

def dummy_qkv()

Expand source code

@pytest.fixture
def dummy_qkv():
    """Random Q, K, V tensors."""
    B, H, N, D = 2, 4, 16, 16
    return (
        torch.randn(B, H, N, D),
        torch.randn(B, H, N, D),
        torch.randn(B, H, N, D),
    )

Random Q, K, V tensors.

def repo_root()

Expand source code

@pytest.fixture
def repo_root():
    """Absolute path to repository root."""
    return Path(__file__).resolve().parents[2]

Absolute path to repository root.

def tmp_results_dir(tmp_path)

Expand source code

@pytest.fixture
def tmp_results_dir(tmp_path):
    """Temporary results directory."""
    d = tmp_path / "results"
    d.mkdir()
    return d

Temporary results directory.