Fast Integration Research in Machine Learning (ML) is fast paced. Combined with the huge
variety of ML applications, this result in a high number of ever-changing codebases. At the same
time, adaptability of new research ideas is highly correlated with their ease of use. Therefore, we
aimed to reduce friction for those integrating JaxPruner into existing codebases. JaxPruner uses
the popular Optax optimization library (Babuschkin et al., 2020) to achieve this, requiring minimal
changes when integrating with existing libraries. State variables (i.e. masks, counters) needed for
pruning and sparse training algorithms are stored together with the optimization state, which makes
parallelization and checkpointing easy.
JAXPRUNER