Oryx

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

Installation

You can install Oryx via pip:

$ pip install oryx

Documentation and Examples

Documentation is available on the Oryx website.

Guides

• A tour of Oryx
• Probabilistic programming

Development

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

Running tests

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

$ poetry run pytest oryx