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 …

more ...