In order to make this repository reusable, it would be good to produce end-to-end documentation specifying requirements from data generation, to model training and inference.
I'm hoping to separate out the data loading/mini-batch preparation and the neural network class, but in order to make sure I understand what's going on I have left some TODOs with requests to update the docstrings and types etc in neural_networks.py on the branch pylint-py.