Is this calculation of the vector-Jacobian product in the PyTorch documention wrong?

Question

In the official PyTorch documentation there is the following calculation (here):

$$ J^{T} \cdot \vec{v}=\left(\begin{array}{ccc} \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{1}} \\ \vdots & \ddots & \vdots \\ \frac{\partial y_{1}}{\partial x_{n}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}} \end{array}\right)\left(\begin{array}{c} \frac{\partial l}{\partial y_{1}} \\ \vdots \\ \frac{\partial l}{\partial y_{m}} \end{array}\right)=\left(\begin{array}{c} \frac{\partial l}{\partial x_{1}} \\ \vdots \\ \frac{\partial l}{\partial x_{n}} \end{array}\right) $$

However, I am wondering why the result isn't as follows:

$$ J^{T} \cdot \vec{v}=\left(\begin{array}{ccc} \frac{\partial y_{1}}{\partial x_{1}} & \cdots & \frac{\partial y_{m}}{\partial x_{1}} \\ \vdots & \ddots & \vdots \\ \frac{\partial y_{1}}{\partial x_{n}} & \cdots & \frac{\partial y_{m}}{\partial x_{n}} \end{array}\right)\left(\begin{array}{c} \frac{\partial l}{\partial y_{1}} \\ \vdots \\ \frac{\partial l}{\partial y_{m}} \end{array}\right)=\left(\begin{array}{c} m\frac{\partial l}{\partial x_{1}} \\ \vdots \\ m\frac{\partial l}{\partial x_{n}} \end{array}\right) $$

As the matrix is multiplied by a vector and the $\partial y_{x}$ terms cancel out it is should be $m$ times $\frac{\partial l}{\partial x_{x}}$.

I know that the official PyTorch docs are probably right but I just can't get my head around why.

Answer 1

I think your error is assuming that $$\frac{\partial y_1}{\partial x_1} \frac{\partial l}{\partial y_1} = \frac{\partial l}{\partial x_1}.$$ While 'cancelling' derivatives like this works sometimes, it is not true in general, and treating these quantities as fractions is not really correct in the usual formulation of calculus in terms of limits.

If you run through the calculation using the chain rule with Jacobian matrices, you should recover the correct product. If you're not clear on the compositions, write the functions $y_i(x_1, \dots, x_n)$ and $l(y_1, \dots, y_m)$. If it helps define $y(x_1, \dots, x_n) = (y_1, \dots, y_m)$ then use the chain rule to compute $l(y(x_1, \dots, x_n))$. If it would help to see the calculation in more detail, do let me know. Otherwise, it's a good exercise.