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.

Tutorials

If you’re ready to explore JAX more deeply, the JAX 101 tutorial goes into much more detail:

• Tutorial: JAX 101
  • JAX As Accelerated NumPy
  • Just In Time Compilation with JAX
  • Automatic Vectorization in JAX
  • Advanced Automatic Differentiation in JAX
  • Pseudo Random Numbers in JAX
  • Working with Pytrees
  • Parallel Evaluation in JAX
  • Stateful Computations in JAX

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.

Finding Help

If you have questions about JAX, we’d love to answer them! Two good places to get your questions answered are:

• JAX GitHub discussions

• JAX on StackOverflow