Replace hand-picked test tolerances#

The problem#

A test compares an implementation with a reference computed in float64, using torch.testing.assert_close or allclose with an rtol and atol someone picked once. The test is flaky: it fails on a correct implementation when the inputs cancel, and someone loosens atol until it passes. After that it also passes real bugs whose effect is smaller than the loosened tolerance.

Why the usual approach falls short#

A fixed relative and absolute tolerance cannot fit every input. How far a correct float32 result can land from the exact answer depends on the inputs: on cancellation, on the number of terms and on their size. A float64 reference is not exact either, and it says nothing about which float32 results are correct.

The solution#

Ask ExactRef, for the inputs you have, which outputs a correct implementation can produce. Here is a correct float32 loop that assert_close rejects:

Code
import torch
import exactref as er

# Three float32 values whose exact sum is 1
x = torch.tensor([1e8, 1.0, -1e8])
# A plain loop in float32: each addition rounds to float32
out = (x[0] + x[1]) + x[2]
print("loop:", out.item())

# The usual reference: the sum in float64, converted to float32
reference = x.double().sum().float()
# assert_close with its default rtol and atol for float32
try:
    torch.testing.assert_close(out, reference)
except AssertionError:
    print("assert_close fails")

# Could a correct float32 implementation produce out?
print("covered:", er.sum(x).covers(out).all())
Output
loop: 0.0
assert_close fails
covered: True

1e8 + 1.0 rounds back to 1e8 in float32, so the loop's 0.0 is a correct answer even though the exact sum is 1. And here is a bug that assert_close passes, on small-magnitude data:

Code
import math
import random

# Seeded, so every run uses the same values
rng = random.Random(0)

# A tensor of the given shape, filled with random values
def randn(*shape, scale=1.0):
    # One value per element: normal, with mean 0 and standard deviation scale
    values = [rng.gauss(0, scale) for _ in range(math.prod(shape))]
    # torch.tensor makes float32, PyTorch's default for Python floats
    return torch.tensor(values).reshape(shape)

# Random matrices: A is 64x256 and B is 256x32
A, B = randn(64, 256, scale=1e-3), randn(256, 32, scale=1e-3)
# The bug: every element is off by a constant 1e-6
buggy = A @ B + 1e-6

# Passes: buggy is within assert_close's default tolerance of a float64 reference
torch.testing.assert_close(buggy, (A.double() @ B.double()).float())

# Could a correct implementation produce every element of buggy?
print("covered:", er.matmul(A, B).covers(buggy).all())
Output
covered: False

The outputs are about 1e-5 in size, so assert_close's default atol=1e-5 hides an error of 1e-6. ExactRef's tolerance comes from these inputs, and it is far smaller.

In a test, replace the comparison with one line:

Code
# The reference for A @ B
r = er.matmul(A, B)
# Every element must be one a correct implementation can produce
assert r.covers(A @ B).all()

Get started has a short helper that lists the failing elements when the check fails.