Arrays

Try x[0] = 1 on a jnp array and JAX refuses outright. Arrays here are immutable, so the way you express an update is different: x.at[0].set(1.0) never touches x, it returns a new array with the change already in it. The copy you might now be picturing on every line like this mostly does not happen: under jit, XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → can see the moment the old value goes dead and reuse its buffer for the output instead of allocating a second one. Donation and aliasing make that reuse explicit at a function's boundaries, and chapter 11 comes back to donation once you have a concrete reason to control it yourself.

grad, vmap, and jit all work the same way: each records what a function does to its inputs, then transforms that recording. A mutation hidden inside the function would make the recording describe something that never happened, a step the transform never saw. That is why immutability here is not a style preference. Values in, values out is what keeps the recording honest.

An array is a value, not a place.

import jax.numpy as jnp

x = jnp.zeros(4)
y = x.at[0].set(1.0)   # a new array; x is unchanged
print(x)               # [0. 0. 0. 0.]
print(y)               # [1. 0. 0. 0.]

Ask jnp.arange(3) for its dtype and you get int32. Add 1.5 to it and you might expect float64, the way NumPy would sometimes give you depending on the values involved; JAX gives you float32 instead, because the Python float 1.5 is weakly typed. A weak scalar adapts to the array it meets rather than promoting the array up to itself.

That adaptation comes from a promotion lattice JAX defines and documents, not from NumPy's older, value-dependent rules. The two disagree exactly where old NumPy habits are most likely to mislead you, so reaching for NumPy intuition on a promotion question here is the wrong instinct.

The default floating dtype is float32, and 64-bit is off unless you turn on jax_enable_x64. That default is where the biggest surprise in this section lives: jnp.array(1e300) does not raise, it overflows to inf, silently, by design. It is the first x64 surprise almost everyone meets, and now you have met it here instead of in a debugger.

import jax.numpy as jnp

print(jnp.arange(3).dtype)           # int32
print((jnp.arange(3) + 1.5).dtype)   # float32: the Python float is weak
print(jnp.array(1e300))              # inf: float32 unless you enable x64

Run y = x @ x on a large array and the line returns almost instantly, long before a matmul that size could actually finish. JAX dispatches operations asynchronously: what you get back is a future-like array while the device keeps computing behind it. Printing it or converting it to NumPy blocks until the value is ready; calling y.block_until_ready() blocks on purpose, explicitly, wherever you need the wait.

This is exactly why wrapping a single op in time.perf_counter() without a block measures dispatch time, not compute time. The clock stops the moment Python gets its future back, not when the device finishes the work. Chapter 11 builds the honest benchmarking harness on top of this one fact, but the fact itself belongs here, at the first place you would be tempted to time something.

import time
import jax.numpy as jnp

x = jnp.ones((2000, 2000))
t0 = time.perf_counter()
y = x @ x                     # returns before the matmul finishes
t1 = time.perf_counter()
y.block_until_ready()         # now it is actually done
t2 = time.perf_counter()
print(f"dispatch {1e3 * (t1 - t0):.2f} ms · compute {1e3 * (t2 - t1):.2f} ms")

Every array you create in JAX, whatever operation produced it, is the same type: jax.Array. It knows where it lives through a .sharding attribute; on a single device that sharding is a SingleDeviceSharding, the simplest case of a much bigger idea. Chapter 10 stretches that same object across a whole mesh of devices without changing the shape of your code at all.

jax.device_put is how you place an array somewhere on purpose, and once it is placed, computation follows the operands: an op involving that array runs where the array lives, not wherever Python happens to be executing.