← Learn JAX

Optax

JAX

Gradient transforms, optimizer chains, schedules, weight decay, EMA, masking. The production optimizer library for JAX.

This track uses JAX. You can make it your practice mode to browse related content across the site.

25 exercises

Explore freely. Sign in to keep your progress.

Start track

Your learning sequence

Check yourself

4 questions · not saved while you're signed out

Use these questions to check your understanding and compare explanations. Your score does not affect track completion.

0 / 4 correct

For a single scalar gradient of 100, what update does clip_first return? Would reversing the transformations change it?

Choose one answer to reveal the explanations. Answers are kept only while this page is open.

import jax.numpy as jnp
import optax

params = jnp.array(0.0)
grad = jnp.array(100.0)
clip_first = optax.chain(optax.clip_by_global_norm(1.0), optax.sgd(0.1))
scale_first = optax.chain(optax.sgd(0.1), optax.clip_by_global_norm(1.0))
for opt in (clip_first, scale_first):
    updates, _ = opt.update(grad, opt.init(params), params)
    print(f"{float(updates):.1f}")

Question 1 of 4