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
Jul 1, 202611 min read148
