Need NumPy-style array operations with access to JAX acceleration? Import jax.numpy, commonly abbreviated jnp, and use its mostly NumPy-compatible API for creating and manipulating JAX Array values. In many programs, replacing import numpy as np with import jax.numpy as np is enough to begin using JAX on CPUs, GPUs, or TPUs. section 1.1.1 section 1.2.1
Use jax.numpy with these practical differences in mind:
- Create arrays through functions such as array() and linspace(); scalar results are also represented as JAX Array values. section 3.2.2
- Convert Python lists or tuples explicitly before passing them to operations such as sum(); unlike NumPy, JAX rejects them in many function arguments to avoid silent performance degradation. section 3.2.2
- Replace in-place updates with functional updates because JAX arrays are immutable; for example, use arr = arr.at[i].add(10). section 1.2.1
- Expect some NumPy differences, including incomplete function coverage, distinct type-promotion behavior, asynchronous dispatch, and device-dependent data placement. section 1.1.1 section 3.6
Convert before summing
import jax.numpy as jnp values = jnp.array([1, 42, 31337]) total = jnp.sum(values) # total is Array(31380, dtype=int32) # jnp.sum([1, 42, 31337]) raises TypeError The resulting Array exposes familiar properties such as ndim, shape, dtype, size, and nbytes, making ordinary array inspection feel close to NumPy.
Figure 1.1 A simple (and pretty useless) function with some amount of computations that we will evaluate with pure NumPy, with JAX using a CPU and GPU, and with JAX-compiled versions of the same function for CPU and GPU backends. We will evaluate the time to calculate f(x).
Compare pure NumPy with JAX on CPU and GPU, including JIT-compiled variants, to see where the same array computation can gain acceleration. section 1.1.1
