JAX is a high-performance research and training framework that combines a NumPy-style API with automatic differentiation, JIT compilation, vectorization, and distributed execution.
JAX Web IDE
Run code in the browser with nothing to install and learn through step-by-step examples.
Open the Web IDE →JAX is a high-performance research and training framework that combines a NumPy-style API with automatic differentiation, JIT compilation, vectorization, and distributed execution. Rather than listing concepts, this guide is organized so you can follow the content in the order you have to make decisions on a real project.
Look beyond the model and the prompt: connect the data flow, evaluation, and post-deployment operating metrics in one view.
To make what you read stick, it helps to capture the flow as a picture first. The two diagrams below are reference maps you can keep coming back to while learning JAX.
When you first open JAX, the big picture matters more than individual commands. This section starts with the problems the concepts ahead were created to solve.
| Function | Description |
|---|---|
| grad | Computes the gradient of a function automatically. |
| jit | Compiles a function with XLA to run it fast. |
| vmap | Vectorizes over the batch dimension automatically. |
| pmap | Runs in parallel across several devices. |
Here you look at Installation alongside real code. Rather than copying the example as-is, read it with an eye on the inputs, the outputs, and the parts most likely to change.
uv venv
source .venv/bin/activate
uv pip install "jax[cpu]"
python -c "import jax; print(jax.devices())"Here you look at Automatic differentiation alongside real code. Rather than copying the example as-is, read it with an eye on the inputs, the outputs, and the parts most likely to change.
import jax
import jax.numpy as jnp
def loss(w):
return jnp.sum((w - 3.0) ** 2)
grad_loss = jax.grad(loss)
print(grad_loss(jnp.array([1.0, 2.0, 4.0])))Here you look at JIT compilation alongside real code. Rather than copying the example as-is, read it with an eye on the inputs, the outputs, and the parts most likely to change.
@jax.jit
def matmul(a, b):
return a @ b
x = jnp.ones((1024, 1024))
print(matmul(x, x).shape)Here you look at Vectorization alongside real code. Rather than copying the example as-is, read it with an eye on the inputs, the outputs, and the parts most likely to change.
def predict(w, x):
return jnp.dot(w, x)
batched_predict = jax.vmap(predict, in_axes=(None, 0))JAX practical design is a point where the options diverge. Using the table to compare what each approach is for and how it differs in operation makes later decisions much easier.
| Decision point | Question to ask | Practical guideline |
|---|---|---|
| Boundaries | Which parts of the JAX code are most likely to change? | Separate input/output, configuration, external integrations, and the core rules. |
| State | Where is state created and where does it go away? | Make the owner and lifecycle of state visible in the code. |
| Failure | What does the caller receive on failure? | Decide the timeout, fallback, and error contract first. |
This section covers JAX operating standards from a practical perspective. Instead of memorizing the concept, focus on the situations in which you would reach for it.
JAX verification strategy is a point where the options diverge. Using the table to compare what each approach is for and how it differs in operation makes later decisions much easier.
| Quality axis | How to verify | Definition of done |
|---|---|---|
| Correctness | Automate the normal/failure cases. | The core scenarios pass reproducibly. |
| Regression prevention | When fixing a bug, leave the same case behind as a test. | The same failure is not deployed again. |
| Operability | Check the logs, metrics, and alerts. | There is a path for tracing the cause when a problem occurs. |