jax_quickstart_matrix_multiply_jit_grad_vmap.py

1import jax.numpy as jnp
2from jax import grad, jit, vmap
3from jax import random
4
5# Create a random key for JAX's PRNG
6key = random.PRNGKey(0)
7
8# Generate some random data
9x = random.normal(key, (3000, 3000))
10
11# 1. Basic matrix multiplication (runs on GPU/TPU if available)
12result = jnp.dot(x, x.T)
13print(f"Matrix multiplication result shape: {result.shape}")
14
15# 2. Just-In-Time (JIT) Compilation
16def selu(x, alpha=1.67, lmbda=1.05):
17    return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
18
19# Compile the function
20selu_jit = jit(selu)
21# Execute on the device
22selu_result = selu_jit(x)
23print("JIT compilation successful.")
24
25# 3. Automatic Differentiation (Grad)
26def sum_logistic(x):
27    return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))
28
29# Compute the gradient of the function
30x_small = jnp.arange(3.)
31derivative_fn = grad(sum_logistic)
32print(f"Gradient at {x_small}: {derivative_fn(x_small)}")
33
34# 4. Auto-vectorization (Vmap)
35mat = random.normal(key, (150, 100))
36batched_x = random.normal(key, (10, 100))
37
38def apply_matrix(v):
39    return jnp.dot(mat, v)
40
41# Vectorize the function over a batch
42batch_apply_matrix = vmap(apply_matrix)
43print(f"Vmap output shape: {batch_apply_matrix(batched_x).shape}")