On using the multi-dimensional chain-rule correctly.

I've recently seem some graduate students have confusion about partial derivatives and applying the chain rule. The following is a quiz I've written for my PhD advisor's Image and Video Processing students. I will use it to illustrate the proper application of the multi-dimensional chain-rule. The quiz question is as follows,

Below is a solution that I would expect students to arrive at. It makes use of the scalar chain rule and doesn't worry about deriving Jacobians, as the course does not emphasize this perspective heavily.

On using the multi-dimensional chain-rule

There are two distinct approaches to deriving this partial derivative: the Jacobian way or the scalar way. I see many students tripping themselves up because they're aware of these two methods but not fully aware on their distinction.

From preschool we are all familiar with the standard scalar chain rule of calculus. For differentiable $f: \mathbb{R} \rightarrow \mathbb{R}$ and $g: \mathbb{R} \rightarrow \mathbb{R}$,

It's this rule above that we directly employ in the above solution, by expanding the loss function in terms of its scalar variables, $\hat{\mathbf{y}}[n]$ and $\mathbf{y}[n]$.

However, students also learn that a similar chain rule exists for vector input/output mappings. Namely, for differentiable mappings $f: \mathbb{R}^K \rightarrow \mathbb{R}^M$ and $g: \mathbb{R}^N \rightarrow \mathbb{R}^K$,

where $\frac{\partial f}{\partial g} \in \mathbb{R}^{M \times K}$ is the Jacobian matrix of $f$, defined element-wise as

and an analagous definition for $\frac{\partial g}{\partial x} \in \mathbb{R}^{K \times N}$. As a sanity check, observe that the shapes of the matrix multiplication in (7) work out, and that the Jacobian of $f \circ g$ is an $M \times N$ matrix.

Trouble then arises when students derive the elements of the Jacobian matrices in (7) separately and then forget to compose them using matrix multiplication, i.e.,

Students who have trouble will often try to do element-wise multiplication of $\frac{\partial f}{\partial g}$ and $\frac{\partial g}{\partial x}$, even when the shapes don't make sense.