Toggle navigation

.rst .pdf

Warning

This page was created from a pull request (#9655).

Tutorial: JAX 101

This is a tutorial developed by engineers and researchers at DeepMind.

Tutorials

• 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
• Introduction to pjit

previous

🔪 JAX - The Sharp Bits 🔪

next

JAX As Accelerated NumPy