### Install and Import Optax

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/gradient_accumulation_and_microbatching.md

Installs the optax library and imports necessary modules for the examples.

```python
# ! pip install -q "optax @ git+https://github.com/google-deepmind/optax"
import jax
import jax.numpy as jnp
import optax
from flax import nnx
import functools
import time
from optax import microbatching
import gc
```

--------------------------------

### Install Documentation Dependencies

Source: https://github.com/google-deepmind/optax/blob/main/README.md

Installs the necessary dependencies for building Optax documentation.

```sh
pip install -e ".[docs]"
```

--------------------------------

### Gradient Transformation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates a simple gradient transformation using Optax. This example shows how to apply a scaling factor to gradients.

```python
tx = optax.scale(0.1)
state = tx.init(params)
updates, state = tx.update(grads, state, params)
```

--------------------------------

### Install Optax from GitHub

Source: https://github.com/google-deepmind/optax/blob/main/docs/index.md

Install the most recent version of Optax directly from its GitHub repository using pip.

```default
pip install git+git://github.com/google-deepmind/optax.git
```

--------------------------------

### Install Optax from PyPI

Source: https://github.com/google-deepmind/optax/blob/main/README.md

Use this command to install the latest released version of Optax.

```sh
pip install optax
```

--------------------------------

### Install NetworkX

Source: https://github.com/google-deepmind/optax/blob/main/examples/linear_assignment_problem.ipynb

Install the NetworkX library for graph manipulation. This is a prerequisite for visualizing the assignment problem as a graph.

```bash
pip install -U networkx
```

--------------------------------

### Install Optax Development Version

Source: https://github.com/google-deepmind/optax/blob/main/README.md

Install the latest development version of Optax directly from GitHub.

```sh
pip install git+https://github.com/google-deepmind/optax.git
```

--------------------------------

### Import Libraries and Install NetworkX

Source: https://github.com/google-deepmind/optax/blob/main/examples/linear_assignment_problem.ipynb

Import necessary libraries including NetworkX, JAX for random number generation, Optax for the algorithm, and Matplotlib for plotting. Installs NetworkX quietly.

```python
!pip install -q -U networkx
import networkx as nx
from jax import random
import optax
from matplotlib import pyplot as plt
```

--------------------------------

### Gradient Transformation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/freezing_parameters.md

Shows how to apply a simple gradient transformation.

```python
import optax

# Example: Gradient transformation that adds a constant value
# This is a simplified example; real transformations are more complex.
# optax.add_constant(value=1e-6)
```

--------------------------------

### Import necessary libraries in Optax

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/muon.md

Imports JAX, JAX NumPy, and Optax for use in the optimizer examples. Ensure these libraries are installed.

```python
from pprint import pprint

import jax
import jax.numpy as jnp
from jax import random

import optax
```

--------------------------------

### Basic Gradient Transformation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/linear_assignment_problem.md

Demonstrates a simple gradient transformation using Optax. Requires importing necessary modules.

```python
import optax

# Create an optimizer
optimizer = optax.adam(learning_rate=0.001)

# Initialize the optimizer state
params = {'w': jnp.array([1.0]), 'b': jnp.array([2.0])}
opt_state = optimizer.init(params)

# Compute gradients (example)
grads = {'w': jnp.array([0.1]), 'b': jnp.array([0.2])}

# Apply gradients and update parameters
updates, opt_state = optimizer.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Basic Gradient Transformation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Demonstrates a simple gradient transformation using Optax. This is useful for basic optimization tasks.

```python
import optax

# Define a simple gradient transformation
transforms = [
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=0.001)
]

# Combine transformations
optimizer = optax.chain(*transforms)

# Initialize optimizer state
params = {'w': jnp.array([1.0]), 'b': jnp.array([2.0])}
opt_state = optimizer.init(params)

# Example gradients
grads = {'w': jnp.array([0.5]), 'b': jnp.array([-0.1])}

# Apply updates
updates, opt_state = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)

print(f"Original params: {params}")
print(f"Gradients: {grads}")
print(f"Updates: {updates}")
print(f"New params: {new_params}")
print(f"Optimizer state: {opt_state}")
```

--------------------------------

### Gradient Accumulation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Shows how to implement gradient accumulation using `optax.accumulate_gradients`. This allows for simulating larger batch sizes.

```python
import optax

# Create an optimizer (e.g., Adam)
base_optimizer = optax.adam(learning_rate=0.001)

# Wrap it with gradient accumulation
accumulating_optimizer = optax.accumulate_gradients(base_optimizer, num_steps=4)

# Initialize state and update gradients
# state = accumulating_optimizer.init(params)
# updates = accumulating_optimizer.update(grads, state, params)

```

--------------------------------

### Gradient Accumulation

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/reduce_on_plateau.md

This example demonstrates gradient accumulation using Optax. It allows for larger effective batch sizes.

```python
import optax

# Define the base optimizer
base_optimizer = optax.adam(learning_rate=0.001)

# Wrap it with gradient accumulation
optimizer = optax.accumulate_gradients(base_optimizer, num_steps=4)
```

--------------------------------

### Gradient Transformation with Wrap (Example)

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates wrapping a gradient transformation, such as chaining `clip_by_global_norm` with `adam`.

```python
import optax

