jax 6
- 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