🖼️ 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 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.

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 learing rate scheduler.

Sharpness-Aware Minimization (SAM)

Sharpness-Aware Minimization (SAM).