JAX Transformers Attention Large Memory Layers with Product Keys Augmenting transformer language models with sparse access of large memory matrices by Madison May
JAX Machine Learning Programming Finetuning Transformers with JAX + Haiku A practical, code-first look at DeepMind's new haiku library. by Madison May
Machine Learning Programming JAX A First Look at JAX Put on your metaphorical safety goggles and start building something weird with JAX. by Madison May