text: 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.
url: https://github.com/ElArkk/jax-unirep