jax

JAX

JAX combines composable function transformations with automatic differentiation (Autograd) and Accelerated Linear Algebra (XLA) compilation for high-throughput machine learning research.

When to Use

• High-Performance Numerical & Scientific Computing: Autograd and XLA compilation for accelerated linear algebra on GPU/TPU.
• Custom Deep Learning Research: Building neural networks with Flax, Haiku, or Equinox without framework bloat.
• Composable Function Transformations: Seamlessly chaining jit (compilation), grad (derivatives), and vmap (vectorization).
• Multi-Device Distributed Training: Parallelizing computation across clusters using pmap and shard maps (jax.experimental.shard_map).

Quick Start

import jax
import jax.numpy as jnp

# JIT-compiled matrix multiplication and automatic differentiation
@jax.jit
def loss_fn(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

# Compute value and gradients simultaneously
grad_fn = jax.grad(loss_fn)

w = jnp.array([1.0, 2.0])
x = jnp.array([[1.0, 0.5], [2.0, 1.0]])
y = jnp.array([2.0, 4.0])

grads = grad_fn(w, x, y)
print("Gradients:", grads)

Core Concepts

Composable Function Transformations: jit, grad & vmap

Combining just-in-time compilation, automatic differentiation, and automated batching:

import jax
import jax.numpy as jnp

# Pure function definition
def loss_fn(w, x, y):
    pred = jnp.dot(x, w)
    return jnp.mean((pred - y) ** 2)

# Transform 1: Automatic gradient with respect to weights
grad_fn = jax.grad(loss_fn)

# Transform 2: JIT compilation via XLA for extreme GPU speedup
fast_grad_fn = jax.jit(grad_fn)

# Transform 3: Automatic vectorization across batches
# batch_loss_fn handles 2D matrices automatically without loops
vmapped_loss = jax.vmap(loss_fn, in_axes=(None, 0, 0))

# Sample computation
key = jax.random.PRNGKey(42)
w = jax.random.normal(key, (3,))
x = jnp.array([[1.0, 2.0, 3.0]])
y = jnp.array([5.0])

grad_val = fast_grad_fn(w, x, y)
print("Computed Gradient:", grad_val)