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