Predict on the breast cancer dataset using a neural network made with Jax

In my last blog post I discussed how to create a neural network from scratch in numpy and predict on the breast cancer dataset, with the link to the post being here:- https://medium.com/@tracyrenee61/predict-on-a-breast-cancer-diagnosis-using-a-neural-network-made-from-numpy-ab2aef30e93e

In this post I intend to discuss how to create a neural network made from scratch in Jax and predict on sklearn’s breast cancer dataset. The breast cancer is different from the other datasets that I have been working with because it is a binary classification set while the others are multiclass labelled sets. It is because the breast cancer dataset is different from the previous ones that I decided to make this post.

I have written the code used in this post in Google Colab, which is a free online Jupyter Notebook hosted by Google. Google Colab is a great platform to write code in Python in, with the exception that it does not have an undo function. Therefore, care needs to be taken not to accidentally delete or overwrite valuable code. If that happens, it may be possible to retrieve the code by going into the script’s save history.

After I created the Jupyter Notebook, I imported the libraries that I would need to execute the program, being:-