JAX 101

These tutorials cover basic usage of JAX and its features, including some of the internal mechanisms that make JAX work. They start with the fundamentals and are meant to be read sequentially. For more in-depth discussions of JAX’s design and implementation, in no particular order, see the advanced guides.

• Just-in-time compilation
• Automatic vectorization
• Automatic differentiation
• Pytrees
• Pseudorandom numbers
• Introduction to parallel programming
• Control flow and logical operators with JIT
• Tracing
• Stateful computations