gemm#

GEMM, short for general matrix multiply, is the classic BLAS routine that multiplies two matrices and adds the result to a third: D = beta * C + alpha * (A @ B), where alpha and beta are scalars. It is the computation behind a linear layer with a bias, and PyTorch exposes it as addmm, with addmv, addr and baddbmm for a vector, an outer product and batches. The scaling and the extra addition give an implementation more choices than a plain product: it may apply alpha to an operand, to each product or to the sum, and add beta * C first or last.

In ExactRef, er.gemm(A, B, C, alpha=1, beta=1), also exactref.linalg.products.gemm, computes beta * C + alpha * (A @ B) exactly and rounds once, with matmul's batching and 1-D promotion and C broadcast to the output's shape.

Arguments#

Argument

Meaning

Default

A, B

the matrices to multiply, with matmul's batching and 1-D promotion

required

C

the matrix added, broadcast to the output's shape

required

alpha

the factor on A @ B, taken exactly

1

beta

the factor on C, taken exactly

1

reduce_batch

sum the output's batch dimensions into one contraction

False

dtype

the output format, of the same kind as the inputs

the operands' promoted format

rounding

the rounding of the reference and of the final step into the output format

er.Rounding.NEAREST_EVEN

overflow

what a result past the output format's range becomes: "ieee", "saturate" or "nan" for floats; "wrap", "error" or "saturate" for integers

each format's own rule; integers wrap

Examples#

Code
import exactref as er

# A shorthand for these examples
def f32(values):
    # er.asarray imports the values as a float32 ExactRef array
    return er.asarray(values, "float32")

# Three 2x2 float32 matrices
A = f32([[1.0, 2.0], [3.0, 4.0]])
B = f32([[0.1, 0.2], [0.3, 0.4]])
C = f32([[1.0, 1.0], [1.0, 1.0]])

# 2 * C + 0.5 * (A @ B), rounded to float32, the operands' format
print("rounded:", er.gemm(A, B, C, alpha=0.5, beta=2).rounded.tolist())
Output
rounded: [[2.3499999046325684, 2.5], [2.75, 3.0999999046325684]]

Each element is beta * C + alpha * (A @ B) computed exactly, then rounded once. alpha and beta are taken exactly as given.

Broadcasting C, and the addmv and addr shapes#

Code
# C as one row, [10, 20], added to each row of A @ B
print("broadcast C:", er.gemm(A, B, f32([10.0, 20.0])).rounded.tolist())

# The addmv shape: A times a vector, plus a vector C
print("addmv:", er.gemm(A, f32([1.0, 1.0]), f32([0.5, 0.5])).rounded.tolist())

# v as a 2x1 column and w as a 1x3 row
v = f32([1.0, 2.0]).reshape((2, 1))
w = f32([3.0, 4.0, 5.0]).reshape((1, 3))
# The addr shape: the outer product of v and w, plus a 2x3 C of zeros
print("addr:", er.gemm(v, w, er.Array.zeros("float32", (2, 3))).rounded.tolist())
Output
broadcast C: [[10.699999809265137, 21.0], [11.5, 22.200000762939453]]
addmv: [3.5, 7.5]
addr: [[3.0, 4.0, 5.0], [6.0, 8.0, 10.0]]

beta = 0 and alpha = 0#

With beta = 0, C is not read, and with alpha = 0, neither A nor B is, as in PyTorch and BLAS, so a NaN there does not reach the output:

Code
# A 2x2 float32 matrix of NaNs
nans = f32([[float("nan")] * 2] * 2)

# beta=0: the output is alpha * (A @ B) alone
print("beta=0:", er.gemm(A, B, nans, beta=0).rounded.tolist())
# alpha=0: the output is beta * C alone, here C itself
print("alpha=0:", er.gemm(nans, B, C, alpha=0).rounded.tolist())
Output
beta=0: [[0.7000000476837158, 1.0], [1.5, 2.200000047683716]]
alpha=0: [[1.0, 1.0], [1.0, 1.0]]

Summing over batches: reduce_batch#

reduce_batch=True adds the output's batch dimensions into the sum, so each element is one exact sum over all batches, rounded once. PyTorch's addbmm rounds after each batch instead, which is a loop of gemm calls on earlier results (PyTorch equivalents).

Code
# Two batches: 1x2 matrices for A, 2x1 matrices for B
batches_A = f32([[[1.0, 2.0]], [[3.0, 4.0]]])
batches_B = f32([[[1.0], [1.0]], [[1.0], [1.0]]])

