Search for a command to run...
Series
These articles are about JAX, a library for array-oriented numerical computation, with automatic differentiation and JIT compilation to enable high-performance machine learning research.
Picture this. You have spent six hours training a model. The loss curve looks beautiful, accuracy is climbing, and you are one epoch away from a result worth writing home about. Then the power goes ou
Over the last seven weeks we have built models, written training loops, composed optimizers, bulletproofed our code with Chex, and learned to debug inside JIT-compiled functions. Every one of those ar
Last week, Chex gave us a firewall against shape mismatches and NaN gradients. That firewall catches the bugs that announce themselves; a tensor with the wrong shape, a gradient that explodes into inf
Here's a scenario every ML engineer eventually lives through: you kick off a training run before leaving the office. Three hours later, you check back in. The loss curve looks like a flat line at NaN.
In our last article on this blog in March, we built a complete training loop. We used optax.adamw() to create an optimizer, wrapped it in nnx.Optimizer, and watched our model learn. But we barely scra
Writing a training loop that you can actually control
