book

Deep Learning with JAX

Name: Deep Learning with JAX
Author: Grigory Sapunov
ISBN: 9781633438880

by Grigory Sapunov

November 2024

Intermediate to advanced

408 pages

12h 7m

English

Manning Publications

Read now

Unlock full access

Deep Learning with JAX
Copyright
dedication
contents
front matter
prefaceacknowledgmentsabout this bookWho should read this book?How this book is organized: A roadmapAbout the codeliveBook discussion forumOther online resourcesabout the authorabout the cover illustration
Part 1. First steps
1 When and why to use JAX
1.1 Reasons to use JAX1.1.1 Computational performance1.1.2 Functional approach1.1.3 JAX ecosystem1.2 How is JAX different from NumPy?1.2.1 JAX as NumPy1.2.2 Composable transformations1.3 How is JAX different from TensorFlow and PyTorch?Summary
2 Your first program in JAX
2.1 A toy ML problem: Classifying handwritten digits2.2 An overview of a JAX deep learning project2.3 Loading and preparing the dataset2.4 A simple neural network in JAX2.4.1 Neural network initialization2.4.2 Neural network forward pass2.5 vmap: Auto-vectorizing calculations to work with batches2.6 Autodiff: How to calculate gradients without knowing about derivatives2.6.1 Loss function2.6.2 Obtaining gradients2.6.3 Gradient update step2.6.4 Training loop2.7 JIT: Compiling your code to make it faster2.8 Saving and deploying the model2.9 Pure functions and composable transformations: Why are they important?Exercise 2.1Summary
Part 2. Core JAX
3 Working with arrays
3.1 Image processing with NumPy arrays3.1.1 Loading an image into a NumPy array3.1.2 Performing basic preprocessing operations with an image3.1.3 Adding noise to the image3.1.4 Implementing image filtering3.1.5 Saving a tensor as an image file3.2 Arrays in JAX3.2.1 Switching to JAX NumPy-like API3.2.2 What is Array?3.2.3 Device-related operations3.2.4 Asynchronous dispatch3.2.5 Running computations on TPU3.3 Differences from NumPy3.3.1 Immutability3.3.2 Types3.4 High-level and low-level interfaces: jax.numpy and jax.lax3.4.1 Control flow primitives3.4.2 Type promotionExercise 3.1Summary

