Deep Learning with JAX: Deep Learning With Jax

← Deep Learning with JAX

Chapter 1 of 74 ·
Free teaser

Deep Learning with JAX

Grigory Sapunov

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

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/. It updates frequently, and you may find answers there. Other important sources of information are GitHub Discussions https://github.com/google/jax/discussions and Issues https://github.com/google/jax/issues.