Every op in a framework has two functions: a forward function and a backward function. The forward function computes the op's output from its inputs to make predictions, while the backward function computes the gradients needed to update the model's weights during training.

For example, consider a matmul op with inputs a and b and output c, as shown in the following snippet. The forward function takes a and b as inputs and produces c. The backward function takes the gradient of c (grad_c), along with the original inputs a and b, and produces the gradients of a and b (grad_a and grad_b).

import numpy as np

def matmul_forward(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    # a: (m, k), b: (k, n) -> c: (m, n)
    return a @ b

def matmul_backward(grad_c: np.ndarray, a: np.ndarray, b: np.ndarray):
    # grad_c: (m, n)
    grad_a = grad_c @ b.T  # (m, n) @ (n, k) -> (m, k)
    grad_b = a.T @ grad_c  # (k, m) @ (m, n) -> (k, n)
    return grad_a, grad_b

The following figure visualizes these two functions as a graph, showing how the forward inputs a and b are reused by the backward function.

Let's see an example with more than one op. For example, if matmul is followed by relu to produce d, we can chain the two ops together, as shown in the following figure. The forward pass runs from left to right, and the backward pass runs from right to left. The gradients of a and b may be used to further backpropagate if a and b are outputs of earlier ops.

In a full neural network, if we chain all the ops together, the forward pass runs from the inputs to the loss, and the backward pass runs from the gradient of the loss, which is a scalar of 1, to the gradients of all the weights. This overall process of computing the gradients is called backpropagation.