Major refactoring of mnist_nn.py, moved tests to separate file