logo
Updated

JAX

transformer from g guy

Patrick C Toulme
Patrick C Toulme@PatrickToulme
Last week I traced Splash Attention through the TPU compiler stack. But before that, I had to figure out how JAX actually compiles to TPU machine code. Google's TPU compiler is closed-source. The IRs are undocumented. So I rented a v6e for $1 and traced every layer: JAX → HLO

https://github.com/MizuhoAOKI/jax_generative_models

https://ekzhang.substack.com/p/how-the-jaxjit-jit-compiler-works

Links

bilal
bilal@bilaltwovec
lots of fun quick tpu/performance engineering interview questions in here :D https://t.co/DLS93zWXNi

https://github.com/patrick-kidger/jaxtyping

https://github.com/HMUNACHI/nanodl

https://github.com/google/aqt

https://github.com/google/flax

https://github.com/NVIDIA/JAX-Toolbox

https://ekzhang.substack.com/p/how-the-jaxjit-jit-compiler-works

https://sankalp.bearblog.dev/einsum-new/

https://github.com/xjdr-alt/simple_transformer

https://github.com/Gsunshine/meanflow

https://github.com/JoWestermann/smolyax