Ramsey: probabilistic modelling using Haiku#

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 its API. Ramsey implements probabilistic models, such as neural processes, Gaussian processes, Bayesian neural networks, Bayesian timeseries models and state-space-models, and more.

Example#

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.data import sample_from_sinus_function
from ramsey.models import NP

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)

(x, y), _ = sample_from_sinus_function(random.PRNGKey(0))

neural_process = hk.transform(neural_process)
params = neural_process.init(
    random.PRNGKey(1), x_context=x, y_context=y, x_target=x
)

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.

Installation#

To install from PyPI, call:

pip install ramsey

To install the latest GitHub <RELEASE>, 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:

Clone Ramsey and install it and its dev dependencies via pip install -e '.[dev]',
create a new branch locally git checkout -b feature/my-new-feature or git checkout -b issue/fixes-bug,
implement your contribution,
test it by calling tox on the (Unix) command line,
submit a PR 🙂

License#

Ramsey is licensed under the Apache 2.0 License.

Ramsey: probabilistic modelling using Haiku

Contents