def test_prob_and_log_prob(self):
# test default group_ndims
distrib = Mock(
is_reparameterized=True,
log_prob=Mock(return_value=tf.constant(1.)),
prob=Mock(return_value=tf.constant(2.)),
)
t = StochasticTensor(distrib, tf.constant(0.))
given = t.tensor
exp_1 = np.exp(1.).astype(np.float32)
with self.test_session():
self.assertEqual(t.log_prob().eval(), 1.)
self.assertEqual(t.log_prob().eval(), 1.)
np.testing.assert_allclose(t.prob().eval(), exp_1)
np.testing.assert_allclose(t.prob().eval(), exp_1)
self.assertEqual(distrib.log_prob.call_args_list, [((given, 0), {
'name': None
})])
self.assertEqual(distrib.prob.call_args_list, [])
# test group_ndims equal to default
distrib.log_prob.reset_mock()
distrib.prob.reset_mock()
with self.test_session():
self.assertEqual(t.log_prob(group_ndims=0).eval(), 1.)
np.testing.assert_allclose(t.prob(group_ndims=0).eval(), exp_1)
distrib.log_prob.assert_not_called()
distrib.prob.assert_not_called()
# test group_ndims different from default
distrib.log_prob.reset_mock()
distrib.prob.reset_mock()
with self.test_session():
self.assertEqual(t.log_prob(group_ndims=1).eval(), 1.)
np.testing.assert_allclose(t.prob(group_ndims=2).eval(), exp_1)
self.assertEqual(distrib.log_prob.call_args_list, [((given, 1), {
'name': None
}), ((given, 2), )])
self.assertEqual(distrib.prob.call_args_list, [])
# test use dynamic group_ndims
t = StochasticTensor(distrib,
tf.constant(0.),
group_ndims=tf.constant(1, dtype=tf.int32))
given = t.tensor
distrib.log_prob.reset_mock()
distrib.prob.reset_mock()
with self.test_session():
self.assertEqual(t.log_prob(group_ndims=t.group_ndims).eval(), 1.)
self.assertEqual(t.log_prob(group_ndims=t.group_ndims).eval(), 1.)
np.testing.assert_allclose(
t.prob(group_ndims=t.group_ndims).eval(), exp_1)
np.testing.assert_allclose(
t.prob(group_ndims=t.group_ndims).eval(), exp_1)
self.assertEqual(distrib.log_prob.call_args_list,
[((given, t.group_ndims), {
'name': None
})])
self.assertEqual(distrib.prob.call_args_list, [])
def sample(
self,
n_samples=None,
is_reparameterized=None,
group_ndims=0,
compute_density=False,
name=None,
):
from tfsnippet.stochastic import StochasticTensor
if n_samples is None or n_samples < 2:
n_samples = 2
with tf.name_scope(name=name, default_name="sample"):
samples = self._distribution.sample(n_samples)
samples = tf.reduce_mean(samples, axis=0)
t = StochasticTensor(
distribution=self,
tensor=samples,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=self.is_reparameterized,
)
if compute_density:
with tf.name_scope("compute_prob_and_log_prob"):
log_p = t.log_prob()
t._self_prob = tf.exp(log_p)
return t
def sample(self,
n_samples=1024,
is_reparameterized=None,
group_ndims=0,
compute_density=False,
name=None):
from tfsnippet.stochastic import StochasticTensor
if n_samples is None:
n_samples = 1
n_samples_is_none = True
else:
n_samples_is_none = False
with tf.name_scope(name=name, default_name='sample'):
noise = self.normal.sample(n_samples=n_samples)
noise = tf.transpose(
noise, [1, 0, 2]) # window_length * n_samples * z_dim
noise = tf.truncated_normal(tf.shape(noise))
# n_sample, batchsize, z_dim
time_indices_shape = tf.convert_to_tensor(
[n_samples, tf.shape(self.input_q)[1], self.z_dim])
samples = tf.scan(
fn=self.sample_step,
elems=(noise, self.input_q), #100*samples*3; 100*50*500
initializer=(tf.zeros(time_indices_shape),
tf.zeros(time_indices_shape),
tf.ones(time_indices_shape)),
back_prop=False)[
0] # time_step * n_samples * batch_size * z_dim
samples = tf.transpose(
samples,
[1, 2, 0, 3]) # n_samples * batch_size * time_step * z_dim
if n_samples_is_none:
t = StochasticTensor(
distribution=self,
tensor=tf.reduce_mean(samples, axis=0),
n_samples=1,
group_ndims=group_ndims,
is_reparameterized=self.is_reparameterized)
else:
t = StochasticTensor(
distribution=self,
tensor=samples,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=self.is_reparameterized)
if compute_density:
with tf.name_scope('compute_prob_and_log_prob'):
log_p = t.log_prob()
t._self_prob = tf.exp(log_p)
return t
def sample(self,
n_samples=None,
is_reparameterized=None,
group_ndims=0,
compute_density=None,
name=None):
from tfsnippet.stochastic import StochasticTensor
if is_reparameterized and not self.is_reparameterized:
raise RuntimeError('Distribution is not re-parameterized')
elif is_reparameterized is False and self.is_reparameterized:
@contextlib.contextmanager
def set_is_reparameterized():
try:
self._distribution._is_reparameterized = False
yield False
finally:
self._distribution._is_reparameterized = True
else:
@contextlib.contextmanager
def set_is_reparameterized():
yield self.is_reparameterized
with tf.name_scope(name=name, default_name='sample'):
with set_is_reparameterized() as is_reparameterized:
samples = self._distribution.sample(n_samples=n_samples)
t = StochasticTensor(
distribution=self,
tensor=samples,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
)
if compute_density:
with tf.name_scope('compute_prob_and_log_prob'):
log_p = t.log_prob()
t._self_prob = tf.exp(log_p)
return t
def test_construction(self):
distrib = Mock(is_reparameterized=True, is_continuous=True)
samples = tf.constant(12345678., dtype=tf.float32)
# test basic construction
t = StochasticTensor(distrib, samples, n_samples=1, group_ndims=2)
self.assertIs(t.distribution, distrib)
self.assertTrue(t.is_reparameterized)
self.assertTrue(t.is_continuous)
self.assertEqual(t.n_samples, 1)
self.assertEqual(t.group_ndims, 2)
self.assertEqual(t.dtype, tf.float32)
self.assertIsInstance(t.tensor, tf.Tensor)
with self.test_session():
self.assertEqual(t.eval(), 12345678.)
self.assertEqual(t.tensor.eval(), 12345678)
# test initializing from TensorWrapper
samples = tf.constant(1.)
t = StochasticTensor(Mock(is_reparameterized=False),
_MyTensorWrapper(samples))
self.assertIs(t.tensor, samples)
# test specifying is_reparameterized
t = StochasticTensor(Mock(is_reparameterized=True),
tf.constant(0.),
is_reparameterized=False)
self.assertFalse(t.is_reparameterized)
# test construction with dynamic group_ndims
t = StochasticTensor(distrib,
samples,
group_ndims=tf.constant(2, dtype=tf.int32))
with self.test_session():
self.assertEqual(t.group_ndims.eval(), 2)
# test construction log_prob
t = StochasticTensor(distrib, tf.constant(0.), log_prob=123.)
with self.test_session() as sess:
self.assertEqual(sess.run(t.log_prob()), 123.)
# test construction with bad dynamic group_ndims
t = StochasticTensor(distrib,
samples,
group_ndims=tf.constant(-1, dtype=tf.int32))
with self.test_session():
with pytest.raises(Exception,
match='group_ndims must be non-negative'):
_ = t.group_ndims.eval()
# test construction with dynamic n_samples
t = StochasticTensor(distrib,
samples,
n_samples=tf.constant(2, dtype=tf.int32))
with self.test_session():
self.assertEqual(t.n_samples.eval(), 2)
# test construction with bad dynamic n_samples
t = StochasticTensor(distrib,
samples,
n_samples=tf.constant(0, dtype=tf.int32))
with self.test_session():
with pytest.raises(Exception, match='n_samples must be positive'):
_ = t.n_samples.eval()
def add(self,
name,
distribution,
n_samples=None,
group_ndims=0,
is_reparameterized=None,
transform=None):
"""
Add a stochastic node to the network.
A :class:`StochasticTensor` will be created for this node.
If `name` exists in `observed` dict, its value will be used as the
observation of this node. Otherwise samples will be taken from
`distribution`.
Args:
name (str): Name of the stochastic node.
distribution (Distribution or zhusuan.distributions.Distribution):
Distribution where the samples should be taken from.
n_samples (int or tf.Tensor): Number of samples to take.
If specified, `n_samples` will be taken, with a dedicated
sampling dimension ``[n_samples]`` at the front.
If not specified, just one sample will be taken, without the
dedicated dimension.
group_ndims (int or tf.Tensor): Number of dimensions at the end of
``[n_samples] + batch_shape`` to be considered as events group.
(default 0)
is_reparameterized: Whether or not the re-parameterization trick
should be applied? (default :obj:`None`, following the setting
of `distribution`)
transform ((Tensor, Tensor) -> (tf.Tensor, tf.Tensor)):
The function to transform (x, log_p) to (x', log_p').
If specified, a :class:`StochasticTensor` will be sampled,
then transformed, then wrapped by a :class:`StochasticTensor`
with :class:`TransformedDistribution`.
Returns:
StochasticTensor: The sampled stochastic tensor.
Raises:
TypeError: If `name` is not a str, or `distribution` is a
:class:`TransformedDistribution`.
KeyError: If :class:`StochasticTensor` with `name` already exists.
ValueError: If `transform` cannot be applied.
See Also:
:meth:`tfsnippet.distributions.Distribution.sample`
"""
if not isinstance(name, six.string_types):
raise TypeError('`name` must be a str')
if name in self._stochastic_tensors:
raise KeyError(
'StochasticTensor with name {!r} already exists in '
'the BayesianNet. Names must be unique.'.format(name))
if isinstance(distribution, TransformedDistribution):
raise TypeError('Cannot add `TransformedDistribution`.')
if transform is not None and \
(not distribution.is_continuous or
not distribution.is_reparameterized or
is_reparameterized is False):
raise ValueError('`transform` can only be applied on continuous, '
're-parameterized variables.')
if transform is not None and name in self._observed:
raise ValueError('`observed` variable cannot be transformed.')
distribution = as_distribution(distribution)
if name in self._observed:
t = StochasticTensor(
distribution=distribution,
tensor=self._observed[name],
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
)
else:
t = distribution.sample(
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
)
assert (isinstance(t, StochasticTensor))
# do transformation
if transform is not None:
t_log_p = t.log_prob()
ft, ft_log_p = transform(t, t_log_p)
ft = tf.convert_to_tensor(ft)
ft_log_p = tf.convert_to_tensor(ft_log_p)
if not ft.dtype.is_floating:
raise ValueError('The transformed samples must be '
'continuous: got {!r}'.format(ft))
t = StochasticTensor(distribution=TransformedDistribution(
origin=t,
transformed=ft,
transformed_log_p=ft_log_p,
is_reparameterized=t.is_reparameterized,
is_continuous=True),
tensor=ft,
n_samples=t.n_samples,
group_ndims=t.group_ndims,
is_reparameterized=t.is_reparameterized)
self._stochastic_tensors[name] = t
return t