The type system

Every value in a StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → program carries an explicit type, and the notation makes that impossible to miss. tensor<32x64xbf16> reads left to right: dimensions first, then the element type, so a 32 by 64 matrix stored in bfloat16 shows up exactly that way in the dump you get out of JAX's lowering. There is no untyped register, no implicit widening. If a computation touches a value, the compiler already knows its shape and its dtype before a single operation runs.

Scalars do not get a separate type. A single number is a tensor with rank zero, written tensor<f32> or tensor<i32>, with no dimension list at all. That is a deliberate simplification: one rule, tensor plus dtype, covers both a 4096 by 4096 matmul operand and a loop counter. Booleans follow the same logic with their own element type, i1, a one-bit integer. A comparison op like stablehlo.compare returns tensor<i1> for a scalar predicate or tensor<32xi1> for a per-row mask, and that mask is exactly what feeds a stablehlo.select or a stablehlo.if condition.

This uniformity is the point of the whole opset. Where older HLO text mixed shape information into operation names, StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → pushes it entirely into the type. An add of two tensor<8x128xbf16> values type-checks the same way as an add of two tensor<f32> scalars, because both operands and the result are just tensors with a shape and a dtype attached. Reading a dump becomes a matter of reading types, not memorizing what each op implicitly assumes about its inputs.

The container structure is plain. A StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → program is a module, and a module holds one or more func.func definitions, each with a typed argument list and a typed return, mirroring the tensor notation used everywhere else. jax.jit traces your Python function once and lowers it to exactly this shape, a module wrapping a func whose body is a sequence of typed ops. Nothing about the module or func syntax is TPU specific; it is the same container StableHLO uses for CPU and GPU lowerings too.

You do not need a compiler background to use this. When a Pallas kernel throws a shape error, the fastest diagnosis is opening the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → dump and reading the type on the operand that misbehaves, because the mismatch is stated directly in the text rather than buried in a stack trace. A tensor<256x128xf32> where you expected tensor<256x128xbf16> tells you exactly where a dtype conversion happened, or didn't, in the lowering.