5 RELATED WORK

JAX (Bradbury et al., 2018) is a recently introduced Python library for high-performance machine
learning research. With Autograd (Baydin et al., 2018), JAX can automatically differentiate native
Python and Numpy functions such as loops and branches in both forward and backward modes.
Also, JAX supports just-in-time compilation of NumPy programs on multiple GPUs or TPUs in par-
allel with XLA (Sabne, 2020). Following the functional programming paradigm all transformations
in JAX work on pure functions and thus can be composed together in an arbitrary order. Since these
features can dramatically facilitate machine learning research, the community has started adopting
JAX to develop new research frameworks in recent years, including for example Optax (Babuschkin
et al., 2020), FedJAX (Ro et al., 2021), Flax (Heek et al., 2020), JaxOpt (Blondel et al., 2021) just
to name a few.

3https://www.tensorﬂow.org/datasets/catalog/c4

6
JAXPRUNER