clipped_adam = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=0.001)
)
# The 'clipped_adam' is itself a gradient transformation.
```

--------------------------------

### Gradient Accumulation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/linear_assignment_problem.md

Shows how to accumulate gradients over multiple steps before applying an update. This simulates a larger batch size.

```python
import optax

# Create a gradient accumulation transformation
accumulator = optax.accumulate_gradients(num_steps=4)

# This transformation can be chained with other optimizers
```

--------------------------------

### Gradient Transformation with Parameters

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

This example shows a gradient transformation that takes parameters, such as the `scale` transformation.

```python
import optax

optimizer = optax.chain(
    optax.scale(0.5),  # Scale gradients by 0.5
    optax.adam(learning_rate=0.001)
)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3., 4.])}
opt_state = optimizer.init(params)

grads = {'w': jnp.array([0.1, 0.2]), 'b': jnp.array([0.3, 0.4])}

updates, opt_state = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)
```

--------------------------------

### Initialize Gradient Transformation

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Creates a basic gradient transformation. This example uses `optax.identity` which applies no transformation.

```python
import optax

transforms = optax.identity()

# Example usage (requires state and gradients)
# state = transforms.init(params)
# updates, state = transforms.update(grads, state, params)
```

--------------------------------

### Gradient Accumulation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Demonstrates gradient accumulation by summing gradients over multiple steps before updating parameters.

```python
import optax
import jax
import jax.numpy as jnp

learning_rate = 0.001
accumulation_steps = 4

optimizer = optax.adam(learning_rate)
opt_state = optimizer.init(params)

accumulated_grads = None

for step in range(total_steps):
    # Compute gradients for the current batch
    loss, grads = compute_loss_and_grads(params)

    if accumulated_grads is None:
        accumulated_grads = jax.tree_map(lambda x: x / accumulation_steps, grads)
    else:
        accumulated_grads = jax.tree_map(lambda x, y: x + y / accumulation_steps, accumulated_grads, grads)

    if (step + 1) % accumulation_steps == 0:
        updates, opt_state = optimizer.update(accumulated_grads, opt_state, params)
        params = optax.apply_updates(params, updates)
        accumulated_grads = None # Reset accumulated gradients

        # Update learning rate if using a schedule
        # opt_state = ...

```

--------------------------------

### Gradient Transformation with Initialization

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

This example shows a gradient transformation that requires initialization, such as momentum buffers in SGD.

```python
import optax

optimizer = optax.sgd(learning_rate=0.01, momentum=0.9)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3., 4.])}
opt_state = optimizer.init(params)

# opt_state now contains the momentum buffers initialized by the optimizer.
```

--------------------------------

### Zeroing Gradients Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/linear_assignment_problem.md

Demonstrates how to zero out gradients. This can be useful in specific scenarios, like when gradients should not contribute to updates.

```python
import optax

# Create a gradient transformation that zeros gradients
zeroing_transform = optax.zero_grad()

# This transformation can be applied using optax.chain
```

--------------------------------

### Optimizer with Momentum

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/reduce_on_plateau.md

This example shows how to create an optimizer with momentum using Optax. It includes the necessary steps for initialization and updating.

```python
import optax

# Create an optimizer with momentum
optimizer = optax.adam(learning_rate=0.001, b1=0.9, b2=0.999)

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

# Apply gradients (example)
# updates, state = optimizer.update(grads, state, params)
# params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Scaling

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/reduce_on_plateau.md

This example demonstrates gradient scaling using Optax, often used in conjunction with mixed precision training.

```python
import optax

# Create a gradient scaler
scaler = optax. স্কেলার(init_value=2.**16)

# Use it with an optimizer
optimizer = optax.chain(
    scaler.add_signed_graph_loss_with_gradient_scaling(),
    optax.adam(learning_rate=0.001)
)
```

--------------------------------

### Zeroing Gradients Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Shows how to use `optax.zero_grads` to reset gradients to zero after they have been applied. This is a common step in training loops.

```python
import optax

# Assume 'updates' are computed gradients
# updates = optimizer.update(grads, state, params)

# Zero out the gradients
# grads = optax.zero_grads(grads)

```

--------------------------------

### Gradient Clipping

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/reduce_on_plateau.md

This example demonstrates how to clip gradients by their global norm using Optax. This is useful for stabilizing training.

```python
import optax

# Create an optimizer with gradient clipping
optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=0.001)
)
```

--------------------------------

### Import necessary libraries for Optax

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/reduce_on_plateau.md

Imports JAX, Flax, Matplotlib, NumPy, Grain, TensorFlow Datasets, and Optax modules. This setup is required for running the examples.

```python
from typing import Sequencenfrom flax import linen as nnnimport jaxnimport jax.numpy as jnpnimport matplotlib.pyplot as pltnimport numpy as npnimport grainnimport tensorflow_datasets as tfdsn
import optaxnimport optax.treenimport from optax import contribn

# Show on which platform JAX is running.nprint("JAX running on", jax.devices()[0].platform.upper())
```

