Intro to Jax Table of Contents Imports Introduction Backend: XLA What is jax? Jax primitives: jit , grad , vmap jit grad vmap Composition Building a neural network from scratch Creating the neural network Training Imports Some of the libraries we will use throughout this post are imported below. import time import numpy as np import matplotlib.pyplot as plt import matplotlib as mpl import seaborn…