Optax: Optimizers You Can Compose Like LEGO

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