def model():
x = numpyro.sample("x", dist.Normal())
with numpyro.handlers.mask(mask=False):
numpyro.sample("y", dist.Normal(x), obs=1)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=1)
mcmc.run(random.PRNGKey(1))
assert_allclose(mcmc.get_samples()['x'].mean(), 0., atol=0.1)
@pytest.mark.parametrize('init_strategy', [
init_to_feasible(),
init_to_median(num_samples=2),
init_to_sample(),
init_to_uniform(radius=3),
init_to_value(values={'tau': 0.7}),
init_to_feasible,
init_to_median,
init_to_sample,
init_to_uniform,
init_to_value,
])
def test_initialize_model_change_point(init_strategy):
def model(data):
alpha = 1 / jnp.mean(data.astype(np.float32))
lambda1 = numpyro.sample('lambda1', dist.Exponential(alpha))
lambda2 = numpyro.sample('lambda2', dist.Exponential(alpha))
tau = numpyro.sample('tau', dist.Uniform(0, 1))
lambda12 = jnp.where(
jnp.arange(len(data)) < tau * len(data), lambda1, lambda2)
return numpyro.sample("obs",
dist.Bernoulli(logits=logits),
obs=labels)
return true_coefs, (data, labels), model
########################################
# Stein Exterior (Smoke tests)
########################################
@pytest.mark.parametrize("kernel", KERNELS)
@pytest.mark.parametrize(
"init_loc_fn",
(init_to_uniform(), init_to_sample(), init_to_median(),
init_to_feasible()),
)
@pytest.mark.parametrize("auto_guide", (AutoDelta, AutoNormal))
@pytest.mark.parametrize("problem", (uniform_normal, regression))
def test_steinvi_smoke(kernel, auto_guide, init_loc_fn, problem):
true_coefs, data, model = problem()
stein = SteinVI(
model,
auto_guide(model, init_loc_fn=init_loc_fn),
Adam(1e-1),
Trace_ELBO(),
kernel,
)
stein.run(random.PRNGKey(0), 1, *data)