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 __init__(self, low=0., high=1., validate_args=None):
self.low, self.high = promote_shapes(low, high)
batch_shape = lax.broadcast_shapes(np.shape(low), np.shape(high))
base_dist = _BaseUniform(batch_shape)
super(Uniform, self).__init__(base_dist,
AffineTransform(low, high - low),
validate_args=validate_args)
def __init__(self, low=0., loc=0., scale=1., validate_args=None):
self.low, self.loc, self.scale = promote_shapes(low, loc, scale)
base_loc = (loc - low) / scale
base_dist = _BaseTruncatedNormal(base_loc)
super(TruncatedNormal, self).__init__(base_dist,
AffineTransform(low, scale),
validate_args=validate_args)
def model(data):
alpha = numpyro.sample('alpha', dist.Uniform(0, 1))
with handlers.reparam(config={'loc': TransformReparam()}):
loc = numpyro.sample('loc', dist.TransformedDistribution(
dist.Uniform(0, 1).mask(False),
AffineTransform(0, alpha)))
numpyro.sample('obs', dist.Normal(loc, 0.1), obs=data)
def reparam_model(dim=10):
y = numpyro.sample('y', dist.Normal(0, 3))
with numpyro.handlers.reparam(config={'x': TransformReparam()}):
numpyro.sample(
'x',
dist.TransformedDistribution(dist.Normal(jnp.zeros(dim - 1), 1),
AffineTransform(0, jnp.exp(y / 2))))
def model(data):
alpha = numpyro.sample("alpha", dist.Uniform(0, 1))
with numpyro.handlers.reparam(config={"loc": TransformReparam()}):
loc = numpyro.sample(
"loc",
dist.TransformedDistribution(
dist.Uniform(0, 1).mask(False), AffineTransform(0, alpha)),
)
numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
def model():
with numpyro.plate_stack("plates", shape):
with numpyro.plate("particles", 100000):
return numpyro.sample(
"x",
dist.TransformedDistribution(
dist.Normal(jnp.zeros_like(loc), jnp.ones_like(scale)),
[AffineTransform(loc, scale),
ExpTransform()]).expand_by([100000]))
def model():
fn = dist.TransformedDistribution(
dist.Normal(jnp.zeros_like(loc), jnp.ones_like(scale)),
[AffineTransform(loc, scale), ExpTransform()]).expand(shape)
if event_shape:
fn = fn.to_event(len(event_shape)).expand_by([100000])
with numpyro.plate_stack("plates", batch_shape):
with numpyro.plate("particles", 100000):
return numpyro.sample("x", fn)
def __init__(self, alpha, scale=1., validate_args=None):
batch_shape = lax.broadcast_shapes(np.shape(scale), np.shape(alpha))
self.scale, self.alpha = np.broadcast_to(scale,
batch_shape), np.broadcast_to(
alpha, batch_shape)
base_dist = Exponential(self.alpha)
transforms = [ExpTransform(), AffineTransform(loc=0, scale=self.scale)]
super(Pareto, self).__init__(base_dist,
transforms,
validate_args=validate_args)
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)
def _get_transform(self):
loc = numpyro.param('{}_loc'.format(self.prefix), self._init_latent)
scale = numpyro.param('{}_scale'.format(self.prefix),
np.full(self.latent_size, self._init_scale),
constraint=constraints.positive)
return AffineTransform(loc, scale, domain=constraints.real_vector)
def reparam_model(dim=10):
y = numpyro.sample('y', dist.Normal(0, 3))
numpyro.sample(
'x',
dist.TransformedDistribution(dist.Normal(np.zeros(dim - 1), 1),
AffineTransform(0, np.exp(y / 2))))
def get_transform(self, params):
loc = params["{}_loc".format(self.prefix)]
scale = params["{}_scale".format(self.prefix)]
return IndependentTransform(AffineTransform(loc, scale), 1)
def get_transform(self, params):
loc = params['{}_loc'.format(self.prefix)]
scale = params['{}_scale'.format(self.prefix)]
return AffineTransform(loc, scale, domain=constraints.real_vector)
def _support(self, *args, **kwargs):
# support of the transformed distribution
_, loc, scale = self._parse_args(*args, **kwargs)
return AffineTransform(loc, scale, domain=self._support_mask).codomain