pyro-ppl / pyro-ppl/numpyro

init_strategy does not work with CircularReparam

Open
#1,614 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
2.8k
Forks
316
Avg merge
3d 9h
Merged PRs (30d)
27

Description

Hi!

I am having an issue with initialising a model with the init_by_median() strategy for my MCMC NUTS model.
I am drawing the location parameter of a VonMises distribution from a VonMises prior and I reparam this parameter with CircularReparam().

Working example:

import numpyro
from numpyro import sample, plate
from numpyro.infer import NUTS, MCMC, init_to_median
from numpyro.distributions import VonMises
from numpyro.handlers import reparam
from jax.random import PRNGKey
import jax.numpy as jnp

@numpyro.handlers.reparam(
    config={"loc_1": numpyro.infer.reparam.CircularReparam()}
)
def tst(x):
    loc_1 = sample("loc_1", VonMises(0, 100))
    conc = 2
    
    with plate("plate", len(x)):
        y = sample("y", VonMises(loc_1, conc), obs=x)
    return loc_1

if __name__ == "__main__":
    kernel = NUTS(tst, init_strategy=init_to_median(num_samples=1000))
    mcmc = MCMC(kernel, num_warmup=1, num_samples=10)

    mcmc.warmup(PRNGKey(42), collect_warmup=True, x=jnp.array([1, 2, 3]))

    print(mcmc.get_samples())

Output:

{'loc_1': Array([1.7704213], dtype=float32), 'loc_1_unwrapped': Array([1.7704215], dtype=float32)}

This should be the median of VonMises(0, 100), which is 0. If you remove the reparameterisation, you get exactly this value.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Run the provided tst example with NUTS, MCMC, init_to_median, and CircularReparam, then compare it with the version that removes the reparameterization. Trace how init_to_median handles loc_1 through CircularReparam and verify that the completed behavior returns the median of VonMises(0, 100), namely 0, without the conflicting loc_1_unwrapped result.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.