Minimal Overhead There are growing number of options for accelerating sparsity in neural net-
works (e.g. N:M sparsity (Mishra et al., 2021), CPU-acceleration (Elsen et al., 2020), activation
sparsity (Gale et al., 2022)). However, integration with existing frameworks is often lacking, mak-
ing these advances relatively difﬁcult to use, especially in research. Given our main goal of facilitat-
ing research, JaxPruner follows the tradition of using binary masks for introducing sparsity, which
introduces some additional operations and requires additional storage for the masks. We aim to
minimize this overhead. Building on this initial version, we hope to introduce further optimizations
as software and hardware support matures over time. As a ﬁrst step towards this goal, we provide
integration with the jax.experimental.sparse1 module, which converts pruned parameters
into sparse representations and reduces the memory footprint signiﬁcantly.

3 QUICK OVERVIEW
JAXPRUNER