Ура! Моя книга
“Deep Learning with JAX” (в девичестве "JAX in Action") вышла в печать! Я только что получил свои бумажные копии
🙂https://www.manning.com/books/deep-learning-with-jax Для тех, кто не следил, JAX -- это питоновская библиотека для высокопроизводительных вычислений и large-scale ML, с отличной поддержкой ускорителей, в частности TPU.
На данный момент JAX является вполне реальной альтернативой TensorFlow и PyTorch (
torch.func, в юности functorch, до сих пор пытается угнаться и всё ещё beta), и многие компании, в частности Google DeepMind, Cohere, xAI и прочие, перешли на него. На JAX созданы такие известные модели как AlphaFold, GraphCast, Gemini, Gemma, Grok, и я уже молчу сколько разного рисёча.
JAX -- это больше, чем библиотека для ML, это библиотека для очень разных высокопроизводительных, параллельных и распределённых вычислений. Не просто так его называют “NumPy на стероидах”. За пределами ML/DL, например, JAX активно используется для физических симуляций, и на GitHub есть уже огромное количество производных библиотек.
Сейчас отличное время, чтобы застолбить себе немного будущего :)
Отдельная радость должна быть для любителей функционального программирования, ибо JAX -- это первый фреймворк с большим охватом, работающий в этой парадигме. Очень прикольно использовать функции для трансформации других функций. Написали функцию для обработки одного элемента -- трансформировали в функцию для обработки батча. Написали сложную математическую функцию -- трансформировали в функцию, вычисляющую её производную. Аналогично с компиляцией и распараллеливанием. Никаких hidden state и side-effects, код чист, красив и понятен. А также БЫСТР! (см.
https://x.com/fchollet/status/1735420737744507374)
Книга состоит из трёх частей на 370+ страницах.
Part 1: First steps.
Верхнеуровневое введение в JAX для менеджеров и вообще всех, рассказывающее, где и почему стоит использовать JAX. Плюс отдельная глава для тех, кто любит видеть код, где показан полный цикл реализации простой нейросети с использованием большинства фишек JAX.
Part 2: Core JAX.
Основная часть книги, где покрыты все основы JAX, шаг за шагом. От работы с массивами (тензорами), autodiff, компиляция, векторизация, параллелизация и шардирование, случайные числа (в функциональном программировании старые приёмы из NumPy не работают эффективно, зато теперь всё наглядно и воспроизводимо!) и pytrees.
Part 3: Ecosystem.
Большая глава с практическим знакомством с экосистемой высокоуровневых библиотек для DL (Flax, Optax, Orbax, CLU, …), а также примеры использования HuggingFace Transformers/Diffusers, которые давно уже добавили поддержку JAX. Также есть отдельная глава с очень верхнеуровневым и широким обзором того, что есть в JAX и вокруг за пределами нейросетевого мейнстрима.
Много крутых и умных людей читало и ревьюило мою книгу, спасибо куче GDE и не только. И отдельное спасибо Франсуа Шолле за добрые слова
🙂“A comprehensive guide to mastering JAX, whether you’re a seasoned deep learning practitioner or just venturing into the realm of differentiable programming and large-scale numerical simulations.”-- François Chollet, Software Engineer, Google
В общем это был прикольный опыт, я доволен результатом, надеюсь, вам тоже понравится.
Ещё отдельное спасибо всем, кто поддерживал GonzoML на Патреоне (
https://www.patreon.com/GonzoML). Всем действующим платным членам нашей тесной группы я отправил коды для получения книги бесплатно (проверьте сообщения!) -- у вас будет постоянно обновляемая версия (a JAX очевидно будет меняться!) в онлайн доступе.