Optax

Warnings

• breaking The signature of `optimizer.update` changed in Optax 0.1.x to 0.2.x to include `params` as the third argument: `optimizer.update(grads, opt_state, params)`. This is crucial for optimizers that might need current parameter values to compute updates (e.g., for weight decay).
  Fix: Ensure your `optimizer.update` calls pass the current model parameters as the third argument. If you're coming from an older version, add `params` to your `update` function signature.
• gotcha Optax, like JAX, operates on immutable data structures. When you call `optimizer.update`, it returns a *new* optimizer state and *new* updates. You must reassign these values (`opt_state = new_opt_state`, `params = optax.apply_updates(params, updates)`) otherwise your training loop will not progress.
  Fix: Always reassign the optimizer state and parameters after calling `optimizer.update` and `optax.apply_updates` respectively, as demonstrated in the quickstart example.
• gotcha Optax expects parameters, gradients, and optimizer states to be JAX PyTrees (e.g., nested dictionaries, lists, tuples, or custom types registered with `jax.tree_util`). Passing non-PyTree structures or incompatible types can lead to errors.
  Fix: Ensure that your model parameters and the gradients produced by `jax.grad` are consistent PyTree structures. When composing optimizers or transformations, confirm that intermediate results maintain PyTree compatibility.

Install

• pip install optax jax jaxlib Basic Installation
• pip install 'optax[accelerate]' jax[cuda12_pip] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html CUDA-enabled Installation (example)

Imports

• optax

  import optax

• adam

  import optax
  optimizer = optax.adam(...)

  Optimizers like adam are directly under the `optax` namespace, not a nested `optimizers` submodule.
• chain

  from optax import chain
  optimizer = chain(...)

  `chain` is a function directly under the `optax` module, not a submodule.

Quickstart

This quickstart demonstrates a basic training loop with Optax and JAX. It defines a simple linear model and a mean squared error loss. An Adam optimizer is initialized, and a `train_step` function is created, leveraging `jax.grad` to compute gradients and Optax to apply updates. The `train_step` is `jax.jit`-compiled for efficiency. The loop iteratively updates parameters and optimizer state to minimize the loss.

import jax
import jax.numpy as jnp
import optax

# 1. Define a simple model and loss function
def model(params, x):
    return params['w'] * x + params['b']

def loss_fn(params, x, y):
    predictions = model(params, x)
    return jnp.mean((predictions - y)**2)

# 2. Initialize parameters
key = jax.random.PRNGKey(0)
params = {
    'w': jax.random.normal(key, ()),
    'b': jax.random.normal(key, ())
}

# 3. Choose an optimizer (e.g., Adam)
learning_rate = 0.01
optimizer = optax.adam(learning_rate)

# 4. Initialize optimizer state
opt_state = optimizer.init(params)

# 5. Sample data for training
x_data = jnp.array([1.0, 2.0, 3.0, 4.0])
y_data = jnp.array([2.0, 4.0, 6.0, 8.0]) # Target: y = 2x

# 6. Define a single training step using JAX's jit for performance
@jax.jit
def train_step(params, opt_state, x, y):
    # Compute loss and gradients
    loss_value, grads = jax.value_and_grad(loss_fn)(params, x, y)

    # Compute updates from gradients and optimizer state
    updates, new_opt_state = optimizer.update(grads, opt_state, params)

    # Apply updates to parameters
    new_params = optax.apply_updates(params, updates)

    return new_params, new_opt_state, loss_value

# 7. Training loop
print(f"Initial parameters: {params}")
for i in range(100):
    params, opt_state, loss_value = train_step(params, opt_state, x_data, y_data)
    if i % 20 == 0:
        print(f"Step {i}, Loss: {loss_value:.4f}")

print(f"Final parameters: {params}")
# Expected output for w: ~2.0, b: ~0.0