numerically sound mean computation

This commit is contained in:
Alexander Soare
2024-04-03 09:44:20 +01:00
parent 7242953197
commit e9eb262293
2 changed files with 27 additions and 17 deletions

View File

@@ -58,9 +58,12 @@ def test_compute_stats():
for k, pattern in buffer.stats_patterns.items():
expected_mean = einops.reduce(all_data[k], pattern, "mean")
assert torch.allclose(computed_stats[k]["mean"], expected_mean)
assert torch.allclose(
computed_stats[k]["std"],
torch.sqrt(einops.reduce((all_data[k] - expected_mean) ** 2, pattern, "mean"))
)
try:
assert torch.allclose(
computed_stats[k]["std"],
torch.sqrt(einops.reduce((all_data[k] - expected_mean) ** 2, pattern, "mean"))
)
except:
breakpoint()
assert torch.allclose(computed_stats[k]["min"], einops.reduce(all_data[k], pattern, "min"))
assert torch.allclose(computed_stats[k]["max"], einops.reduce(all_data[k], pattern, "max"))