So I have a Ph.D. in analysis and that makes me not the best person to ask because I’ll either say that you need to do it yourself or get out pen and paper to do it with you (I also made a 5+ hour video course on autograd, which I won’t advertise here, though I need to update the session on forward differentiation with the progress that has been made, and you might prefer the free material out there).
If you wanted to do it yourself, you could write down the chain rule, for three functions, say (grossly oversimplifying)
f : Parameters → Intermediates, g: Intermediates → Outputs, h: Outputs → Loss.
where Parameters, Intermediates, Outputs are in some n, m, k-dimensional space respectively and the Loss is scalar.
The you spell out the chain rule (writing Df=dIntemediate/dParameter for the m x n matrix with the partial derivatives ∂Intermediate_i/∂Parameter_j) for various things, e.g.
dLoss/dParameter = dLoss/dLoss dLoss/dOutput dOutput/dIntermediate dIntermediate/dParameter
and then you start evaluating this from the left to right in the backward pass (dLoss/dLoss = 1). The intermediate results will be “row vectors” (grad_out and grad_in in autograd-lingo) and the factors are the Jacobians that are applied in the “backward”.
In recent PyTorch, you can call the .grad_fn nodes, so you can indeed compute this step by step for your viewing pleasure (but you have to sumtosize whenever broadcasting happens).
When you do forward AD, you start with a vector (the direction) on the right hand side of the chain rule.
The idea is not to get matrix-valued intermediates (as these cost a lot of memory and compute to process).
Best regards
Thomas