First Steps with Jax, Part 2
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 …