Home Blog Learn

All Courses

Getting Started with JAX

Chapter 1: Introduction to JAX

Core Design Philosophy: Function Transformations

Installation and Setup

Working with JAX Arrays

Device Management: CPU, GPU, TPU

Practice: Basic Array Operations

Chapter 2: Accelerating Functions with JIT Compilation

The Need for Speed: Why Compile?

Introducing jax.jit

How JIT Works: Tracing and Compilation

Python Control Flow and jit

Static vs Traced Values

Common Pitfalls with jit

Hands-on Practical: Applying jit

Chapter 3: Automatic Differentiation with grad

Understanding Gradients

Introducing jax.grad

How Autodiff Works: Reverse Mode

Differentiating with Respect to Arguments

Higher-Order Derivatives (grad of grad)

Value and Gradient (jax.value_and_grad)

Differentiation and Control Flow

Limitations and Considerations

Hands-on Practical: Computing Gradients

Chapter 4: Automatic Vectorization with vmap

The Concept of Vectorization

Introducing jax.vmap

Mapping over Specific Arguments (in_axes, out_axes)

Handling Multiple Batched Arguments

Combining vmap with jit and grad

Performance Considerations for vmap

Hands-on Practical: Vectorizing Functions

Chapter 5: Parallelization Across Devices with pmap

Introduction to Data Parallelism (SPMD)

Introducing jax.pmap

Mapping Data to Devices (in_axes, out_axes)

Device Meshes and Axis Names

Collective Operations (lax.psum, lax.pmean, etc.)

Combining pmap with other Transformations

Debugging pmapped Functions

Hands-on Practical: Parallel Computation

Chapter 6: Managing State in JAX

Functional Purity and Side Effects

The Challenge of State in Functional Code

Pattern: Explicit State Passing

Using PyTrees for Structured State

Example: Stateful Counter

Example: Simple Optimizer State

Combining State Management with Transformations

Practice: Implementing Stateful Functions

Automatic Differentiation with `grad`

Many computational tasks, especially in machine learning model training, rely on calculating the gradient of functions. Automatic differentiation provides an efficient way to compute these derivatives. This chapter introduces jax.grad, the core JAX transformation for obtaining gradient functions from Python code operating on numerical inputs.

You will learn how to:

Apply jax.grad to compute the gradient $∇f(x)$ of a scalar-valued function $f$ .
Understand the basics of reverse-mode automatic differentiation as used by grad.
Control differentiation with respect to specific function arguments.
Compute higher-order derivatives by composing grad.
Efficiently get both the function's output value and its gradient using jax.value_and_grad.
Recognize how control flow interacts with differentiation and identify potential limitations.

By the end of this chapter, you will be able to use jax.grad effectively to differentiate your numerical functions within the JAX framework.

Sections

3.1 Understanding Gradients
3.2 Introducing `jax.grad`
3.3 How Autodiff Works: Reverse Mode
3.4 Differentiating with Respect to Arguments
3.5 Higher-Order Derivatives (`grad` of `grad`)
3.6 Value and Gradient (`jax.value_and_grad`)
3.7 Differentiation and Control Flow
3.8 Limitations and Considerations
3.9 Hands-on Practical: Computing Gradients

© 2025 ApX Machine Learning

;