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

SERIES: Learning RL and JAX in Public - from zero to DeepMind :) On Day 1 I covered what JAX is and the three superpowers: GPU acceleration, grad, and jit. Today I actually used them. And the moment w...