Up and Running with JAX - Fully-Connected Network Forward Pass

In a previous post , I introduced JAX with particular emphasis on JIT compilation, vectorizing transformations and automatic differentiation. In this post, we walkthrough an implementation of the forward pass for a fully-connected neural network with the goal of classifying MNIST handwritten digits, incorporating concepts from the first post. We begin by loading MNIST training and validation sets,…