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, elbo, adam)
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 __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 _get_transform(self):
loc, scale = self._loc_scale()
return AffineTransform(loc, scale, domain=constraints.real_vector)
def reparam_model(dim=10):
y = sample('y', dist.Normal(0, 3))
sample('x', dist.TransformedDistribution(
dist.Normal(np.zeros(dim - 1), 1), AffineTransform(0, np.exp(y / 2))))
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