def sample(self,
n_samples=None,
group_ndims=0,
is_reparameterized=None,
compute_density=None,
name=None):
self._validate_sample_is_reparameterized_arg(is_reparameterized)
if is_reparameterized is None:
is_reparameterized = self.is_reparameterized
with tf.name_scope(name,
default_name=spt.utils.get_default_scope_name(
'sample', self)):
origin_z = self.pz.sample(n_samples=n_samples,
is_reparameterized=is_reparameterized,
compute_density=False)
z = origin_z
for i in range(self._mcmc_iterator):
e_z, grad_e_z, z_prime = self.get_sgld_proposal(z)
e_z_prime, grad_e_z_prime, _ = self.get_sgld_proposal(z_prime)
log_q_zprime_z = tf.reduce_sum(
tf.square(z_prime - z + self._mcmc_alpha * grad_e_z),
axis=-1)
log_q_zprime_z *= -1. / (4 * self._mcmc_alpha)
log_q_z_zprime = tf.reduce_sum(
tf.square(z - z_prime + self._mcmc_alpha * grad_e_z_prime),
axis=-1)
log_q_z_zprime *= -1. / (4 * self._mcmc_alpha)
log_ratio_1 = -e_z_prime + e_z # log [p(z_prime) / p(z)]
log_ratio_2 = log_q_z_zprime - log_q_zprime_z # log [q(z | z_prime) / q(z_prime | z)]
# print(log_ratio_1.mean().item(), log_ratio_2.mean().item())
ratio = tf.clip_by_value(tf.exp(log_ratio_1 + log_ratio_2),
0.0, 1.0)
# print(ratio.mean().item())
rnd_u = tf.random.normal(shape=ratio.shape)
mask = tf.cast(tf.less(rnd_u, ratio), tf.float32)
mask = tf.expand_dims(mask, axis=-1)
z = (z_prime * mask + z * (1 - mask))
t = spt.StochasticTensor(distribution=self,
tensor=z,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized)
return t
def sample(self, n_samples=None, group_ndims=0, is_reparameterized=None,
compute_density=None, name=None):
self._validate_sample_is_reparameterized_arg(is_reparameterized)
if is_reparameterized is None:
is_reparameterized = self.is_reparameterized
with tf.name_scope(name, default_name=spt.utils.get_default_scope_name('sample', self)):
t = spt.StochasticTensor(
distribution=self,
tensor=self.mean,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized
)
return t