PyTorch equivalents#

ExactRef's operations are mathematical operations, not a copy of PyTorch's API. This page lists the ExactRef call that checks each common PyTorch call, and where PyTorch rounds more often than the operation requires. Those extra roundings are described with a chain, any_of, an algorithm or an option. In-place and out= variants map the same way as the plain call.

Matrix products#

PyTorch	ExactRef
`A @ B`	`er.matmul(A, B)`
`torch.matmul(A, B)`
`torch.mm(A, B)`
`torch.bmm(A, B)`
`torch.mv(A, B)`
`torch.dot(A, B)`
`torch.addmm(C, A, B, beta=b, alpha=a)`	`er.gemm(A, B, C, alpha=a, beta=b)`
`torch.addmv(C, A, B, beta=b, alpha=a)`
`torch.baddbmm(C, A, B, beta=b, alpha=a)`
`torch.addr(C, v, w, beta=b, alpha=a)`	`er.gemm(v.reshape((m, 1)), w.reshape((1, n)), C, alpha=a, beta=b)`
`torch.addbmm(C, A, B, beta=b, alpha=a)`	a loop of `er.gemm` calls, below

er.matmul follows NumPy's batching and 1-D promotion, as torch.matmul does. er.gemm takes alpha and beta exactly and broadcasts C to the output's shape; with beta = 0 it does not read C, and with alpha = 0 neither A nor B, as in PyTorch and BLAS. For addr, v and w are ExactRef arrays or earlier results, which have reshape.

Where PyTorch rounds more#

addbmm rounds each batch's product before adding the next one. That is a loop of gemm calls on earlier results:

Code

y = C
for i in range(batches):
    # Add batch i's product to the stored result, rounding once; beta scales
    # only C, in the first call
    y = er.gemm(A[i], B[i], y, alpha=alpha, beta=beta if i == 0 else 1)

er.gemm(A, B, C, alpha=alpha, beta=beta, reduce_batch=True) is the exact form instead: every element is one sum over all batches, rounded once.

Scaling an operand first. An implementation that scales and stores an operand before the product is a chain:

Code

# alpha * A stored first
er.gemm(er.multiply(A, alpha), B, C)

# beta * C stored first
er.gemm(A, B, er.multiply(C, beta), alpha=alpha)

addmv in bfloat16 and float16 may store beta * C in the dtype before the addition, and, with the matrix stored transposed, also alpha * (A @ x). Each form is a chain of er.multiply, er.gemm with beta=0 and er.add, and er.any_of covers both forms.

addr in bfloat16 and float16 computes in the dtype itself. Check it with covers(out, "bfloat16") or covers(out, "float16").

Batched products in the backward. When both operands keep batch dimensions that broadcast against each other, PyTorch rounds each batch's product before summing the copies. The default tolerance does not budget that.

TF32. A GPU in TF32 mode rounds float32 inputs to TF32 before the product. Round the inputs the same way and pass those values to er.matmul: TF32 values are float32 values.

Backend parameters#

Parameter	Express it as
`structure`	ordinary matrices
`offs` (grouped products)	one call per group
`trans_a`, `trans_b`	`op(A)` and `op(B)`; `transpose` works on an earlier result
`triangle`	a full product, comparing only the triangle the implementation computes

Elementwise operations and sums#

PyTorch	ExactRef
`x * y`	`er.multiply(x, y)`
`torch.mul(x, y)`	`er.multiply(x, y)`
`x + y`	`er.add(x, y)`
`torch.add(x, y)`	`er.add(x, y)`
`torch.add(x, y, alpha=a)`	`er.add(x, y, alpha=a)`, with `a` converted as below
`torch.sum(x, dim, keepdim=k)`	`er.sum(x, dim, keepdims=k)`
`torch.sum(x, dtype=d)`	`er.sum(er.multiply(x, 1, dtype=d))`

Broadcasting follows NumPy, and mixed formats are promoted as PyTorch promotes them (bfloat16 with float16 gives float32). An operand of shape () takes part in the promotion only when every operand has shape (), and one in a wider format is first rounded once to the promoted format, as PyTorch casts it. A Python scalar operand is taken exactly and does not change the format. Integer and bool inputs to sum give int64, as in PyTorch. torch.sum(x, dtype=d) converts x first, which is the multiply by 1.

Where PyTorch rounds more#

add with alpha converts alpha before using it: PyTorch's CPU kernels convert it to the tensor's dtype, and its CUDA kernels to the compute type (float32 for float16 and bfloat16). Pass the converted value. On the CPU:

dtype	ExactRef
float32	`er.add(x, y, alpha=a)`
bfloat16	`er.add(x, er.multiply(y, a, dtype="bfloat16"))`, with `alpha * y` stored first
float16	`er.any_of` of both forms: some elements take each

Code

import random
import torch
import exactref as er

# Seeded, so every run uses the same values
rng = random.Random(0)
# Two float16 tensors of 1000 normal values, mean 0 and standard deviation 1
x = torch.tensor([rng.gauss(0, 1) for _ in range(1000)], dtype=torch.float16)
y = torch.tensor([rng.gauss(0, 1) for _ in range(1000)], dtype=torch.float16)

# Alpha as the CPU kernel uses it: 0.3 rounded to float16
a = torch.tensor(0.3, dtype=x.dtype).item()
print("alpha:", a)

# x + a * y, rounded once to float16
plain = er.add(x, y, alpha=a)
# a * y rounded to float16 first, then added to x
stored = er.add(x, er.multiply(y, a, dtype="float16"))

# er.any_of accepts an output of either form, element by element
print("covered:", er.any_of(plain, stored).covers(torch.add(x, y, alpha=0.3)).all())

Output

alpha: 0.300048828125
covered: True

Broadcast operands in the backward. In the backward of add and multiply with a broadcast operand, PyTorch may round each product before summing the copies, an extra rounding the default tolerance does not budget.

Special values#

Complex special values follow IEEE rules in each part, as in PyTorch: (inf + 0j) * (1 + 0j) is inf + nanj. Where mathematics gives an answer, ExactRef gives it even when PyTorch does not: er.add(1+1j, inf+1j, alpha=1j) is infj, where PyTorch, multiplying alpha * y out in complex arithmetic, gives nan+infj.

Attention#

PyTorch	ExactRef
`F.scaled_dot_product_attention(q, k, v, ...)`	`er.scaled_dot_product_attention(q, k, v, ...)`

The arguments and shapes are PyTorch's. ExactRef returns two outputs: output, and the per-row logsumexp that fused kernels return. Dropout is random and has no reference, so dropout_p must be 0.

Where PyTorch rounds differently#

The math path computes the probabilities P, normalized and rounded, before the product with V. No tolerance budgets that rounding.

Fused kernels such as FlashAttention round the weights to bfloat16, rescale per block, or normalize by a reciprocal: use an attention algorithm.

Pre-scaled inputs. The math path scales Q and K by the square root of the scale first. Pass the pre-scaled values with scale=1. For the backward, the reference is then the gradient for the pre-scaled inputs; the math path's final product by the square root of the scale is not covered.

Backend parameters#

Parameter	Express it as
bottom-right causal alignment	a bool mask
`window_size` (sliding windows)	a bool mask
`alibi_slopes` (ALiBi biases)	a float mask
`cum_seq_q`, `cum_seq_k`, `seqused_k` (variable lengths)	one call per sequence
`block_table` (paged layouts)	one call per sequence

Matrix exponential#

PyTorch	ExactRef
`torch.linalg.matrix_exp(A)`	`er.matrix_exp(A)`
its gradient	`er.matrix_exp.backward(A, grad)`, which is `er.matrix_exp_frechet(A^H, grad)`

Guarantees and limits says where PyTorch's and SciPy's results fall outside the tolerance.