Backward

Every operation with a derivative has a backward operation with the same guarantees: each gradient is the exact value rounded once, with a tolerance and, where the forward has one, a normal answer.

Calling a backward

There are two ways. The standalone .backward takes the forward's inputs and the output gradient. A forward result's backward(grad) reuses its inputs:

Code

import exactref as er

# A 2x2 and a 2x1 float32 matrix, and the gradient of their 2x1 product
A = er.asarray([[1.0, 2.0], [3.0, 4.0]], "float32")
B = er.asarray([[0.5], [0.25]], "float32")
grad = er.asarray([[1.0], [1.0]], "float32")

# The standalone backward: the forward's inputs, then the output gradient
g = er.matmul.backward(A, B, grad)
# Each gradient, rounded once to float32, by name
print("grad_A:", g.rounded.grad_A.tolist())
print("grad_B:", g.rounded.grad_B.tolist())

# The same grad_B from the forward result's backward, which reuses A and B
print("grad_B again:", er.matmul(A, B).backward(grad).rounded.grad_B.tolist())

Output

grad_A: [[0.5, 0.25], [0.5, 0.25]]
grad_B: [[4.0], [6.0]]
grad_B again: [[4.0], [6.0]]

The gradient is required; there is no default of ones. Each output is named grad_ plus the input's name:

Backward	Gradients
er.matmul.backward(A, B, grad)	grad_A, grad_B
er.gemm.backward(A, B, C, grad, alpha=, beta=)	grad_A, grad_B, grad_C
er.add.backward(x, y, grad, alpha=)	grad_x, grad_y
er.multiply.backward(x, y, grad)	grad_x, grad_y
er.sum.backward(x, grad, axis=)	grad_x
er.scaled_dot_product_attention.backward(q, k, v, grad)	grad_query, grad_key, grad_value
er.matrix_exp.backward(A, grad)	grad_A

matrix_exp_frechet has no backward. grad must have the forward output's shape. A scalar operand of add or multiply gets no gradient.

Checking against PyTorch's autograd

Run PyTorch's backward, then compare each gradient with covers():

Code

import math
import random
import torch

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

# Like torch.randn, but drawing from rng
def randn(*shape):
    # 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).reshape(shape)

# Random float32 inputs of a 16x32 by 32x8 product, and its output gradient
A, B, grad = randn(16, 32), randn(32, 8), randn(16, 8)

# Copies of A and B that record gradients for PyTorch's autograd
A_, B_ = A.clone().requires_grad_(), B.clone().requires_grad_()
# PyTorch's backward, which stores the gradients in A_.grad and B_.grad
(A_ @ B_).backward(grad)

# ExactRef's gradients for the same inputs and output gradient
g = er.matmul.backward(A, B, grad)
# g["grad_A"] is a result holding just grad_A, so it has covers()
print("grad_A covered:", g["grad_A"].covers(A_.grad).all())
print("grad_B covered:", g["grad_B"].covers(B_.grad).all())

Output

grad_A covered: True
grad_B covered: True

Attention works the same way:

Code

# Query, key and value of shape (batch, heads, positions, features)
q, k, v = randn(1, 2, 8, 16), randn(1, 2, 8, 16), randn(1, 2, 8, 16)
# Copies that record gradients for PyTorch's autograd
q_, k_, v_ = (t.clone().requires_grad_() for t in (q, k, v))
# PyTorch's attention with a causal mask: no query attends to a later key
out = torch.nn.functional.scaled_dot_product_attention(q_, k_, v_, is_causal=True)
# A random output gradient, and PyTorch's backward from it
grad = randn(1, 2, 8, 16)
out.backward(grad)

# ExactRef's forward, then its backward from the same inputs
g = er.scaled_dot_product_attention(q, k, v, is_causal=True).backward(grad)
# Each gradient's name, with the copy whose .grad PyTorch filled
names = [("grad_query", q_), ("grad_key", k_), ("grad_value", v_)]
for name, t in names:
    # g[name] holds one gradient, and covers() judges PyTorch's against it
    print(f"{name} covered:", g[name].covers(t.grad).all())

Output

grad_query covered: True
grad_key covered: True
grad_value covered: True

What each backward computes

• matmul: grad_A = grad @ B^H and grad_B = A^H @ grad, where ^H is the conjugate transpose, as PyTorch uses for complex inputs. Each gradient is one matmul, with the copies that broadcasting made of its operand folded into the sum, so every element is rounded once and its tolerance and normal answer 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 before summing the copies, an extra rounding the default tolerance does not budget.

• gemm: 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.

• add: grad_x = grad and grad_y = conj(alpha) * grad, each summed over its operand's broadcast copies.

• multiply: grad_x = grad * conj(y) and grad_y = grad * conj(x), each summed over its operand's broadcast copies. Where an operand is broadcast, PyTorch rounds each product before summing the copies, an extra rounding the default tolerance does not budget.

• sum: grad_x holds, at every element of x, the gradient of the output it was summed into.

• scaled_dot_product_attention: with P the attention weights, dP = grad V^T, D = sum P * dP per row and dS = P * (dP - D), it returns grad_value = P^T grad, grad_query = scale * dS K and grad_key = scale * dS^T Q. The tolerance follows PyTorch's math-path autograd, and so do the special values. Dropout must be 0, and a float mask gets no gradient. A fused backward kernel rounds differently: its tolerance is AttentionAlgorithm.backward_tolerance.

• matrix_exp: grad_A = matrix_exp_frechet(A^H, grad), which is what PyTorch computes.

One case can raise instead of returning: attention's grad_key, grad_value, and grad_query summed over broadcast copies of a query row, are found exact only when every term is rational. Such a sum that is exactly a rounding boundary, or exactly zero by cancellation, raises ReferenceUnavailableError. See Guarantees and limits.