rocm_jax/notebooks

Notebooks

Use the links below to open any of these for interactive exploration in colab.

• MAML - pedagogical demonstration of Model-Agnostic Meta-Learning in JAX.

• Score Matching - demonstration of Generative Modeling by Estimating Gradients of the Data Distribution

• vmapped log-probabilities - demonstrates the utility of vmap for Bayesian inference.

• 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.

Some additional notebooks (quickstart, auto-diff cookbook and commona gotchas) were recently moved into JAX's online documentation.