def test_elbo_dynamic_support():
x_prior = dist.TransformedDistribution(
dist.Normal(),
[AffineTransform(0, 2),
SigmoidTransform(),
AffineTransform(0, 3)])
x_guide = dist.Uniform(0, 3)
def model():
numpyro.sample('x', x_prior)
def guide():
numpyro.sample('x', x_guide)
adam = optim.Adam(0.01)
# set base value of x_guide is 0.9
x_base = 0.9
guide = substitute(guide, base_param_map={'x': x_base})
svi = SVI(model, guide, adam, ELBO())
svi_state = svi.init(random.PRNGKey(0))
actual_loss = svi.evaluate(svi_state)
assert np.isfinite(actual_loss)
x, _ = x_guide.transform_with_intermediates(x_base)
expected_loss = x_guide.log_prob(x) - x_prior.log_prob(x)
assert_allclose(actual_loss, expected_loss)
def test_elbo_dynamic_support():
x_prior = dist.TransformedDistribution(
dist.Normal(), [AffineTransform(0, 2), SigmoidTransform(), AffineTransform(0, 3)])
x_guide = dist.Uniform(0, 3)
def model():
numpyro.sample('x', x_prior)
def guide():
numpyro.sample('x', x_guide)
adam = optim.Adam(0.01)
x = 2.
guide = substitute(guide, data={'x': x})
svi = SVI(model, guide, adam, Trace_ELBO())
svi_state = svi.init(random.PRNGKey(0))
actual_loss = svi.evaluate(svi_state)
assert jnp.isfinite(actual_loss)
expected_loss = x_guide.log_prob(x) - x_prior.log_prob(x)
assert_allclose(actual_loss, expected_loss)