jax_quickstart_jit_grad_vmap_array_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, (3000, 3000), dtype=jnp.float32)
8# JAX operations are often asynchronous, but let's do a simple dot product
9result = jnp.dot(x, x.T)
10print(f"Matrix multiplication result shape: {result.shape}")
11
12# 2. Use jit() to speed up functions
13def selu(x, alpha=1.67, lmbda=1.05):
14    return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
15
16x = random.normal(key, (1000000,))
17selu_jit = jit(selu)
18# Warm up and run
19_ = selu_jit(x) 
20print("JIT-compiled SELU executed.")
21
22# 3. Take derivatives with grad()
23def sum_logistic(x):
24    return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))
25
26x_small = jnp.arange(3.)
27derivative_fn = grad(sum_logistic)
28print(f"Gradient at {x_small}: {derivative_fn(x_small)}")
29
30# 4. Auto-vectorization with vmap()
31def apply_matrix(v):
32    return jnp.dot(result, v)
33
34# Batch of 10 vectors
35batch_x = random.normal(key, (10, 3000))
36# vmap(apply_matrix) applies the function to every row in the batch
37batch_results = vmap(apply_matrix)(batch_x)
38print(f"Vectorized batch result shape: {batch_results.shape}")