In the second part of this series, we build a simple multilayer perceptron using some of the tools available in the vanilla jax library. Again, I've included a fully-functional notebook you can run on Google Colab and modify in any way you would like; a static copy is reproduced below …

Earlier this summer, I took some time to invest in getting more comfortable with the latest-and-greatest deep learning frameworks that had been released over the past two or so years. While my Keras-fu was serving me well with a lot of the toy deep learning projects I'd mess around …