Notebooks
Use the links below to open any of these for interactive exploration in colab.
-
Quick Start - the first notebook to go through, explores the basic JAX API.
-
Common Gotchas in JAX - answers for the most common problems people have while getting used to JAX's way of doing things.
-
MAML - pedagogical demonstration of Model-Agnostic Meta-Learning in JAX.
-
vmapped log-probabilities - demonstrates the utility of vmap for Bayesian inference.
-
gufuncs via vmap - how to implement NumPy-like gufuncs using vmap.
-
Neural Networks with TFDS Data - training a simple neural net with tensorflow datasets.
-
Neural Networks and Data Loading - training a simple neural net using a pytorch dataloader.
-
XLA in Python - interactive exploration of the XLA compiler and computation model in python.