"""Reproduce the published example and numerical checks on CPU; no weights download."""
import torch
from model import EvoformerCoupling, verify

torch.set_num_threads(1)
torch.manual_seed(0)
torch.use_deterministic_algorithms(True)
inputs_rng = torch.Generator(device="cpu").manual_seed(0)
with torch.no_grad():
    model = EvoformerCoupling().cpu().eval()
    inputs = [
        torch.randn((3, 4, 8), dtype=torch.float32, generator=inputs_rng),
        torch.randn((4, 4, 4), dtype=torch.float32, generator=inputs_rng),
    ]
    output = model(*inputs)
assert list(output.shape) == [4,4,4]
assert torch.isfinite(output).all()
print("Verified output shape:", list(output.shape))
for check in verify():
    print("Passed:", check)