--------------------------------

### Generate data and initialize model/optimizer

Source: https://github.com/google-deepmind/optax/blob/main/examples/gradient_accumulation.ipynb

Generates synthetic image data and corresponding labels for testing. Initializes model parameters and a base SGD optimizer. This setup is used for subsequent examples.

```python
EXAMPLES = jax.random.uniform(jax.random.PRNGKey(0), (9, 28, 28, 1))
LABELS = jax.random.randint(jax.random.PRNGKey(0), (9,), minval=0, maxval=10)

optimizer = optax.sgd(1e-4)
params = net.init(jax.random.PRNGKey(0), EXAMPLES)

```

--------------------------------

### Basic Optimizer Usage

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/ogda_example.md

Demonstrates how to initialize and update an optimizer. Requires importing the desired optimizer from optax.

```python
import optax

# Initialize an optimizer
optimizer = optax.adam(learning_rate=0.001)

# Initialize optimizer state
params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3.])}
opt_state = optimizer.init(params)

# Compute gradients (example)
grads = {'w': jnp.array([0.1, 0.2]), 'b': jnp.array([0.3])}

# Update parameters
updates, opt_state = optimizer.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Transformation with Different Initialization Values for State

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

This example demonstrates how different initializations of optimizer state can affect training.

```python
import optax

optimizer = optax.adam(learning_rate=0.001)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3., 4.])}

# Default initialization
opt_state_default = optimizer.init(params)

# Manually creating a state with different initial values (e.g., for debugging or specific scenarios)
# This is more advanced and depends on the optimizer's state structure.
# For Adam, state includes 'count', 'mu', 'nu'. Let's modify 'mu'.
initial_mu = {'w': jnp.zeros_like(params['w']), 'b': jnp.zeros_like(params['b'])}
initial_nu = {'w': jnp.zeros_like(params['w']), 'b': jnp.zeros_like(params['b'])}
initial_count = jnp.array(100) # Start count from 100

opt_state_custom = optax.ScaleByAdamState(count=initial_count, mu=initial_mu, nu=initial_nu)

grads = {'w': jnp.array([0.1, 0.2]), 'b': jnp.array([0.3, 0.4])}

updates_default, _ = optimizer.update(grads, opt_state_default, params)
updates_custom, _ = optimizer.update(grads, opt_state_custom, params)
```

--------------------------------

### Gradient Transformation with Scale By Schedule (Example)

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates scaling updates using a linear learning rate decay schedule.

```python
import optax

schedule = optax.linear_decay(init_value=0.1, end_value=0.01, transition_steps=1000)
scale_by_schedule_transform = optax.scale_by_schedule(schedule)
# Example usage:
# updates, state = scale_by_schedule_transform.update(grads, state, params)
```

--------------------------------

### Import necessary libraries for DP-SGD

Source: https://github.com/google-deepmind/optax/blob/main/examples/contrib/differentially_private_sgd.ipynb

Imports required libraries including JAX, Optax, DP-accounting, and data loading utilities. This setup is essential for running the differentially private training example.

```python
import warnings
import dp_accounting
import jax
import jax.numpy as jnp
from optax import contrib
from optax import losses
import optax
from jax.example_libraries import stax
import grain
import tensorflow_datasets as tfds
import matplotlib.pyplot as plt
```

--------------------------------

### Multi-Device Training Setup

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/ogda_example.md

Illustrates how to set up Optax optimizers for multi-device training using JAX's `pmap`. This involves replicating the optimizer state across devices.

```python
import optax
from jax import numpy as jnp
from jax.experimental import pmap

# Define a base optimizer
optimizer = optax.adam(learning_rate=0.001)

# Initialize optimizer state on a single device
params = {'w': jnp.array([1.0])}
opt_state = optimizer.init(params)

# Replicate the optimizer state across multiple devices using pmap
# This assumes you have multiple JAX devices available
replicated_opt_state = jax.tree_util.tree_map(lambda x: jax.device_put(x, jax.devices()[0]), opt_state)
replicated_opt_state = jax.pmap(lambda x: x)(replicated_opt_state)

# In a pmapped training step, you would use the replicated state
# @jax.pmap
# def train_step_pmapped(...):
#     ... compute grads ...
#     updates, replicated_opt_state = optimizer.update(grads, replicated_opt_state, params)
#     params = optax.apply_updates(params, updates)
#     return params, replicated_opt_state

print("Optimizer state replicated for multi-device training.")
```

--------------------------------

### Using Optax with JAX

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Illustrates how to integrate Optax with JAX for efficient gradient-based optimization.

```python
import jax
import jax.numpy as jnp
import optax

# Define a simple model (e.g., linear regression)
def predict(params, x):
    return params['w'] * x + params['b']

# Define loss function (e.g., mean squared error)
def mse_loss(params, x, y):
    y_pred = predict(params, x)
    return jnp.mean((y_pred - y) ** 2)

