Advanced Tutorials

This section contains examples and tutorials on more advanced topics, such as Multi Core computation, Custom operations, and more in depth applications

Examples

• Training a Simple Neural Network, with tensorflow/datasets Data Loading
• Training a Simple Neural Network, with PyTorch Data Loading
• Autobatching for Bayesian Inference

Parallel Computation

• Using JAX in multi-host and multi-process environments
• Distributed arrays and automatic parallelization
• Named axes and easy-to-revise parallelism with xmap

Automatic Differentiation

• The Autodiff Cookbook
• Custom derivative rules for JAX-transformable Python functions
• Control autodiff’s saved values with jax.checkpoint (aka jax.remat)

JAX Internals

• How JAX primitives work
• Writing custom Jaxpr interpreters in JAX
• Custom operations for GPUs with C++ and CUDA

Deep Dives

• Generalized Convolutions in JAX