Find the step that loses accuracy#

The problem#

A computation with several steps, such as a small network layer, fails an end-to-end accuracy check. It does X @ W1 + b1, then a product with W2, then a sum. Some step loses more accuracy than it should, and comparing the final output says only that something is wrong.

Why the usual approach falls short#

The error of the final output mixes the roundings of every step, so a check of the output alone cannot say which step is responsible. Checking each step against a float64 version of that step has the usual problems of a hand-picked tolerance, and it ignores that each step works on the previous step's stored, rounded output.

The solution#

Run the same steps in ExactRef, passing each step's result into the next (Chains and any_of), and check every intermediate the implementation stores. Here the implementation has a hidden bug: it stores the first step's output through bfloat16.

Code
import math
import random
import torch
import exactref as er

# 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):
    # One value per element: normal, with mean 0 and standard deviation 1
    values = [rng.gauss(0, 1) for _ in range(math.prod(shape))]
    # torch.tensor makes float32, PyTorch's default for Python floats
    return torch.tensor(values).reshape(shape)

# 8 input rows X, and the layer's weights W1 and W2 and bias b1
X, W1, b1, W2 = randn(8, 32), randn(32, 64), randn(64), randn(64, 16)

# Step 1: X @ W1 + b1. The bug: a round trip through bfloat16.
h = torch.addmm(b1, X, W1).bfloat16().float()
# Step 2: the product with W2
y = h @ W2
# Step 3: the sum of each row
s = y.sum(dim=1)

The ExactRef chain mirrors the steps:

Code
# Step 1: er.gemm(X, W1, b1) is X @ W1 + b1, as in torch.addmm
step1 = er.gemm(X, W1, b1)
# Steps 2 and 3 take the previous step's result, so their bounds cover
# every value that step's bounds allow the implementation to store
step2 = er.matmul(step1, W2)
step3 = er.sum(step2, axis=1)

# The end-to-end check
print("step 3 covered:", step3.covers(s).all())
# Check each intermediate the implementation stores
print("step 1 covered:", step1.covers(h).all())
print("step 2 covered:", step2.covers(y).all())
Output
step 3 covered: False
step 1 covered: False
step 2 covered: False

The first failing step is step 1, so the problem is in how h is computed or stored. To confirm that step 2 is fine on its own, check it on the implementation's actual stored input:

Code
# Step 2 alone, on the h the implementation stored
print("step 2 on h covered:", er.matmul(h, W2).covers(y).all())
Output
step 2 on h covered: True

A step that rounds in a known extra place, such as a deliberate bfloat16 store, is not a bug: describe it as a step of its own, for example with er.multiply(..., 1, dtype="bfloat16"), so the chain budgets it.