TopolearnEvery JAX program starts with jnp arrays; knowing how they differ from NumPy avoids most early surprises.
JAX gives you jax.numpy (imported as jnp), an API that mirrors NumPy: jnp.arange, jnp.dot, jnp.mean, broadcasting and slicing all work as you expect. The result type is jax.Array, which lives on a device (CPU, GPU or TPU) rather than in ordinary host memory.
Differences that matter from day one:
- Default precision is 32-bit.
jnp.array([1.0])isfloat32and integers areint32. Enable 64-bit withjax.config.update("jax_enable_x64", True)at program start if you truly need it. - Arrays are immutable; you update them functionally (next lesson).
- Dispatch is asynchronous. An operation returns immediately while the device computes; the value is materialized when you print it or convert it with
np.asarray. - Out-of-bounds indexing does not raise. Reads are clamped to the last element, because device code cannot throw Python exceptions.
jax.devices() lists the accelerators JAX found. Arrays are placed on the first device by default; jax.device_put(x, device) moves them. Converting back with np.asarray(x) copies to the host.
Install with pip install jax for CPU, or pip install "jax[cuda13]" (or the variant the install guide lists for your CUDA version) for NVIDIA GPUs. The same code runs unchanged on any backend.
Example
import jax
import jax.numpy as jnp
import numpy as np
x = jnp.arange(5)
print(x, x.dtype) # [0 1 2 3 4] int32
y = jnp.linspace(0, 1, 3)
print(y.dtype) # float32 (NumPy would give float64)
print(jnp.dot(x, x)) # 30
print(type(x).__name__) # ArrayImpl (an instance of jax.Array)
print(isinstance(x, jax.Array)) # True
print(jax.devices()) # [CpuDevice(id=0)] on a CPU-only install
print(x[10]) # 4: out-of-bounds reads are clamped
host = np.asarray(x) # copy back to a NumPy array
print(type(host).__name__) # ndarrayRemember
jnpmirrors the NumPy API but returns device-residentjax.Arrays- JAX defaults to float32/int32 unless x64 is enabled
- Out-of-bounds indexing clamps instead of raising
Watch out for
- Comparing JAX and NumPy results bit-for-bit and blaming JAX for float32 rounding differences.