Check a PyTorch operation across dtypes#

The problem#

You maintain a PyTorch backend, or a build of PyTorch, and want to know that its operations are correct in every dtype they support: float32, float16 and bfloat16 at least. A float64 reference converted to the dtype hides exactly the dtype-specific rounding you want to check.

Why the usual approach falls short#

Each dtype rounds differently, and PyTorch sometimes rounds in places the mathematics does not require: it converts alpha before using it, stores some intermediate results in the dtype, and computes some products in the dtype itself. A single tolerance per dtype either misses those or fails them.

The solution#

Loop over dtypes and operations, and check each output with covers(). Use PyTorch equivalents to find the ExactRef call for each PyTorch call and the roundings it adds.

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 and dtype, filled with random values
def randn(*shape, dtype):
    # 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 rounds each value to dtype
    return torch.tensor(values, dtype=dtype).reshape(shape)

# The device under test; set it to "cuda" or "mps" for a GPU
device = "cpu"

for dtype in (torch.float32, torch.float16, torch.bfloat16):
    # Random inputs in this dtype: matrices A, B and C, vectors x and y
    A, B = randn(32, 64, dtype=dtype), randn(64, 16, dtype=dtype)
    C = randn(32, 16, dtype=dtype)
    x, y = randn(1000, dtype=dtype), randn(1000, dtype=dtype)
    # A copy of a tensor on the device
    on = lambda t: t.to(device)
    # Per operation: is the device's output, copied back to the CPU, covered?
    results = {
        "matmul": er.matmul(A, B).covers((on(A) @ on(B)).cpu()).all(),
        # er.gemm is beta * C + alpha * (A @ B), PyTorch's addmm
        "addmm": er.gemm(A, B, C, alpha=0.5, beta=2.0).covers(
            torch.addmm(on(C), on(A), on(B), beta=2.0, alpha=0.5).cpu()).all(),
        "sum": er.sum(x).covers(on(x).sum().cpu()).all(),
        "mul": er.multiply(x, y).covers((on(x) * on(y)).cpu()).all(),
    }
    print(str(dtype), "covered:", results)
Output
torch.float32 covered: {'matmul': True, 'addmm': True, 'sum': True, 'mul': True}
torch.float16 covered: {'matmul': True, 'addmm': True, 'sum': True, 'mul': True}
torch.bfloat16 covered: {'matmul': True, 'addmm': True, 'sum': True, 'mul': True}

torch.add with alpha needs the conversion PyTorch applies to alpha, and in float16 and bfloat16 one more form. The CPU kernel converts alpha to the tensor's dtype:

Code
for dtype in (torch.float32, torch.float16, torch.bfloat16):
    # Random vectors in this dtype
    x, y = randn(1000, dtype=dtype), randn(1000, dtype=dtype)
    # alpha = 0.3 rounded to the dtype, as the CPU kernel uses it
    a = torch.tensor(0.3, dtype=dtype).item()
    # The dtype's name without "torch.", such as "float16"
    name = str(dtype).removeprefix("torch.")

    # The plain form: x + a * y
    plain = er.add(x, y, alpha=a)
    # The stored form: a * y rounded to the dtype first, then added to x
    stored = er.add(x, er.multiply(y, a, dtype=name))
    # any_of widens plain's bounds to cover the stored form too
    covered = er.any_of(plain, stored).covers(torch.add(x, y, alpha=0.3)).all()
    print(name, "covered:", covered)
Output
float32 covered: True
float16 covered: True
bfloat16 covered: True

When a check fails on a backend, move the inputs back to the CPU and look at the failing elements: is it the backend, or a rounding the mapping does not describe yet?