matmul#

Matrix multiplication combines two matrices into a third: each element of A @ B is the dot product of a row of A with a column of B, a sum of products. It is the core operation of linear algebra and of neural networks, where linear layers and attention are built on it. Each output element is a long sum, so how far a float result lands from the exact one depends on how many terms there are, how large they are and the order an implementation adds them in.

In ExactRef, er.matmul(A, B), also exactref.linalg.products.matmul, is the matrix product, with NumPy's batching and 1-D promotion, as in torch.matmul. Both operands must have the same dtype.

Arguments#

Argument

Meaning

Default

A, B

the operands: arrays, tensors or earlier results with the same dtype, matrices or batches of them; a 1-D operand is promoted as in NumPy

required

dtype

the output format, of the same kind as the inputs

the input 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: the values as a float32 array
def f32(values):
    return er.asarray(values, "float32")

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

# ExactRef's reference for A @ B
r = er.matmul(A, B)
# Each element is the exact sum of products, rounded once to float32
print("rounded:", r.rounded.tolist())
# guarantee states what rounded satisfies
print("guarantee:", r.guarantee)
Output
rounded: [[0.7000000476837158, 1.0], [1.5, 2.200000047683716]]
guarantee: correctly_rounded

Shapes#

Batch dimensions broadcast, and a 1-D operand is promoted to a matrix and the added dimension removed afterwards, as in NumPy and PyTorch:

Code
# Shape (2, 3, 4): a batch of two 3x4 matrices of ones
batch = f32([[[1.0] * 4] * 3] * 2)
# Each matrix in the batch times one 4x5 matrix of ones
print("batch shape:", er.matmul(batch, f32([[1.0] * 5] * 4)).rounded.shape)

# A 1-D left operand acts as a row: [1, 2] @ A
print("vector @ A:", er.matmul(f32([1.0, 2.0]), A).rounded.tolist())
# A 1-D right operand acts as a column: A @ [1, 1]
print("A @ vector:", er.matmul(A, f32([1.0, 1.0])).rounded.tolist())
Output
batch shape: (2, 3, 5)
vector @ A: [7.0, 10.0]
A @ vector: [3.0, 7.0]

Output format and rounding#

dtype rounds the exact product once to another format of the same kind; rounding chooses the rounding:

Code
# The exact product rounded once to float16
print("float16:")
for row in er.matmul(A, B, dtype="float16").rounded.tolist():
    print(row)

# The exact product rounded once to float64
print("float64:")
for row in er.matmul(A, B, dtype="float64").rounded.tolist():
    print(row)

# The exact product rounded once to float32, upward (toward +inf)
print("toward +inf:")
for row in er.matmul(A, B, rounding=er.Rounding.TOWARD_POSITIVE).rounded.tolist():
    print(row)
Output
float16:
[0.7001953125, 1.0]
[1.5, 2.19921875]
float64:
[0.700000025331974, 1.0000000149011612]
[1.5000000521540642, 2.2000000327825546]
toward +inf:
[0.7000000476837158, 1.0000001192092896]
[1.5000001192092896, 2.200000047683716]

The float64 result shows the exact products of the float32 inputs: 0.1 in float32 is not exactly 0.1.

Integer inputs#

Integer products are exact, so the guarantee is "exact" and the tolerance 0. They wrap on overflow, as in PyTorch, unless overflow says otherwise:

Code
# Two int8 matrices whose product, 100 + 100 = 200, is past int8's largest value
x = er.asarray([[100, 100]], "int8")
y = er.asarray([[1], [1]], "int8")

# By default 200 wraps around to 200 - 256
print("wrap:", er.matmul(x, y).rounded.tolist())
# overflow="saturate" clamps to the int8 range instead
print("saturate:", er.matmul(x, y, overflow="saturate").rounded.tolist())

# overflow="error" raises instead
try:
    er.matmul(x, y, overflow="error")
except OverflowError as error:
    print(f"{type(error).__name__}: {error}")

# What rounded satisfies, and how far a correct output can be from it
print("guarantee:", er.matmul(x, y).guarantee)
print("tolerance:", er.matmul(x, y).tolerance().tolist())
Output
wrap: [[-56]]
saturate: [[127]]
OverflowError: 200 is outside the int8 range [-128, 127]
guarantee: exact
tolerance: [[0.0]]

Complex inputs#

Code
# A 1x2 and a 2x1 complex64 matrix
u = er.asarray([[1 + 1j, 0.5 - 1j]], "complex64")
w = er.asarray([[1 - 1j], [2 + 0.5j]], "complex64")

# ExactRef's reference for u @ w
z = er.matmul(u, w)
print("rounded:", z.rounded.tolist())
# One bound per part: the real part bounds the real error, the imaginary
# part the imaginary error
print("tolerance:", z.tolerance().tolist())
Output
rounded: [[(3.5-1.75j)]]
tolerance: [[(7.152557373046875e-07+4.76837158203125e-07j)]]

