JAX Core Concepts

import jax.numpy as jnp

# Create arrays (just like NumPy!)
x = jnp.array([1.0, 2.0, 3.0])
zeros = jnp.zeros((3, 4))
rand = jax.random.normal(jax.random.PRNGKey(0), (3, 3))

# NumPy-like operations
y = jnp.dot(x, x)
z = jnp.sin(x) + jnp.cos(x)

# Immutable! Use .at[] for "updates"
# x[0] = 5  # ERROR! JAX arrays are immutable
x_new = x.at[0].set(5.0)  # Creates a new array

import jax

# Method 1: Decorator
@jax.jit
def fast_computation(x, y):
    return jnp.dot(x, y) + jnp.sum(x)

# Method 2: Wrap function
def slow_fn(x):
    return jnp.sum(jnp.sin(x) ** 2)

fast_fn = jax.jit(slow_fn)

# First call: compile + run (slower)
result = fast_fn(jnp.ones(1000000))

# Subsequent calls: just run (much faster!)
result = fast_fn(jnp.ones(1000000))

import jax

# Gradient of a scalar function
def loss_fn(w, x, y):
    predictions = jnp.dot(x, w)
    return jnp.mean((predictions - y) ** 2)

# Compute gradient with respect to first argument (w)
grad_fn = jax.grad(loss_fn)
grads = grad_fn(w, x, y)

# Value and gradient together (more efficient)
loss, grads = jax.value_and_grad(loss_fn)(w, x, y)

# Higher-order derivatives
def f(x):
    return x ** 3

df = jax.grad(f)        # 3x^2
d2f = jax.grad(df)      # 6x
d3f = jax.grad(d2f)     # 6

print(d2f(2.0))  # 12.0 (6 * 2)

import jax

# Function for a single data point
def predict_single(params, x):
    return jnp.dot(params['w'], x) + params['b']

# Automatically vectorize over a batch of inputs
predict_batch = jax.vmap(predict_single, in_axes=(None, 0))
# None = don't vectorize params, 0 = vectorize x along first axis

# Works on entire batches!
params = {'w': jnp.ones(10), 'b': 0.0}
batch_x = jnp.ones((32, 10))  # 32 samples
predictions = predict_batch(params, batch_x)  # (32,)

# Compose with other transforms
fast_batch_predict = jax.jit(jax.vmap(predict_single, in_axes=(None, 0)))

# Per-sample gradients (impossible in PyTorch without tricks!)
per_sample_grads = jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0))