Nice! A small piece of feedback: I would have the dimensions mentioned in the text also annotated on the diagram. It wasn't exactly clear how the input data was flattened for example.
Entirely plausible this is intended for someone more "mathmatical" than myself but appreciate the work regardless.