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 __init__(self,
distribution,
tensor,
n_samples=None,
group_ndims=0,
is_reparameterized=None,
flow_origin=None,
log_prob=None):
"""
Construct the :class:`StochasticTensor`.
Args:
distribution (tfsnippet.distributions.Distribution): The
distribution of this :class:`StochasticTensor`.
tensor (tf.Tensor or TensorWrapper): The samples or observations
of this :class:`StochasticTensor`.
n_samples (tf.Tensor or int): The number of samples taken in
:class:`Distribution.sample`. If not :obj:`None`, the first
dimension of `tensor` should be the sampling dimension.
group_ndims (int or tf.Tensor): The number of dimensions to be
considered as events group in samples. (default 0)
is_reparameterized (bool): Whether or not the samples are
re-parameterized? If not specified, will inherit from
:attr:`tfsnippet.distributions.Distribution.is_reparameterized`.
log_prob (Tensor or None): Pre-computed log-density of `tensor`,
given `group_ndims`.
flow_origin (StochasticTensor): The original stochastic tensor
from the base distribution of a
:class:`tfsnippet.FlowDistribution`.
"""
from tfsnippet.utils import TensorArgValidator, validate_group_ndims_arg
if is_reparameterized is None:
is_reparameterized = distribution.is_reparameterized
if log_prob is not None and not is_tensor_object(log_prob):
log_prob = tf.convert_to_tensor(log_prob)
n_samples = validate_n_samples_arg(n_samples, 'n_samples')
if n_samples is not None:
with tf.name_scope('validate_n_samples'):
validator = TensorArgValidator('n_samples')
n_samples = validator.require_non_negative(
validator.require_int32(n_samples))
group_ndims = validate_group_ndims_arg(group_ndims)
super(StochasticTensor, self).__init__()
self._self_distribution = distribution
self._self_tensor = tf.convert_to_tensor(tensor)
self._self_n_samples = n_samples
self._self_group_ndims = group_ndims
self._self_is_reparameterized = is_reparameterized
self._self_flow_origin = flow_origin
self._self_log_prob = log_prob
self._self_prob = None
def reduce_group_ndims(operation, tensor, group_ndims, name=None):
"""
Reduce the last `group_ndims` dimensions in `tensor`, using `operation`.
In :class:`~tfsnippet.distributions.Distribution`, when computing the
(log-)densities of certain `tensor`, the last few dimensions
may represent a group of events, thus should be accounted together.
This method can be used to reduce these dimensions, for example:
.. code-block:: python
log_prob = reduce_group_ndims(tf.reduce_sum, log_prob, group_ndims)
prob = reduce_group_ndims(tf.reduce_prod, log_prob, group_ndims)
Args:
operation: The operation for reducing the last `group_ndims`
dimensions. It must receive `tensor` as the 1st argument, and
`axis` as the 2nd argument.
tensor: The tensor to be reduced.
group_ndims: The number of dimensions at the end of `tensor` to be
reduced. If it is a constant integer and is zero, then no
operation will take place.
name: TensorFlow name scope of the graph nodes. (default
"reduce_group_ndims")
Returns:
tf.Tensor: The reduced tensor.
Raises:
ValueError: If `group_ndims` cannot be validated by
:meth:`validate_group_ndims`.
"""
group_ndims = validate_group_ndims_arg(group_ndims)
with tf.name_scope(name, default_name='reduce_group_ndims'):
if is_tensor_object(group_ndims):
tensor = tf.cond(
group_ndims > 0,
lambda: operation(tensor, tf.range(-group_ndims, 0)),
lambda: tensor)
else:
if group_ndims > 0:
tensor = operation(tensor, tf.range(-group_ndims, 0))
return tensor
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)