DrJax: Scalable and Differentiable MapReduce Primitives in JAX

Keith Rush, Zachary Charles, Zachary Garrett, Sean Augenstein, Nicole Mitchell
January 2024
PDF arXiv URL