# Define the training step function
@jax.jit
def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(mse_loss)(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

# Initialize parameters and optimizer
params = {'w': jnp.array([1.0]), 'b': jnp.array([0.0])}
optimizer = optax.adam(learning_rate=0.01)
opt_state = optimizer.init(params)

# Example data
x_train = jnp.array([1.0, 2.0, 3.0, 4.0])
y_train = jnp.array([2.0, 4.0, 6.0, 8.0])

# Run a few training steps
for step in range(100):
    params, opt_state, loss = train_step(params, opt_state, x_train, y_train)
    if step % 20 == 0:
        print(f"Step {step}, Loss: {loss:.4f}, Params: {params}")

print(f"Final trained parameters: {params}")
```

--------------------------------

### Initialize AdamW Optimizer and State

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/nanolm.md

Initializes the AdamW optimizer with a specified learning rate and then initializes its state using the model's parameters. This setup is crucial before starting the training loop.

```python
# To run with SGD instead of adam, replace adam with sgd
opt = optax.adamw(learning_rate=LEARNING_RATE)

opt_state = opt.init(var_params)
```

--------------------------------

### Dynamic Parameter Freezing with selective_transform

Source: https://github.com/google-deepmind/optax/blob/main/examples/freezing_parameters.ipynb

Demonstrates dynamically swapping which parameters are frozen by re-creating the optimizer with a new mask. This example starts with param0 trainable and param1 frozen, then swaps at step 200. Ensure `jax`, `optax`, and `jnp` are imported.

```python
# Initialize parameters and mask
params = {"param0": jnp.array([10.0]), "param1": jnp.array([10.0])}
mask = {"param0": False, "param1": True}

# Create optimizer with selective_transform utility
optimizer = selective_transform(optax.adam(0.1), freeze_mask=mask)
opt_state = optimizer.init(params)

# Initialize history lists to track progress
loss_history = []
param0_history = []
param1_history = []

# Training loop with dynamic mask swap at step 200
for step in range(1, 401):
    # At step 200, swap which parameter is frozen
    if step == 200:
        mask = {"param0": True, "param1": False}
        optimizer = selective_transform(optax.adam(0.1), freeze_mask=mask)
        opt_state = optimizer.init(params)  # ⚠️ resets optimizer state

    # Compute loss and gradients
    loss, grads = jax.value_and_grad(loss_fn)(params)

    # Update parameters using optimizer
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)

    # Record loss and parameter values
    loss_history.append(loss)
    param0_history.append(params["param0"][0])
    param1_history.append(params["param1"][0])

    # Print progress every 100 steps
    if step % 100 == 0:
        print(f"Step {step:3d}  loss={loss:.4f}  param0={params['param0'][0]:.4f}  param1={params['param1'][0]:.4f}")

# Verify that param0 reached near zero and param1 was updated
assert jnp.isclose(params["param0"][0], 0.0, atol=1e-2)
assert params["param1"][0] != 10.0
print("✅ Dynamic swap worked param0: train→frozen, param1: frozen→train")
```

--------------------------------

### Basic Gradient Transformation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates a simple gradient transformation using `optax.chain`. This is useful for applying multiple transformations sequentially.

```python
import optax

# Define a sequence of transformations
transforms = [
    optax.scale(0.1),
    optax.clip_by_global_norm(1.0),
]

# Chain them together
optimizer = optax.chain(*transforms)

# Initialize the optimizer state (e.g., for momentum)
# state = optimizer.init(params)

# Apply the transformations to gradients
# updates = optimizer.update(grads, state, params)

```

--------------------------------

### Gradient Accumulation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/ogda_example.md

Demonstrates how to implement gradient accumulation using Optax transformations. This allows for simulating larger batch sizes than memory permits.

```python
import optax
from jax import numpy as jnp

# Define the base optimizer (e.g., SGD)
base_optimizer = optax.sgd(learning_rate=0.01)

# Configure gradient accumulation
# Accumulate gradients over 4 steps
accumulation_steps = 4
optimizer = optax.accumulate_gradients(base_optimizer, num_steps=accumulation_steps)

# Initialize optimizer state
params = {'w': jnp.array([1.0])}
opt_state = optimizer.init(params)

# In a training loop, you would call optimizer.update() multiple times
# before the accumulated gradients are applied and the state reset.
# For example, after 4 calls to update:
# grads_step_1 = {'w': jnp.array([0.1])}
# updates_1, opt_state = optimizer.update(grads_step_1, opt_state, params)
# grads_step_2 = {'w': jnp.array([0.2])}
# updates_2, opt_state = optimizer.update(grads_step_2, opt_state, params)
# ... and so on for 4 steps.
# After the 4th step, 'updates' will contain the averaged gradients,
# and the internal accumulation state will be reset.

print(f"Gradient accumulation configured for {accumulation_steps} steps.")
```

--------------------------------

### Initialize and Update with SGD

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Shows the process of initializing and updating parameters with the Stochastic Gradient Descent (SGD) optimizer. Ensure `optax` is imported.

```python
import optax

# Initialize optimizer state
optimizer = optax.sgd(learning_rate=0.01)
params = ... # Your model parameters
state = optimizer.init(params)

# Compute gradients
grads = ...

