Skip to main content

Applications

You've learned the fundamentals, seen the canonical architectures, and scaled a training run. This section is where it all pays off: small, self-contained implementations of real model families you can build with Flax NNX — each one runnable, each one teaching a distinct idea.

Every guide here follows the same recipe: the math in a few lines, a from-scratch model in idiomatic NNX, a training step, and results you can reproduce. Datasets are tiny (MNIST, synthetic, or built-in) so everything runs on a laptop.

The domains​

🎨 Generative Models​

Learn to generate data, not just classify it: autoencoders, VAEs, GANs, and diffusion models. Showcases nnx.ConvTranspose for learned upsampling.

🔁 Sequence Models & Time Series​

Recurrence and order: RNNs, LSTMs, and GRUs via the nnx.RNN API family, plus sequence-to-sequence and forecasting.

🔬 Graphs, Scientific & Structured​

Beyond grids and sequences: graph neural networks, physics-informed networks (PINNs) that differentiate through the model, Neural ODEs, tabular data, and mixture-of-experts.

🖼️ Advanced Vision​

Vision beyond CNNs: the Vision Transformer (attention on image patches) and the U-Net for dense per-pixel segmentation.

🧬 Multimodal & Adaptation​

Adapt and combine models: parameter-efficient fine-tuning with LoRA (nnx.LoRALinear) and cross-modal alignment.

Prerequisites​

These guides assume you're comfortable with the training loop and Flax NNX state. Individual guides link the specific extra background they need.

How to use this section​

Each sub-category has a suggested reading order that builds one idea at a time (e.g. Autoencoder → VAE → GAN → Diffusion). Jump straight to a domain that interests you, or work through a whole track.

Next steps​