"""Reproduce the published example and numerical checks on CPU; no weights download."""
import torch
from model import AdditiveDecoderStep, 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 = AdditiveDecoderStep().cpu().eval()
    inputs = [
        torch.ones((1, 5), dtype=torch.int64),
        torch.randn((1, 8), dtype=torch.float32, generator=inputs_rng),
        torch.ones((1,), dtype=torch.int64),
    ]
    output = model(*inputs)
assert list(output.shape) == [1,16]
assert torch.isfinite(output).all()
print("Verified output shape:", list(output.shape))
for check in verify():
    print("Passed:", check)