# Both batches' products in one sum, plus a 1x1 C of zero
r = er.gemm(batches_A, batches_B, f32([[0.0]]), reduce_batch=True)
print("rounded:", r.rounded.tolist())
Output
rounded: [[10.0]]

The normal answer#

normal() adds the products in index order, then computes fl(fl(alpha * sum) + fl(beta * C)) in the compute format, the reference BLAS order:

Code
# The first example's normal answer, each step rounded to float32
print("normal:", er.gemm(A, B, C, alpha=0.5, beta=2).normal().tolist())
Output
normal: [[2.3499999046325684, 2.5], [2.75, 3.0999999046325684]]

Checking PyTorch#

Code
import math
import random
import torch

# 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)

# Random matrices: P is 32x64, Q is 64x16 and R is 32x16
P, Q, R = randn(32, 64), randn(64, 16), randn(32, 16)
# PyTorch's 0.5 * R + 2.0 * (P @ Q)
out = torch.addmm(R, P, Q, beta=0.5, alpha=2.0)

# Could a correct implementation produce every element of out?
print("covered:", er.gemm(P, Q, R, alpha=2.0, beta=0.5).covers(out).all())
Output
covered: True

Backward#

er.gemm.backward(A, B, C, grad, alpha=, beta=) returns grad_A = conj(alpha) grad @ B^H, grad_B = conj(alpha) A^H @ grad and grad_C = conj(beta) grad, each summed over its operand's broadcast copies and rounded once.

Code
# An output gradient of ones
ones = f32([[1.0, 1.0], [1.0, 1.0]])
# The gradients of gemm(A, B, C, beta=2) for that output gradient
g = er.gemm.backward(A, B, C, ones, beta=2)

# g has several outputs; listing g.rounded gives their names
print("outputs:", list(g.rounded))
# grad_C is beta times the output gradient: 2 everywhere
print("grad_C:", g.rounded.grad_C.tolist())
Output
outputs: ['grad_A', 'grad_B', 'grad_C']
grad_C: [[2.0, 2.0], [2.0, 2.0]]

Full specification#

er.gemm#

gemm(A, B, C, *, alpha=1, beta=1, reduce_batch: bool = False, dtype=None, rounding: Rounding = Rounding.NEAREST_EVEN, overflow: str | None = None) -> Result#

D = beta C + alpha (A @ B), with matmul's batching and 1-D promotion and C broadcast to the output's shape (PyTorch's addmm, addmv, addr, baddbmm).

Each output element is the exact value rounded once to dtype (default: the operands' promoted format, as add promotes); alpha and beta are taken exactly as given. With beta = 0 C is not read, and with alpha = 0 neither A nor B is, as in PyTorch and BLAS. With reduce_batch=True the output's batch dimensions are summed into the contraction, so each element is one exact sum over all batches (addbmm's exact form; PyTorch's addbmm rounds after each batch, which is a loop of gemm calls on earlier results). addr(C, v, w) is gemm with v as a column (m, 1) and w as a row (1, n).

tolerance() budgets the products rounded or fused in the compute format, alpha applied once on every product's way (to an operand, a product, a partial sum or the whole sum), beta C rounded once and added at any point (first, as PyTorch's CPU kernels do, or last), the additions in normal accumulation, and the final rounding; standard="any_order" allows any order of additions. With an alpha whose imaginary part is nonzero the normal tolerance is the any-order one. normal() adds the products in index order, then computes fl(fl(alpha * sum) + fl(beta * C)) in the compute format (the reference BLAS order). Integer results wrap on overflow unless overflow says otherwise. A non-finite input follows IEEE rules on the exact expression; for complex values in each part separately, as PyTorch's batched products do, with alpha and beta taken as exact coefficients: a zero part of a nonzero alpha or beta adds no term. A, B and C may also be earlier results.

er.gemm.backward#

backward(A, B, C, grad, *, alpha=1, beta=1, reduce_batch: bool = False, dtype=None, rounding: Rounding = Rounding.NEAREST_EVEN, overflow: str | None = None) -> Result#

Gradients of gemm(A, B, C, alpha=alpha, beta=beta) for the output gradient grad.

grad_A = conj(alpha) grad @ B^H, grad_B = conj(alpha) A^H @ grad and grad_C = conj(beta) grad, each summed over the copies that broadcasting (and reduce_batch) made of its operand, in that operand's shape and dtype. grad_A and grad_B are gemm calls with beta = 0 on matmul's backward operands, so each element is one exact value rounded once; PyTorch's product followed by its scaling is covered. grad_C is one product per element, as add's gradients, and takes their zero signs, including where a complex beta other than 1 differs from PyTorch. grad must have the output's shape. A, B, C and grad may also be earlier results.

See also#