--- title: "Deep Learning with JAX: Deep Learning With Jax" id: "12520" type: "page" slug: "01-deep-learning-with-jax" published_at: "2026-07-19T17:18:42+00:00" modified_at: "2026-07-20T00:41:40+00:00" url: "https://preppers-paradise.com/library/deeplearningwithjax/01-deep-learning-with-jax/" markdown_url: "https://preppers-paradise.com/library/deeplearningwithjax/01-deep-learning-with-jax.md" excerpt: "This book provides a practical guide to building and optimizing deep learning models using the JAX library. Readers will learn to leverage its core capabilities such as automatic differentiation, JIT" taxonomy_category: - "AI & Machine Learning" - "Books" - "Free Teaser" taxonomy_post_tag: - "automatic differentiation" - "deep learning" - "flax" - "functional programming" - "high-performance computing" - "jax" - "machine learning" - "neural networks" - "parallel computing" - "python" --- # Deep Learning with JAX: Deep Learning With Jax [← Deep Learning with JAX](/library/deeplearningwithjax/) Chapter 1 of 74 · Free teaser # Deep Learning with JAX Grigory Sapunov ![](https://preppers-paradise.com/wp-content/uploads/x402-books/deeplearningwithjax/_page_2_Picture_2.jpeg "Deep Learning with JAX book cover") *To my parents, who encouraged me to follow my passions, surrounded me with great books, and got me my first computer when it was still a luxury (and mostly an entertainment)* ## Preface ## About This Book ## Acknowledgments ## About the Author ## About the Cover Illustration ## Brief Contents | Part | Chapter | Title | |--------|---------|-----------------------------------------| | Part 1 | | First steps | | | 1 | When and why to use JAX | | | 2 | Your first program in JAX | | Part 2 | | Core JAX | | | 3 | Working with arrays | | | 4 | Calculating gradients | | | 5 | Compiling your code | | | 6 | Vectorizing your code | | | 7 | Parallelizing your computations | | | 8 | Using tensor sharding | | | 9 | Random numbers in JAX | | | 10 | Working with pytrees | | Part 3 | | Ecosystem | | | 11 | Higher-level neural network libraries | | | 12 | Other members of the JAX ecosystem | ## Contents ### PART 1 First Steps * **1 When and why to use JAX** * 1.1 Reasons to use JAX * Computational performance * Functional approach * 1.2 JAX ecosystem * 1.3 How is JAX different from NumPy? * JAX as NumPy * Composable transformations * **2 Your first program in JAX** * 2.1 An overview of a JAX deep learning project * 2.2 Loading and preparing the dataset * 2.3 A simple neural network in JAX * Neural network initialization * Neural network forward pass * 2.4 vmap: Auto-vectorizing calculations to work with batches * 2.5 Autodiff: How to calculate gradients without knowing about derivatives * Loss function * Obtaining gradients * Gradient update step * Training loop * 2.6 JIT: Compiling your code to make it faster * 2.7 Saving and deploying the model * 2.8 Pure functions and composable transformations: Why are they important? * Exercise 2.1 ### PART 2 Core JAX * **3 Working with arrays** * 3.1 Image processing with NumPy arrays * Loading an image into a NumPy array * Performing basic preprocessing operations with an image * Adding noise to the image * Implementing image filtering * Saving a tensor as an image file * 3.2 Arrays in JAX * Switching to JAX NumPy-like API * What is Array? * Device-related operations * Asynchronous dispatch * Running computations on TPU * 3.3 Differences from NumPy * Immutability * Types * 3.4 High-level and low-level interfaces: jax.numpy and jax.lax * Control flow primitives * Type promotion * Exercise 3.1 * **4 Calculating gradients** * 4.1 Different ways of getting derivatives * Manual differentiation * Symbolic differentiation * Numerical differentiation * Automatic differentiation ## 4.2 Calculating gradients with autodiff ### Working with gradients in TensorFlow ### Working with gradients in PyTorch ### Working with gradients in JAX ### Higher-order derivatives ### Multivariable case ## 4.3 Forward- and reverse-mode autodiff ### Evaluation trace ### Forward mode and jvp() ### Reverse mode and vjp() ### Going deeper ## Compiling your code ### 5.1 Using compilation #### Using JIT compilation #### Pure functions and compilation process ### 5.2 JIT internals #### Jaxpr, an intermediate representation for JAX programs #### XLA #### Using AOT compilation ### 5.3 JIT limitations #### Pure and impure functions #### Exact numerics #### Conditioning on input parameter values #### Slow compilation #### Class methods #### Simple functions ### Exercise 5.1 ## Vectorizing your code ### 6.1 Different ways to vectorize a function #### Naive approaches #### Manual vectorization #### Automatic vectorization #### Speed comparisons ### 6.2 Controlling vmap() behavior #### Controlling array axes to map over #### Controlling output array axes #### Using named arguments #### Using decorator style #### Using collective operations ### 6.3 Real-life use cases for vmap() #### Batch data processing #### Batching neural network models #### Per-sample gradients ### Vectorizing loops ## 7 Parallelizing your computations ### 7.1 Parallelizing computations with pmap() #### Setting up a problem #### Using pmap (almost) like vmap ### 7.2 Controlling pmap() behavior #### Controlling input and output mapping axes #### Using named axes and collectives ### 7.3 Data-parallel neural network training example #### Preparing data and neural network structure #### Implementing data-parallel training ### 7.4 Using multihost configurations ## 8 Using tensor sharding ### 8.1 Basics of tensor sharding #### Device mesh #### Positional sharding #### An example with 2D mesh #### Using replication #### Sharding constraints #### Named sharding #### Device placement policy and errors ### 8.2 MLP with tensor sharding #### Eight-way data parallelism #### Four-way data parallelism, two-way tensor parallelism ## 9 Random numbers in JAX ### 9.1 Generating random data #### Loading the dataset #### Generating random noise #### Performing a random augmentation ### 9.2 Differences with NumPy #### How NumPy works #### Seed and state in NumPy #### JAX PRNG #### Advanced JAX PRNG configuration ### 9.3 Generating random numbers in real-life applications #### Building a complete data augmentation pipeline #### Generating random initializations for a neural network ## 10 Working with pytrees ### 10.1 Representing complex data structures as pytrees ### 10.2 Functions for working with pytrees #### Using tree_map() #### Flatten/unflatten a pytree #### Using tree_reduce() #### Transposing a pytree ### 10.3 Creating custom pytree nodes # PART 3 Ecosystem ## 11 Higher-level neural network libraries ### 11.1 MNIST image classification using an MLP #### MLP in Flax #### Optax gradient transformations library #### Training a neural network the Flax way ### 11.2 Image classification using a ResNet #### Managing state in Flax #### Saving and loading a model using Orbax ### 11.3 Using the Hugging Face ecosystem #### Using a pretrained model from the Hugging Face Model Hub #### Going further with fine-tuning and pretraining #### Using the diffusers library ## 12 Members of the JAX ecosystem ### 12.1 Deep learning ecosystem #### High-level neural network libraries #### LLMs in JAX #### Utility libraries ### 12.2 Machine learning modules #### Reinforcement learning #### Other machine learning libraries ### 12.3 JAX modules for other fields ## Appendix A Installing JAX ## Appendix B Using Google Colab ## Appendix C Using Google Cloud TPUs ## Appendix D Experimental parallelization ## Index # Preface JAX is a powerful Python library created by Google for deep learning and highperformance computing. It's widely used in machine learning research and ranks as the third most popular deep learning framework, trailing only behind TensorFlow and PyTorch. Notably, it's the go-to framework for companies like DeepMind, and Google's research increasingly relies on JAX. What I really appreciate about JAX is its emphasis on functional programming in deep learning. It offers robust function transformations, including gradient computation, JIT compilation via XLA, auto-vectorization, and parallelization. JAX supports both GPUs and TPUs, delivering impressive performance. Now is an exciting time to dive into JAX, as its ecosystem is rapidly expanding. Despite being around for a few years, there's a noticeable lack of comprehensive resources for beginners. While JAX's website offers solid documentation and a supportive community, piecing everything together, especially when integrating other libraries, can be daunting. This book is crafted for those eager to master JAX. My goal is to consolidate crucial information in one place and guide you through understanding JAX concepts, enhancing your skills and ability to apply JAX in your projects and research. A basic understanding of deep learning and proficiency in Python are expected. This book doesn't cover deep learning basics, as there are plenty of resources available. Instead, it focuses solely on JAX, although I'll briefly touch on key deep learning concepts when necessary. This should benefit individuals from non–deep learning backgrounds, like physics. JAX is more than just a deep learning framework. Its expanding range of modules beyond deep learning suggests its potential in differentiable programming, large-scale physics simulations, and more. ## Preface My hope is that this book will also serve those interested in these applications. JAX continues to evolve, and I've had to update several chapters significantly. Don't worry about possible changes in the future; the core knowledge you'll gain remains applicable to future versions of JAX. ## Acknowledgments The book took longer than I expected. I changed a few countries along the way; JAX versions also changed. Some chapters had to be rewritten significantly. But now everything is done! First and foremost, I want to thank my family, my wife Mila, and my kids, Danya and Fedya. You suffered from a lack of my attention for so long! Yet you constantly supported me all along the way. I want to thank the people of Armenia, where we lived for some time, for their kindness and hospitality. Special thanks go to the Yerevan tech startup community for their help and support. To Hrant Khachatrian, Zaven Navoyan, Arsen Yeghiazaryan, Andranik Khachatryan, Ashot Arzumanyan, Ash Vardanian, Adam Bittlingmayer, Artur Aleksanyan, Erik Arakelyan, Karén Gyulbudaghyan, and many others—thank you a lot! Thank you, Enterprise Armenia, the National Investment Promotion Agency of Armenia. You do a great job, and your help was invaluable. And, finally, I thank my GDE (Google Developer Experts) friends and Google for supporting such a great initiative. The GDE community is fantastic! Many GDEs looked at early versions and provided helpful feedback. Special thanks go to David Cardozo for his exceptional feedback! ## About This Book Deep Learning with JAX is written to help you understand and start using JAX in your projects and research. It consolidates crucial information in one place and guides you through understanding JAX concepts through a series of easy-to-digest examples that build your intuition in the subject. ### Who Should Read This Book? Deep Learning with JAX targets deep learning practitioners and researchers familiar with frameworks such as PyTorch and TensorFlow who want to start using JAX. Readers should have a basic grasp of deep learning and be proficient in Python. Researchers from other areas (e.g., physics or optimization) or graduate students focused on deep learning, numerical optimization, or distributed computations will also find this book beneficial to their learning and practice. ### How This Book Is Organized: A Roadmap The book has three parts that cover 12 chapters. Part 1 is the introduction and a showcase for JAX: - Chapter 1 answers the crucial question, "Why JAX?" We'll explore what JAX is, its strengths and weaknesses compared to other frameworks like TensorFlow and PyTorch, and when it might be the best tool for your project. - Chapter 2 guides you through your first hands-on experience with JAX. We'll build a simple neural network for image classification, introducing key concepts like JAX transformations for auto-vectorization, gradient calculation, and just-in-time (JIT) compilation. You'll also learn how to save and load models and understand the difference between pure and impure functions in JAX. Part 2 covers the core functionalities of JAX: - Chapter 3 explores the workhorse of deep learning: tensors or multidimensional arrays. We'll compare NumPy arrays and JAX arrays, discuss how to work with them on different hardware like CPUs, GPUs, and TPUs, and explain the nuances of adapting code between NumPy and JAX. - Chapter 4 tackles the critical task of calculating gradients, which is essential for training neural networks. We'll compare various differentiation methods, dive deep into JAX's automatic differentiation capabilities, and explore both forward and reverse autodiff modes. - Chapter 5 teaches how to optimize your code for performance using JIT compilation. We'll examine how JIT works under the hood, its interaction with the XLA compiler, and how to address potential limitations. - Chapter 6 introduces automatic vectorization, a powerful technique for efficiently processing batches of data. We'll explore different vectorization approaches, discuss how to control JAX's vmap() transformation, and analyze real-life scenarios where auto-vectorization shines. - Chapter 7 delves into parallelization, enabling you to run computations on multiple devices simultaneously. We'll discuss how to use the pmap() transformation for parallel execution, control its behavior, and implement data parallel neural network training. We'll also explore running code on multihost configurations for large-scale tasks. - Chapter 8 introduces tensor sharding, a modern and efficient approach to parallelization in JAX. We'll explain how to leverage XLA for automatic parallelization and implement data and tensor parallelism for neural network training, as well as the advantages of this technique. - Chapter 9 tackles the important topic of generating random numbers in JAX. We'll explore the differences between JAX and NumPy in this regard, discuss the role of keys in representing the state of random number generators, and explain how to apply these concepts in real-world applications. - Chapter 10 introduces pytrees, a powerful tool for representing complex data structures in JAX. We'll discuss how to work with pytrees effectively, utilize functions for manipulating them, and even create custom pytree nodes for specialized needs. Part 3 covers the rich and diverse ecosystem of libraries and tools built around JAX: - Chapter 11 introduces higher-level neural network libraries like Flax and Optax, which provide convenient abstractions for building and training complex models. We'll use Flax to build both a simple MLP and a more advanced residual network for image classification and explore how to leverage Hugging Face libraries for working with transformers and diffusion models. - Chapter 12 takes a broader look at the JAX ecosystem, showcasing libraries for various machine learning tasks, including training large language models (LLMs), reinforcement learning, and evolutionary computations. We'll also explore JAX modules for other scientific fields like physics, chemistry, and more. If you're a manager, I suggest reading the first two chapters to get a grasp of JAX's strengths, how it stands out from PyTorch and TensorFlow, and what a typical JAX machine learning project looks like. Chapter 12 is also non-technical and can give you insight into where JAX shines. For developers eager to start crafting neural networks with JAX, focus on chapter 2 for a straightforward deep learning example, chapters 3 to 6 for JAX's foundational concepts, and chapter 11 for an overview of the high-level libraries in the ecosystem. You can tackle the rest of the book in any order, depending on your specific interests. Feel free to skip chapters 7 and 8 if parallelization isn't on your radar yet—you can always circle back. Dive into chapters 9 and 10 if you're curious about random numbers and pytrees, though the earlier chapters provide enough groundwork to get you started. ### About the Code This book contains many examples of source code both in numbered listings and in line with normal text. In both cases, source code is formatted in a fixed-width font like `this` to separate it from ordinary text. Sometimes code is also **in bold** to highlight code that has changed from previous steps in the chapter, such as when a new feature adds to an existing line of code. In many cases, the original source code has been reformatted; we've added line breaks and reworked indentation to accommodate the available page space in the book. In some cases, even this was not enough, and listings include line-continuation markers (➥). Additionally, comments in the source code have often been removed from the listings when the code is described in the text. Code annotations accompany many of the listings, highlighting important concepts. The most important source of information is the JAX documentation [https://jax.readthedocs.io/en/latest/](https://jax.readthedocs.io/en/latest/). It updates frequently, and you may find answers there. Other important sources of information are GitHub Discussions [https://github.com/google/jax/discussions](https://github.com/google/jax/discussions) and Issues [https://github.com/google/jax/issues](https://github.com/google/jax/issues). ## *about the author* ![Grigory Sapunov portrait](https://preppers-paradise.com/wp-content/uploads/x402-books/deeplearningwithjax/_page_19_Picture_1.jpeg) **Grigory Sapunov** is a cofounder and the CTO of Intento. He is a software engineer with more than 20 years of experience, holds a PhD in artificial intelligence, and is a Google Developer Expert in Machine Learning.