I Built a Neural Network Training Loop in 5 Lines Using JAX - Day 2

Short summary

Day 2 of a public learning series on JAX and reinforcement learning, covering derivatives, gradients, and a from-scratch gradient descent training loop in five lines. The author verifies jax.grad against a hand-calculated derivative, builds a simple loss function, and runs five steps of gradient descent to converge a weight toward 5.0. A jit benchmark shows 32x speedup on CPU, with notes on recompilation when input shapes change.

• •jax.grad verified against hand-calculated derivative of x² at x=3, confirming output of 6.0
• •Five-line gradient descent loop converges weight from 1.0 to 5.01 over five steps with learning rate 0.1
• •jit compilation delivers 32x CPU speedup; input shape changes trigger recompilation

Generated with AI, which can make mistakes.