Chains and any_of#

An implementation of a computation with several steps stores each step's output in its dtype, then rounds again in the next step. To check it, run the same steps in ExactRef, passing each earlier result where the implementation used its stored output. The bounds then cover every value the implementation could have stored.

Passing a result as an input#

Code
import math
import random
import torch
import exactref as er

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

# Like torch.randn, but drawing from rng
def randn(*shape, dtype=torch.float32):
    # Normally distributed, mean 0 and standard deviation 1, one per element
    values = [rng.gauss(0, 1) for _ in range(math.prod(shape))]
    return torch.tensor(values, dtype=dtype).reshape(shape)

# Three random 8x8 float32 matrices
A, B, C = randn(8, 8), randn(8, 8), randn(8, 8)

# The reference for A @ B
first = er.matmul(A, B)
# A @ B stored in float32, then @ C
chain = er.matmul(first, C)
# Can a correct implementation of the chain give PyTorch's (A @ B) @ C?
print("covered:", chain.covers((A @ B) @ C).all())
Output
covered: True

What the chained result holds:

  • chain.rounded and chain.exact are the exact product of first.rounded and C, rounded once. They are not the exact value of the whole expression (A @ B) @ C.

  • chain.tolerance() covers every stored A @ B within first.tolerance() of first.rounded, so it is wider than if the stored values were known:

Code
# As if A @ B had been stored as exactly first.rounded
plain = er.matmul(first.rounded, C)
# Whether the chain's largest tolerance is larger than plain's
wider = max(chain.tolerance().scalars()) > max(plain.tolerance().scalars())
print("chain wider:", wider)
Output
chain wider: True

  • chain.normal() is normal accumulation at both steps: it works on first.normal().

How the bounds carry through#

The earlier bounds are taken with the same options: the normal tolerance assumes normal accumulation at every step, the any-order tolerance any order at every step, and compute and flush_to_zero apply to every step. What the next step's bounds assume about each stored value:

The earlier result allows

The next step's bounds

a finite value within its tolerance

cover every such value

an infinity or NaN that may_be() allows

carry it on in may_be(): through an exact coefficient (add's alpha, gemm's beta) with that coefficient's sign, through data as every flag

any value (an infinite tolerance)

give no bound: an infinite tolerance wherever the reference is finite

a possible infinity or NaN, into attention or matrix_exp

give no bound, because exp can turn it back into a finite value

Reading complicated tolerances has an example of the last case. A named algorithm's run() uses the earlier rounded.

Every operation, backward and algorithm accepts results in place of arrays.

One output of several#

A result with several outputs is passed one output at a time, by name:

Code
# The gradients of A @ B for a random output gradient: grad_A and grad_B
g = er.matmul.backward(A, B, randn(8, 8))
# g["grad_A"] is a result holding just grad_A, passed in place of an array
print("shape:", er.matmul(g["grad_A"], C).rounded.shape)
Output
shape: (8, 8)

Attention algorithms take a forward's outputs the same way, as fwd["output"] and fwd["logsumexp"].

Views#

When the next step needs another layout, reshape and transpose move the elements of a one-output result without rounding anything. For example, an outer product of two earlier results, as a column times a row:

Code
# The sums of A's rows, shape (8,)
u = er.sum(A, axis=1)
# The sums of B's columns, shape (8,)
w = er.sum(B, axis=0)
# u as an 8x1 column times w as a 1x8 row
outer = er.matmul(u.reshape((8, 1)), w.reshape((1, 8)))
print("shape:", outer.rounded.shape)
Output
shape: (8, 8)

Either of several expressions: any_of#

Some implementations may compute one of several expressions, for example either association of a product of three matrices. er.any_of(center, *alternatives) covers all of them:

Code
# Seeded again, so these matrices do not depend on the draws above
rng = random.Random(0)
# Three random 4x4 float16 matrices
X, Y, Z = (randn(4, 4, dtype=torch.float16) for _ in range(3))

# (X @ Y) @ Z and X @ (Y @ Z), each storing its first product in float16
left = er.matmul(er.matmul(X, Y), Z)
right = er.matmul(X, er.matmul(Y, Z))
# Whether left's bounds cover PyTorch's X @ (Y @ Z)
print("left covers X @ (Y @ Z):", left.covers(X @ (Y @ Z)).all())

# left, with its bounds widened to cover right as well
either = er.any_of(left, right)
print("either covers (X @ Y) @ Z:", either.covers((X @ Y) @ Z).all())
print("either covers X @ (Y @ Z):", either.covers(X @ (Y @ Z)).all())
Output
left covers X @ (Y @ Z): False
either covers (X @ Y) @ Z: True
either covers X @ (Y @ Z): True

The result is the center's (here left): its rounded, exact, guarantee, normal() and backward(). Only its bounds are widened. Per element, its finite range is the smallest range holding the center's and every alternative's, its infinities and NaN are those any of them allows, and its tolerance is the largest distance from the center's rounded to an end of that range. Every result passed to any_of must have one output, with the same shape and dtype.

PyTorch equivalents lists PyTorch operations whose roundings are chains or any_of of ExactRef calls.