optax

What it is and what it does

Optax is a gradient processing and optimization library built on top of JAX. It provides well-tested, efficient implementations of core optimization components—such as Adam, SGD, and other popular optimizers—along with loss functions and gradient transformation utilities. The library is designed around composability: rather than monolithic optimizer classes, it offers small building blocks that can be recombined in custom ways to create new optimizers or gradient processing pipelines.

The package targets researchers and practitioners working with JAX who need flexible, modular optimization tools. It depends on JAX, jaxlib, numpy, and absl-py. The library evolved from an earlier experimental JAX module and is now maintained as a standalone project by DeepMind. It supports current Python versions and is actively maintained, making it suitable for both research prototyping and production use in JAX-based machine learning workflows.

Use it for:

• Training neural networks with custom optimizer combinations by composing Optax building blocks.
• Implementing gradient clipping, weight decay, or other transformations in a modular way.
• Prototyping new optimization algorithms by combining existing components.
• Using standard optimizers like Adam or RMSprop in JAX-based machine learning projects.
• Computing loss functions like L2 or cross-entropy within JAX training loops.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Optax provides composable building blocks for gradient processing and optimization in JAX, including implementations of popular optimizers and loss functions that can be combined into custom solutions.

Yes. Optax is actively maintained, has no known vulnerabilities, installs with low friction, and is the standard gradient optimization library for JAX. Install it if you are building machine learning systems with JAX and need flexible, composable optimizer and loss components. The only gotcha is the compiled JAX/jaxlib dependency, which may require platform-specific setup.

Install

optax on PyPI

pip install optax

uv add optax

poetry add optax

Installing optax

Before you install

Low friction install as a pure Python wheel. Actively maintained with recent releases; last commit 2026-08-07 and latest release 2026-03-20. Requires JAX and its compiled dependency jaxlib, which may add setup complexity depending on your platform.

License in practice

Licensed under Apache Software License (permissive). No restrictions on commercial or private use.

Quickstart

pip install optax

import optax
import jax
import jax.numpy as jnp

optimizer = optax.adam(learning_rate)
params = {'w': jnp.ones((num_weights,))}
opt_state = optimizer.init(params)

compute_loss = lambda params, x, y: optax.l2_loss(params['w'].dot(x), y)
grads = jax.grad(compute_loss)(params, xs, ys)
updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)

Requires JAX and jaxlib (compiled dependency); Python >= 3.10.

Verify before relying

• Whether the library's optimizer implementations have been benchmarked against standard baselines in recent comparisons.
• Performance characteristics when used with large-scale models or distributed training setups.

Package facts

Evidence: optax-0.2.8-py3-none-any.whl

Keywords: python, machine learning, reinforcement-learning

Tags

More Python Modules packages

Further reading