jax-unirep

My intern, Arkadij Kummer, and I found that the original implementation of UniRep was slow and non-performant, so we reimplemented it in pure NumPy/JAX.

The GitHub repository can be found here.