Despite the API similarity, JAX arrays have some fundamental differences from NumPy arrays. Understanding these differences early will save you from confusing error messages later.
Difference #1: Immutability
This is the biggest difference. JAX arrays cannot be modified in place. NumPy lets you do arr[i] = value, but JAX does not.
import numpy as np
import jax.numpy as jnp
# NumPy: in-place mutation works fine
np_arr = np.array([1, 2, 3])
np_arr[0] = 99
print(np_arr) # [99, 2, 3]
# JAX: in-place mutation raises an error
jnp_arr = jnp.array([1, 2, 3])
# jnp_arr[0] = 99 # ERROR: JAX arrays are immutable
# Instead, use .at[].set() to create a NEW array
new_arr = jnp_arr.at[0].set(99)
print(new_arr) # [99, 2, 3]
print(jnp_arr) # [1, 2, 3] — original unchanged!
The .at[] syntax supports all the same operations NumPy does in place:
import jax.numpy as jnp
x = jnp.array([10, 20, 30, 40, 50])
# Set a value
x_new = x.at[2].set(99) # [10, 20, 99, 40, 50]
# Add to a value
x_add = x.at[2].add(5) # [10, 20, 35, 40, 50]
# Multiply
x_mul = x.at[2].mul(2) # [10, 20, 60, 40, 50]
# Slice updates
x_slice = x.at[1:3].set(0) # [10, 0, 0, 40, 50]
# Conditional update with jnp.where
mask = x > 25
x_where = jnp.where(mask, x * 2, x) # [10, 20, 60, 80, 100]
⚠️ Pure Function Check
Immutability is not a limitation — it's a feature. Because JAX arrays can't be mutated, JAX can safely share array memory across computations, cache compilation results, and reason about data flow for optimization. Every .at[].set() call returns a new array, making data flow explicit and transformation-friendly.
Difference #2: Device Placement
NumPy arrays live in CPU memory. JAX arrays (called jax.Array — formerly known as DeviceArray) live on whatever accelerator is available.
import jax
import jax.numpy as jnp
# Check available devices
print(jax.devices()) # e.g., [CudaDevice(id=0)] or [TpuDevice(id=0)]
print(jax.default_backend()) # 'gpu', 'tpu', or 'cpu'
# Arrays are created on the default device
x = jnp.array([1.0, 2.0, 3.0])
print(x.devices()) # Shows which device(s) the array is on
# Explicitly place on a device
cpu_device = jax.devices('cpu')[0]
x_cpu = jax.device_put(x, cpu_device)
Difference #3: Don't Mix NumPy and JAX NumPy
import numpy as np
import jax.numpy as jnp
x = jnp.array([1.0, 2.0, 3.0])
# This works but BYPASSES JIT — the array is silently moved to CPU
result = np.sum(x) # Uses NumPy, not JAX — no JIT, no GPU
# Always use jnp for JAX arrays
result = jnp.sum(x) # Correct — uses JAX, can be JIT-compiled
💡 Why This Matters
Accidentally using np.something() instead of jnp.something() is one of the most common JAX bugs. Your code will work (NumPy can consume JAX arrays), but it will silently transfer data from GPU to CPU and bypass all JAX optimizations. If your JIT-compiled code is unexpectedly slow, check for stray np. calls.