The basics: interactive NumPy on GPU and TPU¶
In [ ]:
import jax
import jax.numpy as jnp
from jax import random
In [ ]:
key = random.PRNGKey(0)
key, subkey = random.split(key)
x = random.normal(key, (5000, 5000))
print(x.shape)
print(x.dtype)
In [ ]:
y = jnp.dot(x, x)
print(y[0, 0])
In [ ]:
x
In [ ]:
import matplotlib.pyplot as plt
plt.plot(x[0])
In [ ]:
print(jnp.dot(x, x.T))
In [ ]:
print(jnp.dot(x, 2 * x)[[0, 2, 1, 0], ..., None, ::-1])
In [ ]:
import numpy as np
x_cpu = np.array(x)
%timeit -n 5 -r 2 np.dot(x_cpu, x_cpu)
In [ ]:
%timeit -n 5 -r 5 jnp.dot(x, x).block_until_ready()
Automatic differentiation¶
In [ ]:
from jax import grad
In [ ]:
def f(x):
if x > 0:
return 2 * x ** 3
else:
return 3 * x
In [ ]:
key = random.PRNGKey(0)
x = random.normal(key, ())
print(grad(f)(x))
print(grad(f)(-x))
In [ ]:
print(grad(grad(f))(-x))
print(grad(grad(grad(f)))(-x))
Other JAX autodiff highlights:
- Forward- and reverse-mode, totally composable
- Fast Jacobians and Hessians
- Complex number support (holomorphic and non-holomorphic)
- Jacobian pre-accumulation for elementwise operations (like
gelu)
For much more, see the JAX Autodiff Cookbook (Part 1).
End-to-end compilation with XLA with jit¶
In [ ]:
from jax import jit
In [ ]:
key = random.PRNGKey(0)
x = random.normal(key, (5000, 5000))
In [ ]:
def f(x):
y = x
for _ in range(10):
y = y - 0.1 * y + 3.
return y[:100, :100]
f(x)
In [ ]:
g = jit(f)
g(x)
In [ ]:
%timeit f(x).block_until_ready()
In [ ]:
%timeit -n 100 g(x).block_until_ready()
In [ ]:
grad(jit(grad(jit(grad(jnp.tanh)))))(1.0)
Parallelization over multiple accelerators with pmap¶
In [ ]:
jax.device_count()
In [ ]:
from jax import pmap
In [ ]:
y = pmap(lambda x: x ** 2)(jnp.arange(8))
print(y)
In [ ]:
y
In [ ]:
import matplotlib.pyplot as plt
plt.plot(y)
Collective communication operations¶
In [ ]:
from functools import partial
from jax.lax import psum
@partial(pmap, axis_name='i')
def f(x):
total = psum(x, 'i')
return x / total, total
normalized, total = f(jnp.arange(8.))
print(f"normalized:\n{normalized}\n")
print("total:", total)
For more, see the pmap cookbook.
Automatic parallelization with sharded_jit (new!)¶
In [ ]:
from jax.experimental import sharded_jit, PartitionSpec as P
In [ ]:
from jax import lax
conv = lambda image, kernel: lax.conv(image, kernel, (1, 1), 'SAME')
In [ ]:
image = jnp.ones((1, 8, 2000, 1000)).astype(np.float32)
kernel = jnp.array(np.random.random((8, 8, 5, 5)).astype(np.float32))
np.set_printoptions(edgeitems=1)
conv(image, kernel)
In [ ]:
%timeit conv(image, kernel).block_until_ready()
In [ ]:
image_partitions = P(1, 1, 4, 2)
sharded_conv = sharded_jit(conv,
in_parts=(image_partitions, None),
out_parts=image_partitions)
sharded_conv(image, kernel)
In [ ]:
%timeit -n 10 sharded_conv(image, kernel).block_until_ready()