Arrays and jax.numpy

JAX’s basic data type is the array, jax.Array, and its basic API for working with arrays is jax.numpy, which closely mirrors NumPy. This page covers what carries over from NumPy unchanged, and the handful of differences that matter from day one: immutability, default dtypes, and a few indexing behaviors.

From NumPy to JAX

The conventional way to import jax.numpy is under the alias jnp, alongside regular NumPy as np:

import jax
import jax.numpy as jnp
import numpy as np

With this import, array creation, arithmetic, indexing, reductions, and so on all look just like NumPy:

x = jnp.arange(9.0).reshape(3, 3)
y = jnp.ones(3)

print(x @ y)
print(x.sum(axis=0))
print(x[1, :2])

[ 3. 12. 21.]
[ 9. 12. 15.]
[3. 4.]

The arrays themselves are instances of jax.Array:

isinstance(x, jax.Array)

True

If you use type annotations in your code, jax.Array is the appropriate annotation for JAX array values (see jax.typing for more discussion).

Every JAX value also has a JAX type, which you can inspect with jax.typeof(). For an array, the JAX type roughly means its shape and dtype:

print(jax.typeof(x))

float32[3,3]

This compact notation appears all over JAX, and we’ll meet it again when we look at how JAX records whole programs. JAX types can also carry more than shape and dtype (notably sharding, describing how an array is laid out across devices). But when you think “JAX type” you can often think “shape and dtype”.

JAX arrays and NumPy arrays can often be used interchangeably: jax.numpy functions accept NumPy arrays, and most NumPy functions accept JAX arrays (thanks to Python’s duck typing). It’s common to use np for data loading and lightweight host-side manipulation and jnp for number crunching on accelerators.

But there are some important differences.

JAX arrays are immutable

The most important difference from NumPy: once created, a JAX array’s contents cannot be changed. NumPy allows in-place mutation:

x_np = np.arange(10)
x_np[0] = 10
print(x_np)

[10  1  2  3  4  5  6  7  8  9]

The equivalent raises an error in JAX:

x = jnp.arange(10)
x[0] = 10

TypeError: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html

Instead, JAX provides functional updates via the at property, which return a new array with the update applied:

y = x.at[0].set(10)
print(x)  # unchanged
print(y)

[0 1 2 3 4 5 6 7 8 9]
[10  1  2  3  4  5  6  7  8  9]

Alongside set, there’s a family of update operations, including add, multiply, min, and max:

print(x.at[3].add(100))
print(x.at[:3].max(5))

[  0   1   2 103   4   5   6   7   8   9]
[5 5 5 3 4 5 6 7 8 9]

Writing updates functionally may look wasteful, as if every update copies the whole array. Outside of compiled code that’s accurate, but inside jax.jit()-compiled functions, covered in Just-in-time compilation, the compiler can usually perform these updates in place. Other optimizations are possible too: for example, intermediates might not be materialized at all. In general, functional code is amenable to transformation, optimization, and parallelization.

For when mutable arrays are really necessary, JAX provides a distinct array reference type, covered in Refs: mutable arrays.

Default dtypes and precision

NumPy defaults to 64-bit floating point. JAX defaults to 32-bit, which is usually what you want on accelerators like GPUs and TPUs:

print(np.array([1.0, 2.0]).dtype)
print(jnp.array([1.0, 2.0]).dtype)

float64
float32

In fact, by default JAX disables 64-bit dtypes altogether: requesting float64 produces a float32 array (with a warning). If you need 64-bit precision, you can enable it with a config option:

import jax
jax.config.update("jax_enable_x64", True)

For the full story, including how to control default dtypes more finely, see Default dtypes and the X64 flag.

When operations mix dtypes, JAX applies type promotion rules that are similar to NumPy’s but not identical. In particular they’re designed to avoid accidentally promoting everything to 64-bit:

(jnp.arange(3) + 1.5).dtype

dtype('float32')

See Type promotion semantics for the precise promotion semantics.

Indexing differences

JAX supports NumPy-style indexing, including slices and advanced integer array indexing. But some edge-case behaviors differ, because JAX is designed so that every operation can be compiled for accelerators, where raising an exception from inside a computation isn’t an option.

Most notably, out-of-bounds indexing doesn’t raise an error. When reading, the index is clamped to the bounds of the array:

x = jnp.arange(10)
x[11]  # clamped to x[9]

Array(9, dtype=int32)

When writing with .at, out-of-bounds updates are dropped:

x.at[11].set(99)  # update is dropped

Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)

Both behaviors can be adjusted with the mode argument to indexing operations; see jax.numpy.ndarray.at for the options.

Another small difference: jax.numpy functions require arrays (or values of Python’s built-in numeric types) as inputs, rather than silently converting Python lists:

jnp.sum([1, 2, 3])

TypeError: sum requires ndarray or scalar arguments, got <class 'list'> at position 0.

This is deliberate. Silently converting lists to arrays is a common source of hidden performance problems, so JAX asks you to do it explicitly:

jnp.sum(jnp.array([1, 2, 3]))

Array(6, dtype=int32)

Arrays live on devices

A JAX array’s data lives on one or more devices, like CPU, GPU, or TPU. The same JAX code runs on all of them; JAX places arrays on the default device (typically your accelerator, if you have one) and computations follow their data. An array can even be sharded across many devices, so that JAX programs can be written once and run on one chip or thousands. Every array carries a sharding attribute describing exactly how its data is placed:

x.sharding

SingleDeviceSharding(device=CpuDevice(id=0), memory_kind=device)

On a single-device machine this isn’t very interesting, but it’s the foundation of JAX’s approach to parallelism and scaling. That story is told in the performance and scaling docs; see Distributed arrays and automatic parallelization.

Under the hood: jax.numpy, jax.lax, and XLA

jax.numpy is a convenience layer. Its functions are implemented in terms of jax.lax, a stricter, lower-level API. Most jax.numpy functions are thin wrappers that handle NumPy-style conveniences and then defer to their jax.lax counterpart. Here’s jnp.sin, simplified from the real implementation:

def sin(x):
  x = jnp.asarray(x)                            # accept built-in Python numbers
  if not jnp.issubdtype(x.dtype, jnp.inexact):
    x = x.astype(float)                         # promote integers to floating point
  return lax.sin(x)

The jax.numpy layer’s job is NumPy-style argument handling (accepting built-in Python numbers, promoting dtypes) while the computation itself belongs to jax.lax. The jax.lax operations in turn correspond closely to XLA HLO operations, the vocabulary of XLA, the compiler that ultimately runs JAX computations: lax.sin maps essentially one-to-one onto XLA’s Sin.

Being closer to the compiler, jax.lax skips the conveniences. For example, where jax.numpy implicitly promotes mixed types:

jnp.add(1, 1.0)

Array(2., dtype=float32, weak_type=True)

jax.lax requires explicit promotion:

from jax import lax
lax.add(1, 1.0)

TypeError: lax.add requires arguments to have the same dtypes, got int32, float32. (Tip: jnp.add is a similar function that does automatic type promotion on inputs).

In exchange for the strictness, jax.lax exposes operations that are more general than NumPy’s, like arbitrarily strided and dilated convolutions, and structured control flow. You can write a lot of JAX without touching jax.lax, but it’s useful to know it’s there: when you can’t find a jax.numpy function for something, check jax.lax.

Next steps

Arrays and jax.numpy are how you express numerical computation. But JAX is also a system for transforming computations, covered next in Transformations: grad and vmap.