Optimizing Flax NNX Models with Optax (Part 1)

From the creator

This is the first of our two part series on optimizing models with Optax. Let’s explore how to optimize Flax NNX neural network models using Optax, the primary optimization library in the JAX ecosystem. This presentation is designed especially for those of you familiar with PyTorch, and we'll be drawing parallels and highlighting differences to help you transition smoothly. This first episode will cover the core workflow, rom defining a model to running a complete, JIT-compiled training loop. It establishes the foundational knowledge that we’ll need for the more advanced topics in the next episode. Resources: Learn more → https://goo.gle/learning-jax Subscribe to Google for Developers → https://goo.gle/developers Speaker: Robert Crowe

Published 2025-12-04 · Watch on YouTube ↗
Availability and playback restrictions are controlled by the original channel.