custom_call, the escape hatch

StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → defines a closed opset: adds, compares, dots, reductions, control flow, a fixed vocabulary the compiler fully understands. Anything outside that vocabulary still has to travel through the program somehow, and the mechanism is stablehlo.custom_call. A custom call carries an opaque target name, a typed operand list, a typed result list, and a payload the surrounding compiler treats as a black box. StableHLO does not need to understand what runs inside; it only needs to know the types going in and the types coming out.

A Pallas kernel rides through this exact door. When you call pallas_call, JAX does not try to express your kernel body as StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → ops; it wraps the whole thing as a custom_call targeting tpu_custom_call, with the serialized MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → module attached as the payload. From the surrounding StableHLO program's point of view, your kernel is one opaque op with typed inputs and outputs, same as any other custom call. The grid, the block specs, the pipeline, none of it is visible at the StableHLO level. It only becomes visible one layer down, inside Mosaic.

You can see this directly in a real trace. The museum's 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 → overflow capture, from a kernel that blew its VMEM budget, shows custom_call_target "tpu_custom_call" right in the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → op label, because that is literally what the compiler recorded when it lowered the pallas_call. Reading that label confirms the failure is happening inside a MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic →-compiled kernel, not somewhere else in the surrounding StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → graph, which narrows the search before you even open the Mosaic dump.

This is the seam the rest of the stack builds on. Everything past this point, how the MosaicThe MLIR dialect Pallas lowers to, and the last layer of the TPU stack you can read; only LLO below it is closed.taught in /l/mosaic → module inside that payload gets built, scheduled, and lowered to real TPU instructions, is a separate compiler with its own passes and its own failure modes. StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo →'s job ends at the custom call boundary, and Mosaic's job starts there, picking up the same typed operands and turning them into an actual pipelined kernel running on the MXU.