Pith. sign in

REVIEW 1 cited by

JaxPruner: A concise library for sparsity research

Not yet reviewed by Pith; the record is open.

Re-run · record.json Download PDF Read on arXiv ↗

This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.

SPECIMEN: schema-true, not a live event

T0 review · schema-true

One-sentence machine reading of the paper's core claim.

pith:XXXXXXXX · record.json · timestamp

arxiv 2304.14082 v3 pith:KMWRMIX4 submitted 2023-04-27 cs.LG cs.SE

Joo Hyung Lee , Wonpyo Park , Nicole Mitchell , Jonathan Pilault , Johan Obando-Ceron , Han-Byul Kim , Namhoon Lee , Elias Frantar
show 17 more authors
Yun Long Amir Yazdanbakhsh Shivani Agrawal Suvinay Subramanian Xin Wang Sheng-Chun Kao Xingyao Zhang Trevor Gale Aart Bik Woohyun Han Milen Ferev Zhonglin Han Hong-Seok Kim Yann Dauphin Gintare Karolina Dziugaite Pablo Samuel Castro Utku Evci
This is my paper · ORCID
classification cs.LGcs.SE
keywords jaxprunerlibrarypopularresearchsparsealgorithmsconciseintegration
verification ladder T0 review T1 audit T2 compute T3 formal
0 comments
read the original 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.

Discussion (0). Continue with ORCID to comment.

Forward citations

Cited by 1 Pith paper

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score. Full citation record

  1. Accelerating Linear Recurrent Neural Networks for the Edge with Unstructured Sparsity

    cs.LG 2025-02 conditional novelty 6.0 of 10

    Sparse, 8-bit quantized S5 linear RNNs match dense model audio denoising accuracy with 2x less compute and 36% less memory, and run 42x faster with 149x lower energy on Loihi 2 than a dense FP32 model on Jetson Orin Nano.

Pith tools