# google-deepmind/optax

Optax is a gradient processing and optimization library for JAX.

Repository: https://github.com/google-deepmind/optax
Canonical: https://ross.abutalabs.com/products/optax
Homepage: https://optax.readthedocs.io
Language: Python
License: Apache-2.0
License Family: permissive
Topics: machine-learning, optimization
Last push: 2026-08-20T07:32:16+00:00

## Health v2 (maintenance only)
Score: 82/100 (v2, computed 2026-09-03T02:20:16.233290+00:00)
- activity 98, release rhythm 51, longevity 100
- inputs: {"age_days": 2273, "days_push": 13, "days_rel": 166, "gap_med": 119.5, "n_releases_24m": 5}
- flags: none
- formula: round(0.45*activity + 0.35*rhythm + 0.20*longevity); archived -> min(score, 10)

## Adoption (not part of the score)
Stars 2325, forks 361 (observed 2026-08-28T04:06:37.360091+00:00)

## What it is
Optax is a gradient processing and optimization library for JAX, offering composable building blocks like optimizers and loss functions. It is developed by DeepMind to support machine learning research with easily recombined components.

## Use cases
- train neural networks in jax with adam optimizer
- implement custom optimizers from composable gradient transforms
- compute loss functions for machine learning models
- run gradient descent research experiments in jax
- optimize reinforcement learning policies
- schedule learning rates during training

## When to choose
- you use JAX and need well-tested optimizers or loss functions
- you want to compose custom gradient processing pipelines
- you need readable implementations that match standard equations

## When to avoid
- you work in PyTorch or TensorFlow rather than JAX
- you need a full high-level training framework with data loading and experiment tracking

## Facets
- artifact type: library
- maturity: stable
- function: machine-learning, llm-training, deep-learning
- domain: machine-learning, deep-learning, reinforcement-learning
- platform: python, cross-platform
- tags: jax, optimizers, gradient-processing, loss-functions, deepmind, algorithms, gpu

## Member repositories
- google-deepmind/optax (main) score 82

## Provenance
- Observed fields: from GitHub, fetched 2026-08-28T04:06:37.360091+00:00.
- Health v2: computed from the inputs above; adoption is never an input.
- Inferred fields (summary, facets, guidance): AI-extracted, prompt v1, taxonomy v1, on 2026-08-30T02:38:24.040560+00:00, confidence not recorded.
  - readme: https://github.com/google-deepmind/optax (fetched 2026-08-28T04:06:37.360091+00:00, sha ceef840458cf)
  - registry_pypi: https://pypi.org/pypi/optax/json (fetched 2026-08-29T10:19:12.351620+00:00, sha bf601364e27e)
- Data as of 2026-08-30T08:39:29.467469+00:00.