Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Numpyro #18

Draft
wants to merge 5 commits into
base: master
Choose a base branch
from
Draft

Numpyro #18

wants to merge 5 commits into from

Conversation

reubenharry
Copy link
Collaborator

Here's a first attempt at incorporating mclmc into numpyro. The main thing missing is tuning. What I've done is supplied a new kernel for numpyro's MCMC, namely the MCLMC kernel, using dynamics_generalized (i.e. the integrator for the SDE).

This isn't intended to be merged as is (hence the draft status of the pull request), but I'm posting here in case the code is useful.

@JakobRobnik
Copy link
Owner

JakobRobnik commented Sep 8, 2023 via email

@reubenharry
Copy link
Collaborator Author

@JakobRobnik Here's an update with an actual probabilistic program. See the code for details, but in brief:

def m():

    mu = numpyro.sample('mu', dist.Normal(3, 1))
    nu = numpyro.sample('nu', dist.Normal(mu+1, 2))
    

kernel = MCLMC(m)
# kernel = NUTS(m)
mcmc = MCMC(kernel, num_warmup=1000, num_samples=1000)
mcmc.run(random.PRNGKey(0))
posterior_samples = mcmc.get_samples()
mcmc.print_summary()

This shound not be used in practice before some thorough testing is done with a range of distributions (I already caught some silly bugs), and probably there's a more idiomatic way to do things in numpyro, but I'm pushing it since you said some people might find it useful.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@reubenharry @JakobRobnik