Special values#

Infinities and NaN follow IEEE rules on the exact sum, and each complex part follows them on its own real terms:

Code
inf = float("inf")

# The product inf * 0 is NaN, so the sum is NaN
print("[inf, 1] @ [0, 1]:",
      er.matmul(f32([[inf, 1.0]]), f32([[0.0], [1.0]])).rounded.tolist())
# inf * 1 + 1 * 1 is inf
print("[inf, 1] @ [1, 1]:",
      er.matmul(f32([[inf, 1.0]]), f32([[1.0], [1.0]])).rounded.tolist())

# The 1x1 complex matrix [inf + 0j]
c = er.asarray([[complex(inf)]], "complex64")
# The imaginary part is inf * 0 + 0 * 2, and inf * 0 is NaN
print("[inf+0j] @ [2+0j]:",
      er.matmul(c, er.asarray([[2 + 0j]], "complex64")).rounded.tolist())
Output
[inf, 1] @ [0, 1]: [[nan]]
[inf, 1] @ [1, 1]: [[inf]]
[inf+0j] @ [2+0j]: [[(inf+nanj)]]

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

# A 64x256 and a 256x32 matrix
P, Q = randn(64, 256), randn(256, 32)
# Could a correct implementation produce PyTorch's P @ Q? covers() answers per
# element, and all() asks whether every element is covered.
print("covered:", er.matmul(P, Q).covers(P @ Q).all())
Output
covered: True

Backward#

er.matmul.backward(A, B, grad), or er.matmul(A, B).backward(grad), returns grad_A = grad @ B^H and grad_B = A^H @ grad, each one matmul rounded once.

Code
# The gradient of the output: the 2x2 identity matrix
grad = f32([[1.0, 0.0], [0.0, 1.0]])
# The reference backward of A @ B, and the names of its outputs
g = er.matmul.backward(A, B, grad)
print("outputs:", list(g.rounded))
# A^H @ grad: with a real A and the identity, that is A transposed
print("grad_B:", g.rounded.grad_B.tolist())
# The same backward, run from the forward's result
print("result's grad_B:", er.matmul(A, B).backward(grad).rounded.grad_B.tolist())
Output
outputs: ['grad_A', 'grad_B']
grad_B: [[1.0, 3.0], [2.0, 4.0]]
result's grad_B: [[1.0, 3.0], [2.0, 4.0]]

Full specification#

er.matmul#

matmul(A, B, *, dtype=None, rounding: Rounding = Rounding.NEAREST_EVEN, overflow: str | None = None) -> Result#

Matrix product with NumPy batching and 1-D promotion.

Each output element is the exact sum of exact products, rounded once to dtype (default: the input dtype); exact holds the exact sums.

tolerance() budgets the required boundary, accumulation in the compute format, then the rounding to dtype. By default it covers normal accumulation: in index order or in 2 or 4 SIMD-style lanes, each product rounded and then added, starting a new partial sum where adding the next value would overflow (fused multiply-adds are covered too). standard="any_order" covers every order (any tree of additions, with or without fused multiply-add or rounding each product first) that does not overflow. ulps changes nothing here (matmul calls no library function), but must be at least 1/2, as everywhere. normal() is the output of normal accumulation in index order; for complex inputs each part is the sum of its 2k real products, taken in index order.

Non-finite inputs follow IEEE rules applied to the exact sum: NaN if any product is NaN (including inf * 0) or infinities of both signs occur, otherwise the infinity. An exact zero sum is +0, or -0 when rounding toward -inf; an empty sum (k = 0) is +0. Integer results wrap on overflow unless overflow says otherwise. For complex inputs each part follows these rules on its own real terms (a_re b_re and -a_im b_im for the real part, a_re b_im and a_im b_re for the imaginary part): (inf + 0j) @ (2 + 0j) is inf + nanj, as in PyTorch's batched product.

A and B may also be earlier results, as Result describes.

er.matmul.backward#

backward(A, B, grad, *, dtype=None, rounding: Rounding = Rounding.NEAREST_EVEN, overflow: str | None = None) -> Result#

Gradients of matmul(A, B) for the output gradient grad.

grad_A = grad @ B^H and grad_B = A^H @ grad (^H: conjugate transpose, as PyTorch uses for complex inputs), each summed over the batch copies that broadcasting made of its operand, with that operand's shape. Each gradient is one matmul: the broadcast copies are folded into its contraction, so every element is one exact sum rounded once, and exact, tolerance() and normal() are that matmul's. PyTorch folds the same way except when both operands keep batch dimensions that broadcast against each other; then it rounds each batch's product to the dtype before summing the copies, an extra rounding this tolerance does not budget. grad must have the output's shape and the inputs' dtype. A, B and grad may also be earlier results.

See also#