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