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 test_equality(self):
distrib = Mock(is_reparameterized=False)
samples = tf.constant(0.)
t = StochasticTensor(distrib, samples)
self.assertEqual(t, t)
self.assertEqual(hash(t), hash(t))
self.assertNotEqual(StochasticTensor(distrib, samples), t)
def sample(self,
n_samples=None,
group_ndims=0,
is_reparameterized=None,
compute_density=None,
name=None):
group_ndims = validate_group_ndims_arg(group_ndims)
if not compute_density and compute_density is not None:
raise RuntimeError('`FlowDistribution` requires `compute_prob` '
'not to be False.')
with tf.name_scope(name, default_name='FlowDistribution.sample'):
# sample from the base distribution
x = self._distribution.sample(
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
compute_density=True)
# now do the transformation
is_reparameterized = x.is_reparameterized
y, log_det = self._flow.transform(x) # y, log |dy/dx|
if not is_reparameterized:
y = tf.stop_gradient(y) # important!
# compose the transformed tensor
return StochasticTensor(
distribution=self,
tensor=y,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
# log p(y) = log p(x) - log |dy/dx|
log_prob=x.log_prob() - log_det)
def sample(self,
n_samples=None,
is_reparameterized=None,
group_ndims=0,
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)
return StochasticTensor(
distribution=self,
tensor=samples,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
)
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='DiscretizedLogistic.sample'):
# sample from uniform distribution
sample_shape = self.batch_shape
static_sample_shape = self.get_batch_shape()
if n_samples is not None:
sample_shape = tf.concat([[n_samples], sample_shape], 0)
static_sample_shape = tf.TensorShape(
[None if is_tensor_object(n_samples) else n_samples]). \
concatenate(static_sample_shape)
u = tf.random_uniform(shape=sample_shape,
minval=self._epsilon,
maxval=1. - self._epsilon,
dtype=self._param_dtype)
u.set_shape(static_sample_shape)
# inverse CDF of the logistic
inverse_logistic_cdf = maybe_check_numerics(
tf.log(u) - tf.log(1. - u), 'inverse_logistic_cdf')
# obtain the actual sample
scale = maybe_check_numerics(tf.exp(self.log_scale, name='scale'),
'scale')
sample = self.mean + scale * inverse_logistic_cdf
if self.discretize_sample:
sample = self._discretize(sample)
sample = maybe_check_numerics(sample, 'sample')
sample = convert_to_tensor_and_cast(sample, self.dtype)
if not is_reparameterized:
sample = tf.stop_gradient(sample)
t = StochasticTensor(distribution=self,
tensor=sample,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized)
# compute the density
if compute_density:
compute_density_immediately(t)
return t
def sample(self,
n_samples=None,
group_ndims=0,
is_reparameterized=None,
name=None):
return StochasticTensor(
self,
x_samples,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
)
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,
group_ndims=0,
is_reparameterized=None,
compute_density=None,
name=None):
self._validate_sample_is_reparameterized_arg(is_reparameterized)
#######################################################################
# slow routine: generate the mixture by one_hot * stack([c.sample()]) #
#######################################################################
with tf.name_scope(name or 'Mixture.sample'):
cat = self.categorical.sample(n_samples, group_ndims=0)
mask = tf.one_hot(cat,
self.n_components,
dtype=self.dtype,
axis=-1)
if self.value_ndims > 0:
static_shape = (mask.get_shape().as_list() +
[1] * self.value_ndims)
dynamic_shape = concat_shapes(
[get_shape(mask), [1] * self.value_ndims])
mask = tf.reshape(mask, dynamic_shape)
mask.set_shape(static_shape)
mask = tf.stop_gradient(mask)
# derive the mixture samples
c_samples = [
c.sample(n_samples, group_ndims=0) for c in self.components
]
samples = tf.reduce_sum(
mask * tf.stack(c_samples, axis=-self.value_ndims - 1),
axis=-self.value_ndims - 1)
if not self.is_reparameterized:
samples = tf.stop_gradient(samples)
t = StochasticTensor(distribution=self,
tensor=samples,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized)
if compute_density:
compute_density_immediately(t)
return t
def test_is_tensor_object(self):
for obj in [
tf.constant(0.), # type: tf.Tensor
tf.get_variable('x', dtype=tf.float32, shape=()),
TensorWrapper(),
StochasticTensor(Mock(is_reparameterized=False),
tf.constant(0.))
]:
self.assertTrue(
is_tensor_object(obj),
msg='{!r} should be interpreted as a tensor object'.format(
obj))
for obj in [1, '', object(), None, True, (), {}, [], np.zeros([1])]:
self.assertFalse(
is_tensor_object(obj),
msg='{!r} should not be interpreted as a tensor object'.format(
obj))
def sample(self,
n_samples=None,
group_ndims=0,
is_reparameterized=None,
compute_density=None,
name=None):
group_ndims = validate_group_ndims_arg(group_ndims)
if not compute_density and compute_density is not None:
raise RuntimeError('`FlowDistribution` requires `compute_prob` '
'not to be False.')
with tf.name_scope(name, default_name='FlowDistribution.sample'):
# x and log p(x)
ndims_diff = (self.flow.x_value_ndims -
self.base_distribution.value_ndims)
x = self._distribution.sample(
n_samples=n_samples,
group_ndims=ndims_diff,
is_reparameterized=is_reparameterized,
compute_density=True)
log_px = x.log_prob()
# y, log |dy/dx|
is_reparameterized = x.is_reparameterized
y, log_det = self._flow.transform(x)
if not is_reparameterized:
y = tf.stop_gradient(y) # important!
# compute log p(y) = log p(x) - log |dy/dx|
# and then apply `group_ndims` on log p(y)
log_py = reduce_group_ndims(tf.reduce_sum, log_px - log_det,
group_ndims)
# compose the transformed tensor
return StochasticTensor(distribution=self,
tensor=y,
n_samples=n_samples,
group_ndims=group_ndims,
is_reparameterized=is_reparameterized,
log_prob=FlowDistributionDerivedTensor(
tensor=log_py, flow_origin=x),
flow_origin=x)
def log_prob(self, given, group_ndims=0, name=None):
given = tf.convert_to_tensor(given)
with tf.name_scope(name,
default_name='FlowDistribution.log_prob',
values=[given]):
# x, log |dx/dy|
x, log_det = self._flow.inverse_transform(given)
# log p(x)
ndims_diff = (self.flow.x_value_ndims -
self.base_distribution.value_ndims)
log_px = self._distribution.log_prob(x, group_ndims=ndims_diff)
# compute log p(y) = log p(x) + log |dx/dy|,
# and then apply `group_ndims` on log p(x)
log_py = reduce_group_ndims(tf.reduce_sum, log_px + log_det,
group_ndims)
return FlowDistributionDerivedTensor(
tensor=log_py,
flow_origin=StochasticTensor(distribution=self.base_distribution,
tensor=x))
def sample(self,
n_samples=None,
is_reparameterized=None,
group_ndims=0,
compute_density=None,
name=None):
from tfsnippet.stochastic import StochasticTensor
self._validate_sample_is_reparameterized_arg(is_reparameterized)
if 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._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:
compute_density_immediately(t)
return t
def add(self,
name,
distribution,
n_samples=None,
group_ndims=0,
is_reparameterized=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: If observation is not given for `name`, this
argument will be used to determine whether or not
re-parameterization trick should be applied when taking
samples from `distribution` (if not specified, use
`distribution.is_reparameterized`).
If observation is given for `name`, and this argument is
set to :obj:`True`, it will be used to validate the observation.
If this argument is set to :obj:`False`, `tf.stop_gradient`
will be applied on the observation.
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,
or `is_reparameterized = True`, but the observation is not
re-parameterized.
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))
distribution = as_distribution(distribution)
if name in self._observed:
ob_tensor = self._observed[name]
if isinstance(ob_tensor, StochasticTensor):
if is_reparameterized and not ob_tensor.is_reparameterized:
raise ValueError(
'`is_reparameterized` is True, but the observation '
'for `{}` is not re-parameterized: {}'.format(
name, ob_tensor))
if is_reparameterized is None:
is_reparameterized = ob_tensor.is_reparameterized
if not is_reparameterized:
ob_tensor = tf.stop_gradient(ob_tensor)
t = StochasticTensor(
distribution=distribution,
tensor=ob_tensor,
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))
self._stochastic_tensors[name] = t
return t
def add(self, name, distribution, n_samples=None, group_ndims=0,
is_reparameterized=None, flow=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`)
flow (BaseFlow): If specified, transform `distribution` by `flow`.
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 flow is not None and name in self._observed and \
not flow.explicitly_invertible:
raise TypeError('The observed variable {!r} expects `flow` to be '
'explicitly invertible, but it is not: {!r}.'.
format(name, flow))
distribution = as_distribution(distribution)
if flow is not None:
distribution = FlowDistribution(distribution, flow)
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))
self._stochastic_tensors[name] = t
return t
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 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 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
def test_repr(self):
t = StochasticTensor(
Mock(is_reparameterized=False),
Mock(spec=tf.Tensor, __repr__=Mock(return_value='repr_output')))
self.assertEqual(repr(t), 'StochasticTensor(repr_output)')