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.
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
Flax NNX: Neural Networks That Feel Like PyTorch
