jax_quickstart_jit_compilation_and_autodiff.py

1import jax.numpy as jnp
2from jax import grad, jit, vmap
3from jax import random
4
5# Generate key for random number generation
6key = random.PRNGKey(0)
7
8# Create a large matrix
9x = random.normal(key, (3000, 3000), dtype=jnp.float32)
10
11# Define a simple function
12def selu(x, alpha=1.67, lmbda=1.05):
13    return lmbda * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)
14
15# Jit-compile the function for speed
16selu_jit = jit(selu)
17
18# Demonstrate differentiation
19def sum_logistic(x):
20    return jnp.sum(1.0 / (1.0 + jnp.exp(-x)))
21
22x_small = jnp.arange(3.)
23derivative_fn = grad(sum_logistic)
24print(derivative_fn(x_small))