# Update parameters and optimizer state
updates, state = optimizer.update(grads, state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Build Optax Documentation

Source: https://github.com/google-deepmind/optax/blob/main/README.md

Commands to navigate to the docs directory and build the HTML documentation.

```sh
cd docs
make html
```

--------------------------------

### Learning Rate Schedules

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Illustrates how to use learning rate schedules with Optax optimizers. This example uses a linear warm-up followed by a cosine decay.

```python
import optax

# Define a learning rate schedule
num_train_steps = 10000
warmup_steps = 1000

schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=0.001,
    warmup_steps=warmup_steps,
    decay_steps=num_train_steps - warmup_steps,
    end_value=0.0001
)

# Use the schedule with an optimizer
optimizer = optax.adam(learning_rate=schedule)

# Initialize and use as usual
params = ...
opt_state = optimizer.init(params)
grads = ...
updates, opt_state = optimizer.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Initialize Optimizer and Parameters

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/differentially_private_sgd.md

Initializes the optimizer (either DP-SGD or standard SGD) and the model parameters. Requires JAX, Optax, and relevant constants like LEARNING_RATE, L2_NORM_CLIP, NOISE_MULTIPLIER, and DPSGD flag.

```python
if DPSGD:
  tx = contrib.dpsgd(
      learning_rate=LEARNING_RATE, l2_norm_clip=L2_NORM_CLIP,
      noise_multiplier=NOISE_MULTIPLIER, seed=1337)
else:
  tx = optax.sgd(learning_rate=LEARNING_RATE)

_, params = init_random_params(jax.random.PRNGKey(1337), (-1, 28, 28, 1))
opt_state = tx.init(params)
```

--------------------------------

### Matplotlib Figure and Axes Setup

Source: https://github.com/google-deepmind/optax/blob/main/examples/contrib/rosenbrock_ademamix.ipynb

This snippet shows the basic setup for creating a Matplotlib figure and axes, preparing for plotting.

```python
fig = plt.figure()
ax = fig.subplots(1, 2)
ax[0].set_xlabel("x")
ax[0].set_ylabel("y")
```

--------------------------------

### Initialize and Update with RMSprop

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Illustrates initializing and updating model parameters using the RMSprop optimizer. Make sure to import `optax`.

```python
import optax

# Initialize optimizer state
optimizer = optax.rmsprop(learning_rate=0.001)
params = ... # Your model parameters
state = optimizer.init(params)

# Compute gradients
grads = ...

# Update parameters and optimizer state
updates, state = optimizer.update(grads, state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Transformation Pipeline Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

A comprehensive example showing a pipeline of multiple gradient transformations for a typical training scenario.

```python
import optax
import jax.numpy as jnp

# Define a complex pipeline: clip, scale, schedule, and Adam
transforms = [
    optax.clip_by_global_norm(1.0),
    optax.scale(0.5), # Scale gradients before Adam
    optax.scale_by_schedule(optax.cosine_decay(init_value=1.0, decay_steps=1000)),
    optax.adam(learning_rate=0.001)
]
optimizer = optax.chain(*transforms)

# Initialize optimizer state
params = {'w': jnp.array([1.0])}
opt_state = optimizer.init(params)

# Example gradients
grads = {'w': jnp.array([0.5])}

# Apply updates
updates, opt_state = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)

print(f"Original params: {params}")
print(f"Gradients: {grads}")
print(f"Updates: {updates}")
print(f"New params: {new_params}")
print(f"Optimizer state: {opt_state}")
```

--------------------------------

### Define Training Step for Adversarial Examples

Source: https://github.com/google-deepmind/optax/blob/main/examples/adversarial_training.ipynb

Defines a JAX-jitted training step that generates adversarial images using PGD attack and then trains the model on these adversarial examples. It also applies L2 regularization.

```python
@jax.jit
def train_step(params, opt_state, batch):
  images, labels = batch
  # convert images to float as attack requires to take gradients wrt to them
  images = images.astype(jnp.float32) / 255
  adversarial_images_train = pgd_attack(images, labels, params, epsilon=EPSILON)
  # train on adversarial images
  loss_grad_fun = jax.grad(loss_fun)
  grads = loss_grad_fun(params, L2_REG, (adversarial_images_train, labels))
  updates, opt_state = optimizer.update(grads, opt_state)
  params = optax.apply_updates(params, updates)
  return params, opt_state
```

--------------------------------

### Gradient Transformation with State (Shampoo with State)

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates how to use Shampoo with explicit state management.

```python
optimizer = optax.chain(optax.shampoo(learning_rate=1e-3, weight_decay=0.01))
state = optimizer.init(params)

# Inside training loop:
grads = ...
updates, state = optimizer.update(grads, state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Transformation with Different Masks

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

This example demonstrates applying gradients with different masks.

```python
import optax

optimizer = optax.adam(learning_rate=0.001)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3., 4.])}
opt_state = optimizer.init(params)

grads = {'w': jnp.array([0.1, 0.2]), 'b': jnp.array([0.3, 0.4])}

mask1 = {'w': jnp.array([1., 0.]), 'b': jnp.array([0., 1.])}
mask2 = {'w': jnp.array([0., 1.]), 'b': jnp.array([1., 0.])}

