Fantastic work! Extremely neat and clear implementation! Interesting note on the backward pass - what do you think are the main blockers for a backward pass?
I need to get more comfortable with matrix derivatives before I can confidently reimplement it in the same minimal way as I did with the forward pass.