A library for probabilistic modelling using Haiku and JAX
Project description
Ramsey
Probabilistic modelling using Haiku and JAX
About
Ramsey is a library for probabilistic modelling using Haiku and JAX. It builds upon the same module system that Haiku is using and is hence fully compatible with Haiku's and NumPyro's API.
Example usage
Ramsey uses to Haiku's module system to construct probabilistic models and define parameters. For instance, a simple neural process can be constructed like this:
import haiku as hk
import jax.random as random
from ramsey import NP
from ramsey.data import sample_from_sine_function
def neural_process(**kwargs):
dim = 128
np = NP(
decoder=hk.nets.MLP([dim] * 3 + [2]),
latent_encoder=(
hk.nets.MLP([dim] * 3), hk.nets.MLP([dim, dim * 2])
)
)
return np(**kwargs)
key = random.PRNGKey(23)
(x, y), _ = sample_from_sine_function(key)
neural_process = hk.transform(neural_process)
params = neural_process.init(key, x_context=x, y_context=y, x_target=x)
Installation
To install from PyPI, call:
pip install ramsey
To install the latest GitHub , just call the following on the command line:
pip install git+https://github.com/ramsey-devs/ramsey@<RELEASE>
See also the installation instructions for Haiku and JAX, if you plan to use Ramsey on GPU/TPU.
Contributing
Contributions in the form of pull requests are more than welcome. A good way to start is to check out issues labelled "good first issue".
In order to contribute:
- Install Ramsey and dev dependencies via
pip install -e '.[dev]'
, - test your contribution/implementation by calling
tox
on the (Unix) command line before submitting a PR.
Why Ramsey
Just as the names of other probabilistic languages are inspired by researchers in the field (e.g., Stan, Edward, Turing), Ramsey takes its name from one of my favourite philosophers/mathematicians, Frank Ramsey.
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.