JaxPruner: A concise library for sparsity research

Joo Hyung Lee, Wonpyo Park, Nicole Mitchell, Jonathan Pilault, Johan Obando-Ceron +20 more
4/27/2023
cs.LGcs.SE

No official code found. Best implementation: google-research/jaxpruner (236 stars, Apache-2.0, last commit 2025-02-12).

Abstract

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.

View on arXivView PDF

Code Implementations(1)

google-research/jaxprunerCommunity100%
Apache-2.0
23621Shell, PythonMar 19, 20231 years agoApache-2.0

Similar Papers

FasterAI: A Lightweight Library for Creating Sparse Neural Networks

79% match2022

Controlled Sparsity via Constrained Optimization or: How I Learned to Stop Tuning Penalties and Love Constraints

79% match2 repos1662022

DASS: Differentiable Architecture Search for Sparse neural networks

78% match1 repo32022

The State of Sparsity in Deep Neural Networks

78% match1 repo3132019

Sparse by Rule: Probability-Based N:M Pruning for Spiking Neural Networks

78% match2025

JaxDecompiler: Redefining Gradient-Informed Software Design

78% match1 repo162024

Cite this paper

@article{lee2023jaxpruner,
  title  = {JaxPruner: A concise library for sparsity research},
  author = {Joo Hyung Lee and Wonpyo Park and Nicole Mitchell and Jonathan Pilault and Johan Obando-Ceron and Han-Byul Kim and Namhoon Lee and Elias Frantar and Yun Long and Amir Yazdanbakhsh and Shivani Agrawal and Suvinay Subramanian and Xin Wang and Sheng-Chun Kao and Xingyao Zhang and Trevor Gale and Aart Bik and Woohyun Han and Milen Ferev and Zhonglin Han and Hong-Seok Kim and Yann Dauphin and Gintare Karolina Dziugaite and Pablo Samuel Castro and Utku Evci},
  year   = {2023},
  eprint = {2304.14082},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2304.14082v3}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2304.14082.svg)](https://codeofpaper.com/paper/2304.14082)

Discussion