This is a follow up question I have on this excellent answer: https://stats.stackexchange.com/a/235758/307400. I will save me writing down any details about reverse-mode automatic differentiation, the given answer gives a nice introduction. I also found this resource very helpful: https://marksaroufim.medium.com/automatic-differentiation-step-by-step-24240f97a6e6

I am implementing Automatic Differentiation of functions on matrices: $f: \mathbb{R}^{k\times l} \rightarrow \mathbb{R}^{n\times m}$. As mentioned in the answer linked above, if I calculate the partial derivative of such a function $f: \mathbb{R}^{k\times l} \rightarrow \mathbb{R}^{n\times m}$ with respect to another function $g: \mathbb{R}^{r\times s} \rightarrow \mathbb{R}^{p\times q}$, as we do in reverse-mode differentiation, I get a 4D array. The array has the shape $\frac{\partial f}{\partial g}: n\times m \times p \times q$. So far so good.

The formula for computing the adjoint $\overline{v_i}$ given the adjoints of $v_i$'s parents in the computation graph is $$ \overline{v_i} = \sum_{j\in parents(i)} \overline{v_j} \frac{\partial v_j}{\partial v_i} $$ The shapes are:

• $\overline{v_p}: p\times q \times n \times m$
• $\frac{\partial v_p}{\partial v_i}: n\times m \times k \times l$

So here we have to perform "matrix multiplication" of two 4D arrays. How does this work?

I have come up with an approach to multiply these 4D arrays in the given context that makes some intuitive sense to me, but I have no idea if it is theoretically sound:
First thing I noticed is that we want to multiply over both dimensions $n$ and $m$: We multiply an array of shape $p\times q \times n \times m$ with another one of shape $n\times m \times k \times l$ and want to obtain an array of shape $p\times q \times k \times l$, as we want to obtain another adjoint, that is a partial derivative of the final "output" function that I left implicit here but has range $\mathbb{R}^{p\times q}$ with respect to $v_i$ with range $\mathbb{R}^{k\times l}$. We want to "multiply out" the dimensions $n\times m$.

What I did then was to consider each two dimensions together as one, implying that matrices are collapsed into vectors, e.g. by iterating them in row major order. This simplifies the shapes in the multiplication to $$ \underbrace{p\times q}_{a} \times \underbrace{n\times m}_{b} \;\;\cdot\;\; \underbrace{n\times m}_{b} \times \underbrace{k\times l}_{c} $$ Now I just did ordinary matrix multiplication with these new indices (using $x\in \{1, ..., a\}$; $y \in \{1, ..., b\}$ and $z \in \{1, ..., c\}$): $$ \left(\overline{v_p} \frac{\partial v_p}{\partial v_i} \right)_{x,z} = \sum_{y=1}^{b} \bigg(\;\overline{v_p}\;\bigg)_{x,y} \bigg(\frac{\partial v_p}{\partial v_i}\bigg)_{y,z} $$

This sums over the two dimensions we want to get rid of, but is this also theoretically sound?