Optax

As explained in Deep Learning with JAX

A JAX-based library of gradient transformations and optimizer-building blocks used to define and combine training updates.

Optax is a JAX-based library for gradient transformations and optimizer-building blocks. It supplies ready-to-use optimizers and a composable framework for creating custom update rules from reusable pieces. section 11.1.2 section 12.1.3

The practical role of Optax here is to simplify training loops: create an optimizer, initialize its state, turn gradients into updates, and apply those updates to model parameters. In the surrounding workflow, it plugs into TrainState to hide most of that boilerplate. section 11.1.2 section 11.1.3

Optax is tied here to a small set of recurring ideas:

• gradient transformations and reusable optimizer pieces. section 11.1.2
• TrainState for storing optimizer state and applying gradients. section 11.1.3
• JAX training loops and functional model updates. section 11.4

Using an Optax optimizer in a training step

An example workflow creates an Adam or SGD optimizer, initializes state from model parameters, computes gradients with jax.grad or jax.value_and_grad, then calls optimizer.update() and applies the resulting parameter updates. A TrainState-based loop hides the state initialization and update application inside apply_gradients(). section 11.1.2 section 11.1.3

Figure 11.2 The process of using an optimizer in Optax (only a single gradient update step is shown)

Shows the one-step optimizer flow from gradient computation to parameter updates. section 11.1.2

How Deep Learning with JAX covers it

How deep it goes

Explains predefined optimizers, the GradientTransformation interface, and the init-update-apply workflow, but stops short of deeper optimizer internals.

Examples

Uses Adam, SGD with momentum, TrainState.create(), apply_gradients(), and a one-step optimizer diagram.

What it assumes

Assumes JAX gradients and a basic model training loop, while introducing optimizer state and update flow from the ground up.

What you can do after

Apply it via Optax or TrainState to run training updates with less boilerplate and reusable optimizer logic.

How it connects

Connects directly to JAX, Flax, TrainState, gradients, and training loops.

liveBookmanning.com

Read Deep Learning with JAX on liveBook Get the book at manning.com

Related concepts

• TrainStateused by · A Flax dataclass-like container that bundles model parameters, optimizer state, step count, and an apply function for training.
• Flaxused by · A high-level neural-network library built on top of JAX that provides a functional core and a module-based API.
• optimizationimplements · Optimization is the systematic selection or adjustment of values, actions, or system designs to improve a specified objective, often subject to constraints.
• adamimplements · Adam is an adaptive stochastic-gradient optimization algorithm that updates model parameters using moving estimates of the first and second moments of gradients.

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.