Namhoon Lee helped with reviewing and writing the manuscript.

added the parallelization example to the quick start colab. Tested the code and ﬁxed

helped with the jax2tf integration, Tensorﬂow N:M weight/activation sparsity imple-

Elias Frantar
bugs.

Yun Long
mentation.

Amir Yazdanbakhsh implemented N:M sparsity support.

12

Published as a conference paper at ICLR 2023

Shivani Agrawal helped with the quantization-aware training example.

Suvinay Subramanian helped with documentation and testing the library/colabs.

Xin Wang

helped with t5x integration.

Sheng-Chun Kao implemented global sparsity.

Xingyao Zhang conducted investigation on converting pruned JAX model parameters to use
jax.sparse.BCOO with Utku.

Trevor Gale

tested and open-sourced the Colabs.

Aart Bik works on sparse compiler and sparse JAX. Contributed to the related work section in
paper.

Woohyun Han implemented block sparsity.

Milen Ferev implemented the tﬂite conversion colab.
JAXPRUNER