This paper introduces JaxPruner, an open-source JAX-based pruning and sparse training library for machine learning research. JaxPruner aims to accelerate research on sparse neural networks by providing concise implementations of popular pruning and sparse training algorithms with minimal memory and latency overhead. Algorithms implemented in JaxPruner use a common API and work seamlessly with the popular optimization library Optax, which, in turn, enables easy integration with existing JAX based libraries. We demonstrate this ease of integration by providing examples in four different codebases: Scenic, t5x, Dopamine and FedJAX and provide baseline experiments on popular benchmarks.

翻译：本文介绍JaxPruner，一个基于JAX的开源剪枝与稀疏训练库，用于机器学习研究。JaxPruner通过提供流行的剪枝和稀疏训练算法的简洁实现，并以最小内存和延迟开销加速稀疏神经网络研究。该库中实现的算法采用通用API，并能与主流优化库Optax无缝协作，从而便于与现有基于JAX的库集成。我们通过四个不同代码库（Scenic、t5x、Dopamine和FedJAX）的示例展示了这一集成便利性，并在主流基准测试上提供了基线实验。