29 lines
923 B
Python
29 lines
923 B
Python
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
|
|
)
|
|
|