File size: 670 Bytes
176b11a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | 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)
|