jax-ml/oryx

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