WKWesley Kambaleinkambale.dev·Mar 4 · 10 min readCustom training loop from scratch in JAXFor the past three weeks, we've been building up to this moment. Week 1 taught us that JAX is fast. Week 2 showed us how to eliminate loops with vmap and compute gradients with grad. Week 3 gave us Fl00
