def _sample_latent(self, *args, **kwargs):
init_loc = self._init_latent
loc = param('{}_loc'.format(self.prefix), init_loc)
scale = param('{}_scale'.format(self.prefix),
np.ones(np.size(self._init_latent)),
constraint=constraints.positive)
# TODO: Support for `.to_event()` to treat the batch dim as event dim.
return sample("_{}_latent".format(self.prefix),
dist.Normal(loc, scale))
def guide(batch, **kwargs):
encode = kwargs['encode']
encoder_params = param('encoder', None)
batch = np.reshape(batch, (batch.shape[0], -1))
z_loc, z_std = encode(encoder_params, batch)
z = sample('z', dist.Normal(z_loc, z_std))
return z
def model(batch, **kwargs):
decode = kwargs['decode']
decoder_params = param('decoder', None)
z_dim = kwargs['z_dim']
batch = np.reshape(batch, (batch.shape[0], -1))
z = sample('z', dist.Normal(np.zeros((z_dim,)), np.ones((z_dim,))))
img_loc = decode(decoder_params, z)
return sample('obs', dist.Bernoulli(img_loc), obs=batch)
def _get_transform(self):
flows = []
for i in range(self.num_flows):
arn_params = param('{}_arn__{}'.format(self.prefix, i), None)
if i > 0:
flows.append(
PermuteTransform(np.arange(self.latent_size)[::-1]))
flows.append(
InverseAutoregressiveTransform(self.arns[i], arn_params))
return ComposeTransform(flows)
def guide():
alpha = param('alpha', 0.5, constraint=constraints.unit_interval)
param('loc', 0, constraint=constraints.interval(0, alpha))
def model(data):
# NB: model's constraints will play no effect
loc = param('loc', 0., constraint=constraints.interval(0, 0.5))
sample('obs', dist.Normal(loc, 0.1), obs=data)
def guide():
alpha_q = param("alpha_q", 1.0, constraint=constraints.positive)
beta_q = param("beta_q", 1.0, constraint=constraints.positive)
sample("beta", dist.Beta(alpha_q, beta_q))
def guide():
guide_loc = param("guide_loc", 0.)
guide_scale = np.exp(param("guide_scale_log", 0.))
sample("loc", dist.Normal(guide_loc, guide_scale))
def model(data):
mean = param('mean', 0.)
std = param('std', 1., constraint=constraints.positive)
return sample('obs', dist.Normal(mean, std), obs=data)
def _loc_scale(self):
loc = param('{}_loc'.format(self.prefix), None)
scale = param('{}_scale'.format(self.prefix),
None,
constraint=constraints.positive)
return loc, scale
def model(data):
alpha = sample('alpha', dist.Uniform(0, 1))
loc = param('loc', 0., constraint=constraints.interval(0., alpha))
sample('obs', dist.Normal(loc, 0.1), obs=data)