masked_grads1 = optax.apply_mask(grads, mask1)
masked_grads2 = optax.apply_mask(grads, mask2)

updates1, opt_state1 = optimizer.update(masked_grads1, opt_state, params)
updates2, opt_state2 = optimizer.update(masked_grads2, opt_state, params)
```

--------------------------------

### Basic Gradient Transformation Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/linear_assignment_problem.md

Demonstrates a simple gradient transformation using Optax. This is useful for understanding the basic flow of applying transformations to gradients.

```python
import optax

# Define a simple gradient transformation
transforms = [
    optax.clip_by_global_norm(1.0),
    optax.scale(0.1)
]

# Compose the transformations
optimizer = optax.chain(*transforms)

# Initialize the optimizer state
params = {'w': 1.0, 'b': 2.0}
opt_state = optimizer.init(params)

# Example gradients
grads = {'w': 0.5, 'b': -0.3}

# Apply the transformations to the gradients
updates, opt_state = optimizer.update(grads, opt_state, params)

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

print(f"Original params: {params}")
print(f"Gradients: {grads}")
print(f"Updates: {updates}")
print(f"New params: {new_params}")
```

--------------------------------

### Define a Pytree Input

Source: https://github.com/google-deepmind/optax/blob/main/examples/perturbations.ipynb

Example of creating a pytree structure with JAX arrays for input.

```python
tree_a = (jnp.array((0.1, 0.4, 0.5)),
          {'k1': jnp.array((0.1, 0.2)),
           'k2': jnp.array((0.1, 0.1))},
          jnp.array((0.4, 0.3, 0.2, 0.1)))

```

--------------------------------

### Initialize and Update with Adam

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates how to initialize and update model parameters using the Adam optimizer. Requires importing `optax`.

```python
import optax

# Initialize optimizer state
optimizer = optax.adam(learning_rate=0.001)
params = ... # Your model parameters
state = optimizer.init(params)

# Compute gradients (e.g., using JAX's grad)
grads = ...

# Update parameters and optimizer state
updates, state = optimizer.update(grads, state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Averaging

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Averages gradients across multiple devices. Useful for distributed training setups.

```python
import optax

# Create a gradient averaging transformation
optimizer = optax.chain(
    optax.all_reduce(num_devices=8, reduction=optax.Reduction.MEAN),
    optax.adamw(learning_rate=0.001)
)
```

--------------------------------

### Learning Rate Scheduling with Optax

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/linear_assignment_problem.md

Shows how to use a learning rate schedule with an optimizer. This example uses a linear decay schedule.

```python
import optax

# Define a learning rate schedule
schedule = optax.linear_decay(
    init_value=0.1, end_value=0.01, transition_steps=1000
)

# Create an optimizer with the schedule
optimizer = optax.adam(learning_rate=schedule)

# Initialization and update steps would follow as in the basic example
```

--------------------------------

### Basic Optimizer Usage

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates how to initialize and use a basic optimizer like Adam. Requires importing the optimizer and applying it to gradients.

```python
import optax

# Initialize an optimizer
optimizer = optax.adam(learning_rate=0.001)

# Initialize optimizer state
params = ... # Your model parameters
opt_state = optimizer.init(params)

# Compute gradients (e.g., using JAX's grad)
grads = ...

# Update parameters
updates, opt_state = optimizer.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Transformation with State

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/linear_assignment_problem.md

Illustrates a gradient transformation that maintains its own state. Momentum is an example of such a transformation.

```python
import optax

# Create a momentum optimizer
optimizer = optax.chain(
    optax.momentum(learning_rate=0.01, momentum=0.9),
    optax.adam(learning_rate=0.001) # Example of chaining
)

# The opt_state will contain momentum information
```

--------------------------------

### Single Batch Training Loop

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/gradient_accumulation.md

Runs a training loop over a single batch containing all examples.

```python
new_params_single_batch = fit(
    optimizer,
    params,
    batches=[dict(image=EXAMPLES, label=LABELS,)]
)
```

--------------------------------

### Optax.MultiSteps Optimizer Setup

Source: https://github.com/google-deepmind/optax/blob/main/examples/gradient_accumulation_and_microbatching.ipynb

Wraps an existing optimizer with optax.MultiSteps to handle gradient accumulation automatically. Simplifies the update function by managing accumulated gradients internally.

```python
import functools
import jax
import optax

multi_adam = optax.MultiSteps(adamw, accumulation_steps)

@functools.partial(jax.jit, donate_argnums=(0, 2))
def update_fn_v2(params, batch, opt_state):
  grads = jax.grad(loss_fn)(params, batch)
  updates, opt_state = multi_adam.update(grads, opt_state, params)
  params = optax.apply_updates(params, updates)
  return params, opt_state

multi_opt_state = multi_adam.init(params)
update_fn_v2.lower(params, batch, multi_opt_state).compile()
```

--------------------------------

### Import necessary libraries

Source: https://github.com/google-deepmind/optax/blob/main/examples/ogda_example.ipynb

Imports required libraries for JAX, Optax, and plotting. Ensure these are installed before running.

```python
import functools
import jax
import optax
import matplotlib.pyplot as plt
from jax import lax, numpy as jnp
```

--------------------------------

### Import necessary libraries

Source: https://github.com/google-deepmind/optax/blob/main/examples/meta_learning.ipynb

Imports required libraries for JAX, Optax, and plotting. Ensure these are installed before running.

```python
from typing import Iterator, Tuple
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
```

--------------------------------

### Import necessary libraries

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/meta_learning.md

Imports essential libraries for JAX, NumPy, and Optax. Ensure these are installed before running.

```python
from typing import Iterator, Tuplenimport jaxnimport jax.numpy as jnpnimport matplotlib.pyplot as pltnimport numpy as npnimport optax
```

--------------------------------

### Training Loop Setup

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/lookahead_mnist.md

Sets up variables to track training accuracy and losses, and computes initial test set statistics before the training loop begins. This snippet is a precursor to the main training loop.

```python
train_accuracy = []
train_losses = []

# Computes test set accuracy at initialization.
test_stats = dataset_stats(params.slow, net, test_loader_batched)
test_accuracy = [test_stats["accuracy"]]
test_losses = [test_stats["loss"]]


@jax.jit
```

--------------------------------

### Learning Rate Scheduling

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Applies a learning rate schedule. This example uses a cosine decay schedule.

```python
schedule = optax.cosine_decay_schedule(init_value=0.001, decay_steps=100000, alpha=0.1)
tx = optax.adam(learning_rate=schedule)
```

--------------------------------

### Gradient Transformation with State Initialization

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/ogda_example.md

Demonstrates how to initialize state for a gradient transformation that requires it.

```python
import optax

# Use RMSprop optimizer, which requires state initialization
optimizer = optax.rmsprop(learning_rate=0.001)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3.])}

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

