jax 9
- Ten Ways to Break XLA
- Hacking on XLA: Building It, and Teaching It a New Pass
- Types Without Places, Places Without Types
- How JAX Shards a Computation Across a Mesh
- When XLA Isn't Enough: Pallas, Mosaic, and Triton
- What JAX Traces, and What It Refuses
- Programming the TPU: What Its Open-Source Compiler Already Tells You
- XLA Up Close: The Performance Bargain and Its Rigid Price
- Three Ways to Take a Gradient: Tape, Trace, and Source Transform