4 Calculating gradients
4.1 Different ways of getting derivatives4.1.1 Manual differentiation4.1.2 Symbolic differentiation4.1.3 Numerical differentiation4.1.4 Automatic differentiation4.2 Calculating gradients with autodiff4.2.1 Working with gradients in TensorFlow4.2.2 Working with gradients in PyTorch4.2.3 Working with gradients in JAX4.2.4 Higher-order derivatives4.2.5 Multivariable case4.3 Forward- and reverse-mode autodiff4.3.1 Evaluation trace4.3.2 Forward mode and jvp()4.3.3 Reverse mode and vjp()4.3.4 Going deeperSummary
5 Compiling your code
5.1 Using compilation5.1.1 Using JIT compilation5.1.2 Pure functions and compilation process5.2 JIT internals5.2.1 Jaxpr, an intermediate representation for JAX programs5.2.2 XLA5.2.3 Using AOT compilation5.3 JIT limitations5.3.1 Pure and impure functions5.3.2 Exact numerics5.3.3 Conditioning on input parameter values5.3.4 Slow compilation5.3.5 Class methods5.3.6 Simple functionsExercise 5.1Summary
6 Vectorizing your code
6.1 Different ways to vectorize a function6.1.1 Naive approaches6.1.2 Manual vectorization6.1.3 Automatic vectorization6.1.4 Speed comparisons6.2 Controlling vmap() behavior6.2.1 Controlling array axes to map over6.2.2 Controlling output array axes6.2.3 Using named arguments6.2.4 Using decorator style6.2.5 Using collective operations6.3 Real-life use cases for vmap()6.3.1 Batch data processing6.3.2 Batching neural network models6.3.3 Per-sample gradients6.3.4 Vectorizing loopsSummary
7 Parallelizing your computations
7.1 Parallelizing computations with pmap()7.1.1 Setting up a problem7.1.2 Using pmap (almost) like vmap7.2 Controlling pmap() behavior7.2.1 Controlling input and output mapping axes7.2.2 Using named axes and collectives7.3 Data-parallel neural network training example7.3.1 Preparing data and neural network structure7.3.2 Implementing data-parallel training7.4 Using multihost configurationsSummary
8 Using tensor sharding
8.1 Basics of tensor sharding8.1.1 Device mesh8.1.2 Positional sharding8.1.3 An example with 2D mesh8.1.4 Using replication8.1.5 Sharding constraints8.1.6 Named sharding8.1.7 Device placement policy and errors8.2 MLP with tensor sharding8.2.1 Eight-way data parallelism8.2.2 Four-way data parallelism, two-way tensor parallelismSummary
9 Random numbers in JAX
9.1 Generating random data9.1.1 Loading the dataset9.1.2 Generating random noise9.1.3 Performing a random augmentation9.2 Differences with NumPy9.2.1 How NumPy works9.2.2 Seed and state in NumPy9.2.3 JAX PRNG9.2.4 Advanced JAX PRNG configuration9.3 Generating random numbers in real-life applications9.3.1 Building a complete data augmentation pipeline9.3.2 Generating random initializations for a neural networkSummary
10 Working with pytrees
10.1 Representing complex data structures as pytrees10.2 Functions for working with pytrees10.2.1 Using tree_map()10.2.2 Flatten/unflatten a pytree10.2.3 Using tree_reduce()10.2.4 Transposing a pytree10.3 Creating custom pytree nodesSummary
Part 3. Ecosystem
11 Higher-level neural network libraries
11.1 MNIST image classification using an MLP11.1.1 MLP in Flax11.1.2 Optax gradient transformations library11.1.3 Training a neural network the Flax way11.2 Image classification using a ResNet11.2.1 Managing state in Flax11.2.2 Saving and loading a model using Orbax11.3 Using the Hugging Face ecosystem11.3.1 Using a pretrained model from the Hugging Face Model Hub11.3.2 Going further with fine-tuning and pretraining11.3.3 Using the diffusers librarySummary
12 Other members of the JAX ecosystem
12.1 Deep learning ecosystem12.1.1 High-level neural network libraries12.1.2 LLMs in JAX12.1.3 Utility libraries12.2 Machine learning modules12.2.1 Reinforcement learning12.2.2 Other machine learning libraries12.3 JAX modules for other fieldsSummary
Appendix A. Installing JAX
A.1 Installing JAX on CPUsA.2 Installing JAX on GPUsA.2.1 pip installation with CUDAA.2.2 pip installation with self-installed CUDA/CuDNNA.2.3 Using Docker containersA.3 Installing JAX on TPUs
Appendix B. Using Google Colab
Appendix C. Using Google Cloud TPUs
C.1 Setting up a Cloud TPU projectC.2 Running a Cloud TPU instance and connecting it to Google ColabC.2.1 Running and deleting a Cloud TPU nodeC.2.2 Preparing a Cloud TPU nodeC.2.3 Connecting to a Cloud TPU node from a Colab notebookC.3 Resources
Appendix D. Experimental parallelization
D.1 Using xmap() and named-axis programmingD.1.1 Working with named axesD.1.2 Parallelism and hardware meshesD.2 Using pjit() for tensor parallelismD.2.1 Basics of pjit()D.2.2 MLP example with pjit()Summary
index

Content preview from Deep Learning with JAX

front matter

preface

JAX is a powerful Python library created by Google for deep learning and high-performance 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 ...

Become an O’Reilly member and get unlimited access to this title plus top books and audiobooks from O’Reilly and nearly 200 top publishers, thousands of courses curated by job role, 150+ live events each month,
and much more.

Read now

Unlock full access

More than 5,000 organizations count on O’Reilly

O’Reilly covers everything we've got, with content to help us build a world-class technology community, upgrade the capabilities and competencies of our teams, and improve overall team performance as well as their engagement.

Julian F.

Head of Cybersecurity

I wanted to learn C and C++, but it didn't click for me until I picked up an O'Reilly book. When I went on the O’Reilly platform, I was astonished to find all the books there, plus live events and sandboxes so you could play around with the technology.

Addison B.

Field Engineer

I’ve been on the O’Reilly platform for more than eight years. I use a couple of learning platforms, but I'm on O'Reilly more than anybody else. When you're there, you start learning. I'm never disappointed.

Amir M.

Data Platform Tech Lead

I'm always learning. So when I got on to O'Reilly, I was like a kid in a candy store. There are playlists. There are answers. There's on-demand training. It's worth its weight in gold, in terms of what it allows me to do.

Mark W.

Embedded Software Engineer

Publisher Resources

ISBN: 9781633438880Publisher Support Publisher Website

Cloud Computing

Data Engineering

Data Science

AI & ML

Programming Languages

Software Architecture

IT/Ops

Security

Design

Business

Soft Skills

Deep Learning with JAX

by Grigory Sapunov

front matter

preface

Become an O’Reilly member and get unlimited access to this title plus top books and audiobooks from O’Reilly and nearly 200 top publishers, thousands of courses curated by job role, 150+ live events each month,
and much more.