Oryx is a library for probabilistic programming and deep learning built on top of Jax. - View it on GitHub