Create a neural network using Jax and make predictions on the seeds dataset

In my last blog post I discussed how to make a neural network from scratch using numpy with the seeds dataset, and that post can be found here:- https://medium.com/@tracyrenee61/create-a-neural-network-in-numpy-and-make-predictions-on-seeds-a0e4955c6ce1

In this post I intend to discuss how to make a neural network from scratch, using the Jax library. Jax is a Python library that is set to overtake tensorflow (for research purposes), so I thought it would be a good idea to learn as much about Jax as I can. Jax has more powerful capability than numpy and that is one reason why it is a good idea to learn how to use this library and take advantage of its versatility. Jax is written very similar to the numpy Application Programming Interface (API), but there are a few areas where the two libraries do not match up. For instance, the arrays in Jax are immutable, so the programmer has to work around this to accomplish the same tasks in Jax that are easily performed in numpy. In addition, Jax’s random number generator has different arguments than numpy, so the coding for this must change as well.

I have created a Jupyter Notebook with the code in Python in Google Colab, which is a free online Jupyter Notebook hosted by Google. One main flaw in this product is the fact that Google Colab does not have…