---
id: optax
version: "0.2.8"
license: unclear
license_treatment: permissive
maintenance: active
---

# optax — A gradient processing and optimization library in JAX.

License: permissive · Maintenance: active · Downloads: 3.8M/mo

## 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 above — 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

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

- License: not declared (permissive)
- Python support: supports_current
- Install friction: low
- Maintenance: active
- Downloads: 3.8M/month (top 5,000 on PyPI)
- Known vulnerabilities: none known

## Tags

JAX optimizer library, gradient processing JAX, Adam optimizer JAX, custom optimizers JAX, machine learning optimization, neural network training JAX, loss functions JAX, jax-ecosystem, gradient-optimization, composable-components

[View on SkillFed](https://skillfed.io/packages/optax) · [View on PyPI](https://pypi.org/project/optax/)