A performant reimplementation of the UniRep model.
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.