I have a product of two matrices, defined as $Y=AB$. I need to find $\nabla_A J$ where $J(Y)$ is a scalar-valued loss function dependent on $Y$. As I understand, this gradient can be written as $\nabla_A J = \nabla_A Y \cdot B^T$. I am having some difficulty proving this to myself.
Any element in $Y$ can be written down as $y_{ij} = \sum_{k=1}^M a_{ik}b_{kj}$ and it's derivative can be written as $\frac{\partial y_{ij}}{\partial a_{ik}} = b_{kj}$.
Some intense googling led me to this page which uses concrete matrices to demonstrate. They use the chain rule element by element like so:
\begin{equation} \label{eq:dLdX11-chain} \frac{\partial L}{x _ {1,1}} = \sum _ {i=1}^N\sum _ {j=1}^M \frac{\partial L}{\partial y _ {i,j}}\frac{\partial y _ {i,j}}{\partial x _ {1,1}} = \frac{\partial L}{\partial Y} \cdot \frac{\partial Y}{x _ {1,1}} \end{equation}
However, I do not know how to extend this generally to an entire matrix.
$\def\p#1#2{\frac{\partial #1}{\partial #2}}$There is a problem with your first derivative $$\eqalign{ y_{ij} &= a_{ik}b_{kj} \\ \p{y_{ij}}{a_{mn}} &= \left(\p{a_{ik}}{a_{mn}}\right)b_{kj} \\ &= \delta_{im}\delta_{kn}\;b_{kj} \\ &= \delta_{im}\,b_{nj} \\ }$$ If you switch the indices to $\Big((m,n)=(i,k)\Big)$, then you obtain $$\eqalign{ \p{y_{ij}}{a_{ik}} &= \delta_{ii}\,b_{kj} = N\,b_{kj} \\ }$$ The first rule of index notation is that repeated indexes indicate summation. Whenever you violate this rule you obtain invalid results. You did not intend to sum over the $i$ index on the LHS, but that is what happened.
Now let's to the whole calculation in index notation. $$\eqalign{ \p{L}{a_{mn}} &= \left(\p{L}{y_{ij}}\right) \left(\p{y_{ij}}{a_{mn}}\right) \\ &= \left(\p{L}{y_{ij}}\right) \left( \delta_{im}\,b_{nj} \right) \\ &= \left(\p{L}{y_{mj}}\right) b_{nj} \\ &= \left(\p{L}{y_{mj}}\right) b^T_{jn} \\ \p{L}{A} &= \left(\p{L}{Y}\right) B^T \\ }$$
In the above, $\delta_{im}$ denotes a Kronecker delta symbol. It has a special property under summation. It substitutes its indices onto adjacent terms in the expression.