Getting Started with JAX#
Welcome to JAX! The JAX documentation contains a number of useful resources for getting started. JAX Quickstart is the easiest place to jump-in and get an overview of the JAX project.
If you’re accustomed to writing NumPy code and are starting to explore JAX, you might find the following resources helpful:
How to Think in JAX is a conceptual walkthrough of JAX’s execution model.
🔪 JAX - The Sharp Bits 🔪 lists some of JAX’s sharp corners.
JAX Frequently Asked Questions (FAQ) answers some frequent jax questions.
If you’re ready to explore JAX more deeply, the JAX 101 tutorial goes into much more detail:
If you prefer a video introduction here is one from JAX contributor Jake VanderPlas:
Building on JAX#
JAX provides the core numerical computing primitives for a number of tools developed by the larger community. For example, if you’re interested in using JAX for training neural networks, two well-supported options are Flax and Haiku.
For a community-curated list of JAX-related projects across a wide set of domains, check out Awesome JAX.
If you have questions about JAX, we’d love to answer them! Two good places to get your questions answered are: