JAX JIT in the browser

Edit the Python below and ask JAX to compile it with XLA. Python, JAX, XLA and the generated program all run locally in this browser tab.

Loading the WebAssembly Python runtime…

Python code

Output

Waiting for Python…

The first compilation takes longer; later runs reuse the same Python kernel.