bioRxivBioengineeringPreprint
Reimplementing Unirep in JAX
Abstract
UniRep is a recurrent neural network model trained on 24 million protein sequences, and has shown utility in protein engineering. The original model, however, has rough spots in its implementation, and a convenient API is not available for certain tasks. To rectify this, we reimplemented the model in JAX/NumPy, achieving near-100X speedups in forward pass performance, and implemented a convenient API for specialized tasks. In this article, we wish to document our model reimplementation process with the goal of educating others interested in learning how to dissect a deep learning model, and engineer it for robustness and ease of use.
Reproducibility ledger
Nobody has reported trying to reproduce this paper yet. Did you replicate an experiment or run the code? Your report saves the next lab a week.
Sign in to file a reproducibility report.
Related papers
- Unified rational protein engineering with sequence-only deep representation learning
Ethan C. Alley, Grigory Khimulya, … George McDonald Church·Nature MethodsPeer-reviewed·Oct 21, 2019
- Model reconstruction from Small angle X-ray Scattering data using deep learning methods
Hao He, Can Liu, Haiguang Liu·iSciencePeer-reviewed·Feb 12, 2020
- NeRFax: An efficient and scalable conversion from the internal representation to Cartesian space
Oliver Dutton, Falk Hoffmann, Kamil Tamiola·bioRxiv·May 25, 2022
- Jaxley: Differentiable simulation enables large-scale training of detailed biophysical models of neural dynamics
Michael Deistler, Kyra L. Kadhim, … Jakob H. Macke·Nature MethodsPeer-reviewed·Nov 13, 2025
- jaxon: a differentiable, GPU-native simulator for peripheral-nerve fiber models
David Lung, Max Haberbusch·bioRxiv·Jul 31, 2026
- DeepImageJ: A user-friendly environment to run deep learning models in ImageJ
E. Gomez-de-Mariscal, C. Garcia-Lopez-de-Haro, … D. Sage·bioRxiv·May 6, 2021