Refs, and why mutation entered a functional language

Hand a Ref to jnp.exp and it will not trace. The design document states the rule as a property of the type: refs cannot be passed into the usual set of JAX primitives without being read from first. Read one and you get a JAX Array back. Write into one and what you write must be an Array. Those two sentences are the entire contract, and every kernel you have written obeys them whether or not you noticed.

A read is not notation. On TPU it means loading into vector registers, on GPU into the lowest level of the hierarchy, and a write is the same move in reverse. So x_ref[...] + y_ref[...] is two loads and an add, in that order, and the kernel jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → shows all three.

This site's own capture makes it visible. The binders are memrefs, not arrays, and each <- is a real transfer between memory and registers.

{ lambda ; a:MemRef<None>{bfloat16[256,256]} b:MemRef<None>{bfloat16[256,256]} c:MemRef<None>{bfloat16[256,256]}. let
    d:bf16[256,256] <- a[:,:]
    e:bf16[256,256] <- b[:,:]
    f:bf16[256,256] = add d e
    c[:,:] <- f
  in () }

The design document is careful about the provenance here: refs are not a Pallas-specific concept, they were introduced to JAX to represent stateful computations, and Pallas leverages them for kernels that operate on mutable memory. That matters for how you should read the restriction. A ref refusing to enter jnp.exp is not a rule Pallas bolted onto JAX. It is what the type has always meant.

The practical consequence is that a kernel body has no hidden loads. Every trip between memory and registers is a term in the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, which is why counting loads in a printed kernel is a legitimate way to reason about a body that feels slower than it should. Read the same ref twice in a body and you get two loads, because you asked for two.

The memory-spaces lesson before this one covers VMEMThe TPU’s software-managed vector scratchpad, about 128 MiB. Blocks must be staged here before compute touches them; what is resident is what your schedule staged.taught in /l/tpu → as the default, SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → for scalars, and ANYThe memory space that tells the compiler not to place a ref at all: a promise that the kernel will move the data itself with a manual DMA.taught in /l/pallas → as the promise that you will move the data yourself. The detail the guide leaves implicit is what ANY forbids. A buffer in the ANY space cannot be dereferenced with ordinary indexing at all: x_ref[...] on it is not a slow read, it is not a read. You copy into a VMEM or SMEM buffer first with pltpu.sync_copy or pltpu.async_copy, and only then does indexing mean anything.

ANY is also documented as a hint rather than an address. It tells the compiler the memory space is unconstrained, and in most cases XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → will place the buffer in HBM. You get a promise about who is responsible, not a guarantee about where the bytes sit.

The TPU reference draws the boundary by rank, not by size. Every 0D array is stored in scalar registers and its operations run on the scalar core. Everything else runs on the vector core, and the document says so explicitly for the case people get wrong: even a single-element array of rank 1 or more goes to the vector unit.

That rule has a price tag attached, because all vector computation is padded up to the tile. Adding two 1x1 arrays costs what adding two 8x128 arrays costs. So a sequence length kept as a shape-() scalar in SMEMScalar memory: lengths, flags, and indices live here, feeding control flow without ever entering the vector datapath.taught in /l/tpu → is a scalar-core decision; the same number wrapped in a length-1 vector is a full tile of vector work to compute one comparison. The shape you chose picked the processor.

input_output_aliases tells pallas_call that an output may reuse an input's buffer, which is donation at kernel granularity. The aliasing lesson later in this arc covers the mechanics. What it does not spell out is the ordering consequence: once the two share memory, a write to the output ref is a write to the input, so within a grid step the order of your reads and writes is now semantics. Read the aliased input after writing the output and you read what you just wrote.

The sparse-kernel guide uses aliasing for something other than saving an allocation, and it is worth stealing. In its block-sparse matmul, some output blocks are never visited by the grid at all, so their buffer would hold whatever was there. It passes an array of zeros in and aliases it onto the output, which makes "never visited" mean zero instead of uninitialized. Aliasing as an initialization strategy, not a memory optimization.

kernel = pl.pallas_call(
  dsd_kernel,
  grid_spec=grid_spec,
  out_shape=out_shape,
  # We use input-output aliases to zero-out o_ref for blocks that we never
  # visit. By passing in an array of zeros we avoid having o_ref start with
  # uninitialized values.
  input_output_aliases={4: 0},  # Map zeros to o_ref.
)