# Now you can use the optimizer for updates
grads = {'w': jnp.array([0.1, 0.2]), 'b': jnp.array([0.3])}
updates, opt_state = optimizer.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Transformation with Different Decay Values

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

This example demonstrates using different decay values for RMSprop.

```python
import optax

# Using different decay values for RMSprop
optimizer_rmsprop_decay1 = optax.rmsprop(learning_rate=0.001, decay=0.9)
optimizer_rmsprop_decay2 = optax.rmsprop(learning_rate=0.001, decay=0.99)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3., 4.])}
opt_state1 = optimizer_rmsprop_decay1.init(params)
opt_state2 = optimizer_rmsprop_decay2.init(params)

grads = {'w': jnp.array([0.1, 0.2]), 'b': jnp.array([0.3, 0.4])}

updates1, opt_state1 = optimizer_rmsprop_decay1.update(grads, opt_state1, params)
updates2, opt_state2 = optimizer_rmsprop_decay2.update(grads, opt_state2, params)
```

--------------------------------

### Gradient Transformation with Different Inputs

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

This example demonstrates how a gradient transformation can handle different input gradients.

```python
import optax

optimizer = optax.adam(learning_rate=0.001)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3., 4.])}
opt_state = optimizer.init(params)

grads1 = {'w': jnp.array([0.1, 0.2]), 'b': jnp.array([0.3, 0.4])}
grads2 = {'w': jnp.array([0.5, 0.6]), 'b': jnp.array([0.7, 0.8])}

updates1, opt_state1 = optimizer.update(grads1, opt_state, params)
updates2, opt_state2 = optimizer.update(grads2, opt_state1, params)

new_params1 = optax.apply_updates(params, updates1)
new_params2 = optax.apply_updates(params, updates2)
```

--------------------------------

### Gradient Transformation with State (Zero Initialization)

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates how to use zero initialization with an optimizer.

```python
optimizer = optax.chain(optax.sgd(learning_rate=0.1), optax.zero_init())
state = optimizer.init(params)

# Inside training loop:
grads = ...
updates, state = optimizer.update(grads, state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Clipping by Value

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

This example shows how to clip gradients by their individual values, rather than by their global norm.

```python
import optax

optimizer = optax.adam(learning_rate=0.001, clip_by_value_clip_threshold=1.0)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3., 4.])}
opt_state = optimizer.init(params)

grads = {'w': jnp.array([0.1, 5.0]), 'b': jnp.array([0.3, 0.4])}

updates, opt_state = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Transformation with State

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

This example shows a gradient transformation that maintains state, such as momentum in the SGD optimizer.

```python
import optax

optimizer = optax.sgd(learning_rate=0.01, momentum=0.9)

params = {'w': jnp.array([1., 2.]), 'b': jnp.array([3., 4.])}
opt_state = optimizer.init(params)

grads = {'w': jnp.array([0.1, 0.2]), 'b': jnp.array([0.3, 0.4])}

updates, opt_state = optimizer.update(grads, opt_state, params)
new_params = optax.apply_updates(params, updates)
```

--------------------------------

### Import necessary libraries

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/perturbations.md

Imports JAX, JAX NumPy, Optax Treen, and Optax Perturbations for use in the notebook.

```python
import jax
import jax.numpy as jnp
from jax import tree_util as jtu

import optax.treen
from optax import perturbations
```

--------------------------------

