Example applications

The following pages provide examples of common applications of the JAX AI stack:

• Train a miniGPT language model with JAX
• Basic text classification with 1D CNN
• Text classification with a transformer language model using JAX
• Machine Translation with encoder-decoder transformer model
• Image segmentation with UNETR model
• Image Captioning with Vision Transformer (ViT) model
• Train a Vision Transformer (ViT) for image classification with JAX
• Time series classification with CNN