Name		Name	Last commit message	Last commit date
parent directory ..
Common_Gotchas_in_JAX.ipynb		Common_Gotchas_in_JAX.ipynb
README.md		README.md
XLA_in_Python.ipynb		XLA_in_Python.ipynb
autodiff_cookbook.ipynb		autodiff_cookbook.ipynb
gufuncs.ipynb		gufuncs.ipynb
maml.ipynb		maml.ipynb
neural_network_and_data_loading.ipynb		neural_network_and_data_loading.ipynb
neural_network_with_tfds_data.ipynb		neural_network_with_tfds_data.ipynb
quickstart.ipynb		quickstart.ipynb
vmapped log-probs.ipynb		vmapped log-probs.ipynb

README.md

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.