Examples

This directory contains examples using the optax library.

• Adversarial training
• ResNet on CIFAR10 with Flax NNX and Optax.
• Simple NN with Flax.
• Freezing Parameters in Optax
• From-scratch Implementation
• Gradient Accumulation
• Summary
• Setup
• Part 1: Microbatching for Gradient Accumulation
• Part 2: microbatching.micro_vmap
• Part 3: microbatching.micro_grad
• L-BFGS
• Linear assignment problem
• Optimal transport
• Lookahead Optimizer on MNIST
• Meta-Learning
• MLP MNIST
• Character-level Transformer on Tiny Shakespeare
• Hyperparameters and dataset download
• Data preparation
• NanoLM Model Definition
• State, Optimizer, and Loss Definition
• Model training
• Text generation
• Optimistic Gradient Descent in a Bilinear Min-Max Problem
• Perturbed optimizers
• Argmax one-hot
• Differentiable ranking
• General input / outputs (Pytrees)