tests/conftest.py

import pytest
import torch


@pytest.fixture(scope="session")
def device():
    return "cuda" if torch.cuda.is_available() else "cpu"


@pytest.fixture(scope="session")
def dtype():
    return torch.float32


def make_attn_weights(B, H, N_total, device):
    """Create valid softmax attention weights."""
    torch.manual_seed(42)
    A = torch.rand(B, H, N_total, N_total, device=device)
    return A / A.sum(dim=-1, keepdim=True)


def make_softmax_matrix(B, N_text, N_vis, device):
    """Create valid softmax matrix simulating attention maps."""
    torch.manual_seed(0)
    P = torch.rand(B, N_text, N_vis, device=device)
    return P / P.sum(dim=-1, keepdim=True)