Oryx is a library for probabilistic programming and deep learning built on top of Jax. The approach is to expose a set of function transformations that compose and integrate with JAX's existing transformations (e.g. jit, grad, and vmap).

This is not an official Google product

You can install Oryx via pip:

$ pip install oryx

Documentation is available on the Oryx website.

• A tour of Oryx
• Probabilistic programming

To develop and modify Oryx, you need to install poetry, a tool for Python packaging and dependency management.

To install the development dependencies of Oryx, you can run

$ poetry install

and to enter a virtual environment for testing or debugging, you can run:

$ poetry shell

Oryx uses Pytest for testing. To run the tests, use the following command:

$ poetry run pytest oryx