Annotated JAX

Deep learning examples using the JAX ecosystem of libraries. Each page is a notebook from the repository, rendered with its saved outputs. Setup instructions are in the README.

• Basic Matrix
• Differentiation
• Diff On Vectors
• Random Numbers And Arrays
• Jit
• Vmap
• Autodiff And Linear Regression
• Ann Optax Flax
• Lin Reg Optax
• Mini Batches Dataloading
• Pos Tag Lstm
• Cnn Cifar10
• Model Checkpointing