### Gradient Transformation with State (Lion with State)

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/adversarial_training.md

Demonstrates how to use Lion with explicit state management.

```python
optimizer = optax.chain(optax.lion(learning_rate=1e-4, weight_decay=0.01))
state = optimizer.init(params)

# Inside training loop:
grads = ...
updates, state = optimizer.update(grads, state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### RMSprop Optimizer

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Example of using the RMSprop optimizer, another adaptive learning rate optimization algorithm.

```python
import optax

# Initialize the RMSprop optimizer
optimizer = optax.rmsprop(learning_rate=0.001)

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

# Update parameters and optimizer state
updates, state = optimizer.update(grads, state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Adam Optimizer

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Example of using the Adam optimizer, a popular adaptive learning rate optimization algorithm.

```python
import optax

# Initialize the Adam optimizer
optimizer = optax.adam(learning_rate=0.001)

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

# Update parameters and optimizer state
updates, state = optimizer.update(grads, state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Demonstrate data sampling

Source: https://github.com/google-deepmind/optax/blob/main/examples/meta_learning.ipynb

Shows how to sample data points from the defined generator. This is useful for verifying the generator's output.

```python
g = generator()

for _ in range(5):
  x, y = next(g)
  print(f"Sampled y = {y:.3f}, x = {x:.3f}")
```

--------------------------------

### Gradient Scaling

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/reduce_on_plateau.md

This example demonstrates gradient scaling, which can be useful for mixed-precision training to prevent underflow of gradients.

```python
import optax

# Define gradient scaling
scaler = optax.scale_by_adam()

# Combine with other transformations
optimizer = optax.chain(
    optax.scale_by_float64_dtype(),
    optax.clip_by_global_norm(1.0),
    scaler,
    optax.scale(0.01) # learning rate
)

# Initialize and update (similar to the basic example)
```

--------------------------------

### Initialize Model and Optimizer

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/gradient_accumulation_and_microbatching.md

Sets up model configuration, initializes a Transformer model, generates dummy data, and initializes an AdamW optimizer state. This prepares the components for microbatching demonstrations.

```python
hidden_size = 512
num_heads = 8
num_layers = 12
vocab_size = 10000
data_size = 8192
sequence_length = 256
batch_size = 32
accumulation_steps = 16

model = Transformer(vocab_size, num_layers, hidden_size, num_heads, rngs=nnx.Rngs(0))
key = jax.random.key(0)
batch = jax.random.randint(key, (batch_size, sequence_length), 0, vocab_size)

graphdef, params = nnx.split(model, nnx.Variable)
print(optax.tree.size(params))

adamw = optax.adamw(0.01)
opt_state = adamw.init(params)
```

--------------------------------

### Initialize Network Parameters and Optimizer

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/gradient_accumulation.md

Generates random image-like data and initializes network parameters and a base optimizer for testing.

```python
EXAMPLES = jax.random.uniform(jax.random.PRNGKey(0), (9, 28, 28, 1))
LABELS = jax.random.randint(jax.random.PRNGKey(0), (9,), minval=0, maxval=10)

optimizer = optax.sgd(1e-4)
params = net.init(jax.random.PRNGKey(0), EXAMPLES)
```

--------------------------------

### Learning Rate Scheduling

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Applies a learning rate schedule. This example uses a linear decay schedule.

```python
import optax

num_train_steps = 1000

cosine_decay_scheduler = optax.cosine_decay_schedule(
    init_value=0.001,
    decay_steps=num_train_steps,
    alpha=0.1
)

optimizer = optax.adam(learning_rate=cosine_decay_scheduler)
```

--------------------------------

### Optimizer State Update Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/freezing_parameters.md

Updates the optimizer state after applying updates. Crucial for adaptive optimizers.

```python
import optax

optimizer = optax.adam(learning_rate=0.001)
state = optimizer.init(params)
updates, state = optimizer.update(grads, state, params)
```

--------------------------------

### Basic Optimizer Usage

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/contrib/rosenbrock_ademamix.md

Demonstrates the basic usage of an optimizer in Optax. This involves creating an optimizer, initializing its state, and applying updates to model parameters.

```python
import optax

# Create an optimizer
optimizer = optax.adam(learning_rate=0.001)

# Initialize optimizer state
params = ... # Your model parameters
opt_state = optimizer.init(params)

# Compute gradients (e.g., during training)
grads = ...

# Apply updates to parameters
updates, opt_state = optimizer.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
```

--------------------------------

### Gradient Clipping Example

Source: https://github.com/google-deepmind/optax/blob/main/docs/_collections/examples/linear_assignment_problem.md

Demonstrates how to clip gradients by their global norm. This helps prevent exploding gradients.

```python
import optax

# Create a gradient clipper
clipper = optax.clip_by_global_norm(1.0)

# Use this clipper within an optax.chain or directly if needed
```

--------------------------------

### Import necessary libraries

Source: https://github.com/google-deepmind/optax/blob/main/examples/gradient_accumulation.ipynb

Imports required libraries for Flax, JAX, NumPy, and Optax. Ensure these are installed before running.

```python
from typing import Iterable

import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
import optax

```