JAXRENS is the main software package I developed in the course of my PhD. It is a fully-fledged atomistic nested sampling code for JAX-based machine learning potentials, that can be run in parallel on multiple GPUs. The predecessor of JAXRENS has already undergone significant stress tests in our RENS and RENS AL studies. It resembles the by far most sophisticated software project I have ever developed myself. Figuring out how to squeeze out good performance on (multiple) hardware accelerators was not always straightforward and required me to navigate around various technical problems and looking at NVIDIA Nsight profiles for weeks.

Find JAXRENS here: https://github.com/nunglert/jaxrens.