jax_quickstart_grad_jit_vmap_matrix_operations.py

1import jax.numpy as jnp
2from jax import grad, jit, vmap
3from jax import random
4
5# 1. Multiply Matrices
6key = random.PRNGKey(0)
7x = random.normal(key, (1000, 1000))
8# JAX operations are transparently accelerated on GPU/TPU
9result = jnp.dot(x, x.T)
10print(f"Matrix multiplication result shape: {result.shape}")
11
12# 2. Automatic Differentiation with grad
13def sum_logistic(x):
14    return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))
15
16x_small = jnp.arange(3.)
17derivative_fn = grad(sum_logistic)
18print(f"Gradient at {x_small}: {derivative_fn(x_small)}")
19
20# 3. Compilation with jit
21def selu(x, alpha=1.67, lmbda=1.05):
22    return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
23
24# Create a just-in-time compiled version of the function
25selu_jit = jit(selu)
26x = random.normal(key, (1000000,))
27# Warm up and run
28%timeit selu_jit(x).block_until_ready()
29
30# 4. Auto-vectorization with vmap
31mat = random.normal(key, (150, 100))
32batched_x = random.normal(key, (10, 100))
33
34def apply_matrix(v):
35    return jnp.dot(mat, v)
36
37# vmap allows us to apply a function over a batch dimension automatically
38def batch_apply_matrix(batched_v):
39    return vmap(apply_matrix)(batched_v)
40
41print(f"Batched output shape: {batch_apply_matrix(batched_x).shape}")