File size: 1,348 Bytes
2ea2835 | 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 28 29 30 31 32 33 34 35 | import pytest
import torch
from model.wan_kv_mmd import _accumulate_fp32_pixel_cotangent
def test_pixel_cotangents_accumulate_immediately_in_fp32():
accumulated = None
shape = None
expected = torch.zeros(2, 3, 4, dtype=torch.float32)
norms = []
for index, dtype in enumerate(
(torch.bfloat16, torch.float16, torch.float32), start=1
):
value = torch.full((2, 3, 4), float(index), dtype=dtype)
expected.add_(value.float())
accumulated, shape, norm = _accumulate_fp32_pixel_cotangent(
accumulated, value, shape, branch_name=f"branch{index}")
norms.append(norm)
assert accumulated.dtype == torch.float32
torch.testing.assert_close(accumulated, expected)
assert shape == (2, 3, 4)
assert all(item.ndim == 0 for item in norms)
def test_pixel_cotangent_accumulation_rejects_shape_and_nonfinite():
first, shape, _ = _accumulate_fp32_pixel_cotangent(
None, torch.ones(1, 2), None, branch_name="first")
with pytest.raises(ValueError, match="shapes differ"):
_accumulate_fp32_pixel_cotangent(
first, torch.ones(1, 3), shape, branch_name="second")
with pytest.raises(FloatingPointError, match="non-finite"):
_accumulate_fp32_pixel_cotangent(
None, torch.tensor([float("nan")]), None, branch_name="bad")
|