from __future__ import annotations

import torch

from reference.torch_online_softmax import torch_online_softmax
from reference.torch_row_softmax import torch_row_softmax


def test_row_softmax_handles_large_values():
    x = torch.tensor([[10000.0, 10001.0, 9999.0]], dtype=torch.float32)
    out = torch_row_softmax(x)
    torch.testing.assert_close(out.sum(dim=1), torch.ones(1), atol=1e-6, rtol=1e-6)
    assert torch.isfinite(out).all()


def test_online_softmax_handles_large_negative_values():
    x = torch.tensor([[-10000.0, -9998.0, -9999.0]], dtype=torch.float32)
    out = torch_online_softmax(x)
    torch.testing.assert_close(out.sum(dim=1), torch.ones(1), atol=1e-6, rtol=1e-6)
    assert torch.isfinite(out).all()


def test_row_and_online_softmax_agree():
    x = torch.randn(10, 40) * 8.0
    torch.testing.assert_close(
        torch_row_softmax(x), torch_online_softmax(x), atol=1e-5, rtol=1e-5
    )