🖼️ Example gallery

Gradient Accumulation

Gradient Accumulation.

Meta-Learning

Meta-Learning.

Optimistic Gradient Descent in a Bilinear Min-Max Problem

Optimistic Gradient Descent in a Bilinear Min-Max Problem.

ResNet on CIFAR10 with Flax NNX and Optax.

ResNet on CIFAR10 with Flax and Optax.

Simple NN with Flax.

Train the parameters of a Flax module.

Lookahead Optimizer on MNIST

Adam optimizer and lookahead wrapper on the MNIST dataset.

MLP MNIST

Train an MLP classifier on MNIST using Optax.

Adversarial training

Adversarial training of CNN on MNIST.

Character-level Transformer on Tiny Shakespeare

Character-level Transformer on Tiny Shakespeare.

L-BFGS

Using LBFGS and linesearch.

Perturbed optimizers

Differentiable functions with perturbations.

Linear assignment problem

Solving the linear assignment problem.

Contrib Examples

Examples that make use of the 🔧 Contrib module.

Differentially private convolutional neural network on MNIST.

Differentially private convolutional neural network on MNIST.

Reduce on Plateau Learning Rate Scheduler

Example usage of reduce_on_plateau learning rate scheduler.

Sharpness-Aware Minimization (SAM)

Sharpness-Aware Minimization (SAM).

Recreate AdeMAMix Rosenbrock Plot from Paper

AdEMAMix.