Here is the setup that we work with - a function is shown as a (input, output) edge. The Jacobian of this function is written on the edge itself:$$x\in R^n\underset{\text{Shape (m,n)}}{\overset{\partial y/\partial x}{\to }}y\in R^m$$This representation is naturally extendable to a composition of functions $z=g(y),y=f(x)$. Next we see how chain rule looks like on this graph.

Consider two cases:

In the first case, we have a simple composition of functions: $z=g(y),y=f(x)$

Next we look at a DAG which has two inputs and one output. In this case, the Jacobian wrt input is the sum of Jacobians of each path: $\frac{dz}{dx}=(\underset{J_3}{\underbrace{\frac{\partial z}{\partial y}}})(\underset{J_2}{\underbrace{\frac{\partial y}{\partial x}}})+\underset{J_1}{\underbrace{\frac{\partial z}{\partial x}}}$. The first path is $J_2\to J_3$ and the second path is $J_1$.

This is all we need to fully define the chain rule on a graph. In the next example we combine both these cases.

Consider a DAG of a function $w=f(x),w\in R,x\in R^n$.

This leads to a insight about autodiff on a graph:

A path product is the product of all Jacobians along a path.

Let's revisit the DAG of the variance function:

There are two paths from $x$ to ${\sigma }^2$. We calculate the Jacobians along each path and add them up to get the Jacobian of ${\sigma }^2$ wrt $x$: $$\begin{array}{rl}{J}_5& =1/2\\ J_4& =2y^T\\ J_3& =\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]\\ J_2& =\left[\begin{array}{c}-1\\ -1\end{array}\right]\\ J_1& =\left[\begin{array}{cc}1/2& 1/2\end{array}\right]\\ \\ \partial {\sigma }^2/\partial x& =J_5{J}_4{J}_2{J}_1+J_5{J}_4{J}_3\in R^{1\times 2}\\ & =y^T\left[\begin{array}{cc}1/2& 1/2\end{array}\right]\left[\begin{array}{c}-1\\ -1\end{array}\right]+y^T\left[\begin{array}{cc}1& 0\\ 0& 1\end{array}\right]\\ & =\frac{1}{2}{x}^T\left[\begin{array}{cc}1& -1\\ -1& 1\end{array}\right]\end{array}$$ This is a linear-algebriac way of representing:

Consider multiplication of three matrices $J_1,J_2,J_3$ of shapes $(m,n),(p,m),(q,p)$. There are two ways to multiply them:

$$\underset{qpm+qmn}{\underbrace{(\underset{(q,p)}{J_3}\underset{(p,m)}{J_2})\underset{(m,n)}{J_1}}}$$$$\underset{qpn+pmn}{\underbrace{\underset{(q,p)}{J_3}(\underset{(p,m)}{J_2}\underset{(m,n)}{J_1})}}$$One can think of them as Jacobians corresponding to functions $a\stackrel{J_1}{\to }b,b\stackrel{J_2}{\to }c,c\stackrel{J_3}{\to }d$. Their product is then the Jacobian corresponding to $a\in R^n\to d\in R^q$.

In most cases in scientific computing, the input dimension $n$ is large and the output dimension $q$ is small (mostly 1). In this case, the first way (multiplying from output towards input) is more efficient. In cases where the output dimension is large or comparable to the input dimension (eg when calculating Hessian), the second way (multiplying from input towards output) is more efficient.

This is something that is often skipped or just skimmed over in a lot of text on autodiff. While we assumed the existence of Jacobians to explain the ideas of differentiation on a graph, in practice it is neither cheap not necessary to actually calculate and store the entire Jacobian matrix for each edge. Let's consider the following function to illustrate the issue:$$\underset{(n,n)}{\underbrace{x}}\stackrel{J_1}{\to }\underset{(=x^{-1})}{\underbrace{y}}\stackrel{J_2}{\to }\underset{(=\Vert y{\Vert }_F)}{\underbrace{z}}$$Here the Jacobian $J_1$ is of shape $(n^2,n^2)$ and Jacobian $J_2$ is of shape $(1,n^2)$. Let's try to calculate $\partial z/\partial x$ just by calculating and then multiplying the Jacobians. This is not efficient but we are doing it just to see why it is not efficient.

We know that:$$\begin{array}{rl}dy& =-x^{-1}(dx)x^{-1}\\ \Rightarrow \text{vec}(dy)& =\underset{J_1}{\underbrace{({(x^{-1})}^T\otimes x^{-1})}}\text{vec}(dx)\end{array}$$Similarly:$$\begin{array}{rl}\mathrm{\nabla }{z}_y& =\frac{y}{z}\\ \Rightarrow J_2& =\frac{1}{z}(\text{vec}(y){)}^T\end{array}$$Now that we have both the Jacobian matrices, we can calculate the product of them to get the Jacobian of $z$ wrt $x$. The product will be of shape $(1,n^2)$. It will need to be reshaped to a $(n,n)$ matrix. In the next part, we will implement it in a more efficient way.

import torch

# perform backprop to calculate `x.grad`
x = torch.randn(100,100, requires_grad=True, dtype=torch.float64)
y = torch.linalg.inv(x)
z = torch.linalg.norm(y)
z.backward()

# calculate the Jacobians 
x_inv = torch.linalg.inv(x).contiguous()
x_inv_t = torch.linalg.inv(x).t()
j1 = torch.kron(-x_inv_t, x_inv)

# need to transpose `y` first
# since pytorch uses row-major order
j2 = (y.t()/z).reshape(1,-1)

# calculate the jacobian and compare with actual gradient
j3 = (j2 @ j1).reshape(100,100).t()
assert torch.allclose(j3, x.grad)