jax.numpy: The NumPy-Compatible API

If you know NumPy, you're already 80% of the way to knowing JAX's computation API. The jax.numpy module (conventionally imported as jnp) provides nearly every function from NumPy with the same names, same argument signatures, and same behavior — but with the ability to run on GPUs and TPUs.

import numpy as np
import jax.numpy as jnp

# Creating arrays — identical syntax
np_arr = np.array([1.0, 2.0, 3.0])
jnp_arr = jnp.array([1.0, 2.0, 3.0])

# Zeros, ones, ranges — identical
np_zeros = np.zeros((3, 4))
jnp_zeros = jnp.zeros((3, 4))

np_range = np.arange(0, 10, 2)
jnp_range = jnp.arange(0, 10, 2)

np_lin = np.linspace(0, 1, 50)
jnp_lin = jnp.linspace(0, 1, 50)

# Math — identical
print(np.sum(np_arr ** 2))     # 14.0
print(jnp.sum(jnp_arr ** 2))  # 14.0

The coverage is extensive. Here are the most commonly used creation functions:

import jax.numpy as jnp

# From Python data
a = jnp.array([[1, 2], [3, 4]])

# Filled arrays
zeros = jnp.zeros((3, 4))           # 3x4 of zeros
ones = jnp.ones((2, 3))             # 2x3 of ones
full = jnp.full((2, 2), fill_value=7.0)  # 2x2 of 7.0

# Ranges
r = jnp.arange(0, 10, 0.5)          # [0.0, 0.5, 1.0, ..., 9.5]
l = jnp.linspace(0, 1, 5)           # [0.0, 0.25, 0.5, 0.75, 1.0]

# Identity and diagonal
eye = jnp.eye(3)                     # 3x3 identity matrix
diag = jnp.diag(jnp.array([1, 2, 3]))  # 3x3 diagonal matrix

# Random arrays (JAX uses explicit PRNG keys, covered in Track 3)
import jax
key = jax.random.PRNGKey(0)
rand = jax.random.normal(key, shape=(3, 4))  # 3x4 standard normal

Array properties work the same way too:

x = jnp.array([[1.0, 2.0, 3.0],
                [4.0, 5.0, 6.0]])

print(x.shape)   # (2, 3)
print(x.dtype)   # float32
print(x.ndim)    # 2
print(x.size)    # 6