Stax

References for the whole chapter, plus appendices: a source-code index, measurement scripts, and worked solutions for selected exercises.

Content-addressed caching, parallel compilation, fallback semantics, design comparisons against PyTorch 2 and XLA, measurement models, exercises, and labs.

GPU execution and Triton programming models, reduction kernels, multi-reduction fusion, launch configurations, epilogue fusion, and CUDA Graphs.

Buffer dependencies, the greedy fusion algorithm, buffer reuse, affine dependences and tiling, plus the full CPU codegen path from vectorization to wrappers.

Reverse-mode autodiff over the operator-level graph, the minimum-cut model for save-versus-recompute, and lowering to a loop-level intermediate representation.