Research First Often research projects require running multiple algorithms and baselines, and so
they beneﬁt greatly from rapid prototyping. JaxPruner achieves this by committing to a generic API
shared among different algorithms, which in return makes it extremely easy to switch between differ-
ent algorithms. We provide implementations for common baselines and aim to make our algorithms
easy to modify. Furthermore it is extremely easy to switch between common sparsity structures
(unstructured, N:M, Block). A quick overview of such features are discussed in next section.
JAXPRUNER