JAX GitHub Repository: https://github.com/google/jax
JAX Quickstart Notebook: https://jax.readthedocs.io/en/latest/notebooks/quickstart.html
Using JAX to Accelerate Our Research by DeepMind: https://deepmind.google/discover/blog/using-jax-to-accelerate-our-research/
Kaggle Learn Guide: https://www.kaggle.com/learn-guide/jax
Tensorflow JAX video: https://www.youtube.com/watch?v=WdTeDXsOSj4&ab_channel=TensorFlow
Tensorflow Probability on JAX: TensorFlow Probability on JAX Guide
Oryx Probabilistic Programming: A Tour of Oryx Guide
Oryx Repo: https://github.com/jax-ml/oryx
NumPyro: https://github.com/pyro-ppl/numpyro
An Astronomer's Introduction to NumPyro: https://dfm.io/posts/intro-to-numpyro
DeepMind Haiku GitHub Repository: https://github.com/google-deepmind/dm-haiku
Optax by DeepMind: https://github.com/google-deepmind/optax
Trax by Google Brain: https://github.com/google/trax
Google Flax GitHub Repository: https://github.com/google/flax
Flax Documentation: https://flax.readthedocs.io/en/latest/
Flax Tutorial on Hugging Face: https://huggingface.co/blog/afmck/flax-tutorial
Intro to JAX & Flax: https://phlippe.github.io/media/GDE_Talk_Intro_to_JAX_Flax_2022_12_06.pdf
Jaxtronomy: https://github.com/google-research/google-research/tree/master/jaxstronomy