167 citations · 393 across the 12 of their papers we have counts for
1 paper · 1 filter
Keith Rush, Zachary Charles, Zachary Garrett +2
We present DrJAX, a JAX-based library designed to support large-scale distributed and parallel machine learning algorithms that use MapReduce-style operations. DrJAX leverages JAX'…