stalhabukhari/JAX-DNN-MNIST

A demonstration of NN training in JAX (via STAX).

Jupyter Notebook