Optimizing Flax NNX Models with Optax (Part 1)
35
Google Developers2.57 млн
Следующее
3 дня – 1357:13Популярные
4 дня – 6 44013:26
24 дня – 3 1455:20Опубликовано 4 декабря 2025, 5:00
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 → goo.gle/learning-jax
Subscribe to Google for Developers → goo.gle/developers
Speaker: Robert Crowe
Resources:
Learn more → goo.gle/learning-jax
Subscribe to Google for Developers → goo.gle/developers
Speaker: Robert Crowe
Свежие видео
4 дня – 5 2270:20
5 дней – 29 9130:35
8 дней – 299 65811:45
10 дней – 71 8410:28
11 дней – 7802:20
50 дней – 36 1978:49Случайные видео
82 дня – 4 2580:38
154 дня – 2801:02
279 дней – 248 4119:46
11.01.22 – 26 63545:48
04.03.08 – 11 6325:18
10 дней – 10 1560:33
2 дня – 9 9910:32
3 дня – 2 05110:31
3 дня – 15412:02
4 дня – 3003:10
4 дня – 2510:52
5 дней – 1 0660:49
5 дней – 1302:21
2 дня – 84 16719:10