def __init__(self, samples, log_weights, validate_args=None):
self._samples = samples
self._log_weights = log_weights
sample_shape, weight_shape = samples.size(), log_weights.size()
if weight_shape > sample_shape or weight_shape != sample_shape[:len(weight_shape)]:
raise ValueError("The shape of ``log_weights`` ({}) must match "
"the leftmost shape of ``samples`` ({})".format(weight_shape, sample_shape))
self._aggregation_dim = log_weights.dim() - 1
event_shape = sample_shape[len(weight_shape):]
self._categorical = Categorical(logits=self._log_weights)
super(TorchDistribution, self).__init__(batch_shape=weight_shape[:-1],
event_shape=event_shape,
validate_args=validate_args)
def _finalize(self):
"""
Appends values collected in the samples/weights buffers to their
corresponding tensors.
"""
if not self._samples_buffer:
return
self._samples = self._append_from_buffer(self._samples,
self._samples_buffer)
self._log_weights = self._append_from_buffer(self._log_weights,
self._weights_buffer)
self._categorical = Categorical(logits=self._log_weights)
# Reset buffers.
self._samples_buffer, self._weights_buffer = [], []
def _finalize(self):
"""
Appends values collected in the samples/weights buffers to their
corresponding tensors.
"""
if not self._samples_buffer:
return
self._samples = self._append_from_buffer(self._samples, self._samples_buffer)
self._log_weights = self._append_from_buffer(self._log_weights, self._weights_buffer)
self._categorical = Categorical(logits=self._log_weights)
# Reset buffers.
self._samples_buffer, self._weights_buffer = [], []
def filter(self, value):
"""
Compute posterior over final state given a sequence of observations.
:param ~torch.Tensor value: A sequence of observations.
:return: A posterior distribution
over latent states at the final time step. ``result.logits`` can
then be used as ``initial_logits`` in a sequential Pyro model for
prediction.
:rtype: ~pyro.distributions.Categorical
"""
# Combine observation and transition factors.
value = value.unsqueeze(-1 - self.observation_dist.event_dim)
observation_logits = self.observation_dist.log_prob(value)
logp = self.transition_logits + observation_logits.unsqueeze(-2)
# Eliminate time dimension.
logp = _sequential_logmatmulexp(logp)
# Combine initial factor.
logp = (self.initial_logits.unsqueeze(-1) + logp).logsumexp(-2)
# Convert to a distribution.
return Categorical(logits=logp, validate_args=self._validate_args)
class Empirical(TorchDistribution):
r"""
Empirical distribution associated with the sampled data.
:param torch.Tensor samples: samples from the empirical distribution.
:param torch.Tensor log_weights: log weights (optional) corresponding
to the samples. The leftmost shape of ``log_weights`` must match
that of samples
"""
arg_constraints = {}
support = constraints.real
has_enumerate_support = True
def __init__(self, samples, log_weights, validate_args=None):
self._samples = samples
self._log_weights = log_weights
sample_shape, weight_shape = samples.size(), log_weights.size()
if weight_shape > sample_shape or weight_shape != sample_shape[:len(
weight_shape)]:
raise ValueError("The shape of ``log_weights`` ({}) must match "
"the leftmost shape of ``samples`` ({})".format(
weight_shape, sample_shape))
self._aggregation_dim = log_weights.dim() - 1
event_shape = sample_shape[len(weight_shape):]
self._categorical = Categorical(logits=self._log_weights)
super(TorchDistribution, self).__init__(batch_shape=weight_shape[:-1],
event_shape=event_shape,
validate_args=validate_args)
@property
def sample_size(self):
"""
Number of samples that constitute the empirical distribution.
:return int: number of samples collected.
"""
return self._log_weights.numel()
def sample(self, sample_shape=torch.Size()):
sample_idx = self._categorical.sample(sample_shape)
return self._samples[sample_idx]
def log_prob(self, value):
"""
Returns the log of the probability mass function evaluated at ``value``.
Note that this currently only supports scoring values with empty
``sample_shape``.
:param torch.Tensor value: scalar or tensor value to be scored.
"""
if self._validate_args:
if value.shape != self.batch_shape + self.event_shape:
raise ValueError(
"``value.shape`` must be {}".format(self.batch_shape +
self.event_shape))
if self.batch_shape:
value = value.unsqueeze(self._aggregation_dim)
selection_mask = self._samples.eq(value)
# Get a mask for all entries in the ``weights`` tensor
# that correspond to ``value``.
for _ in range(len(self.event_shape)):
selection_mask = selection_mask.min(dim=-1)[0]
selection_mask = selection_mask.type(self._categorical.probs.type())
return (self._categorical.probs * selection_mask).sum(dim=-1).log()
def _weighted_mean(self, value, keepdim=False):
weights = self._log_weights.reshape(
self._log_weights.size() +
torch.Size([1] * (value.dim() - self._log_weights.dim())))
dim = self._aggregation_dim
max_weight = weights.max(dim=dim, keepdim=True)[0]
relative_probs = (weights - max_weight).exp()
return (value * relative_probs).sum(
dim=dim, keepdim=keepdim) / relative_probs.sum(dim=dim,
keepdim=keepdim)
@property
def event_shape(self):
return self._event_shape
@property
def mean(self):
if self._samples.dtype in (torch.int32, torch.int64):
raise ValueError(
"Mean for discrete empirical distribution undefined. " +
"Consider converting samples to ``torch.float32`` " +
"or ``torch.float64``. If these are samples from a " +
"`Categorical` distribution, consider converting to a " +
"`OneHotCategorical` distribution.")
return self._weighted_mean(self._samples)
@property
def variance(self):
if self._samples.dtype in (torch.int32, torch.int64):
raise ValueError(
"Variance for discrete empirical distribution undefined. " +
"Consider converting samples to ``torch.float32`` " +
"or ``torch.float64``. If these are samples from a " +
"`Categorical` distribution, consider converting to a " +
"`OneHotCategorical` distribution.")
mean = self.mean.unsqueeze(self._aggregation_dim)
deviation_squared = torch.pow(self._samples - mean, 2)
return self._weighted_mean(deviation_squared)
@property
def log_weights(self):
return self._log_weights
def enumerate_support(self, expand=True):
# Empirical does not support batching, so expanding is a no-op.
return self._samples
class Empirical(TorchDistribution):
r"""
Empirical distribution associated with the sampled data. Note that the shape
requirement for `log_weights` is that its shape must match the leftmost shape
of `samples`. Samples are aggregated along the ``aggregation_dim``, which is
the rightmost dim of `log_weights`.
Example:
>>> emp_dist = Empirical(torch.randn(2, 3, 10), torch.ones(2, 3))
>>> emp_dist.batch_shape
torch.Size([2])
>>> emp_dist.event_shape
torch.Size([10])
>>> single_sample = emp_dist.sample()
>>> single_sample.shape
torch.Size([2, 10])
>>> batch_sample = emp_dist.sample((100,))
>>> batch_sample.shape
torch.Size([100, 2, 10])
>>> emp_dist.log_prob(single_sample).shape
torch.Size([2])
>>> # Vectorized samples cannot be scored by log_prob.
>>> with pyro.validation_enabled():
... emp_dist.log_prob(batch_sample).shape
Traceback (most recent call last):
...
ValueError: ``value.shape`` must be torch.Size([2, 10])
:param torch.Tensor samples: samples from the empirical distribution.
:param torch.Tensor log_weights: log weights (optional) corresponding
to the samples.
"""
arg_constraints = {}
support = constraints.real
has_enumerate_support = True
def __init__(self, samples, log_weights, validate_args=None):
self._samples = samples
self._log_weights = log_weights
sample_shape, weight_shape = samples.size(), log_weights.size()
if (
weight_shape > sample_shape
or weight_shape != sample_shape[: len(weight_shape)]
):
raise ValueError(
"The shape of ``log_weights`` ({}) must match "
"the leftmost shape of ``samples`` ({})".format(
weight_shape, sample_shape
)
)
self._aggregation_dim = log_weights.dim() - 1
event_shape = sample_shape[len(weight_shape) :]
self._categorical = Categorical(logits=self._log_weights)
super().__init__(
batch_shape=weight_shape[:-1],
event_shape=event_shape,
validate_args=validate_args,
)
@property
def sample_size(self):
"""
Number of samples that constitute the empirical distribution.
:return int: number of samples collected.
"""
return self._log_weights.numel()
def sample(self, sample_shape=torch.Size()):
sample_idx = self._categorical.sample(
sample_shape
) # sample_shape x batch_shape
# reorder samples to bring aggregation_dim to the front:
# batch_shape x num_samples x event_shape -> num_samples x batch_shape x event_shape
samples = (
self._samples.unsqueeze(0)
.transpose(0, self._aggregation_dim + 1)
.squeeze(self._aggregation_dim + 1)
)
# make sample_idx.shape compatible with samples.shape: sample_shape_numel x batch_shape x event_shape
sample_idx = sample_idx.reshape(
(-1,) + self.batch_shape + (1,) * len(self.event_shape)
)
sample_idx = sample_idx.expand((-1,) + samples.shape[1:])
return samples.gather(0, sample_idx).reshape(sample_shape + samples.shape[1:])
def log_prob(self, value):
"""
Returns the log of the probability mass function evaluated at ``value``.
Note that this currently only supports scoring values with empty
``sample_shape``.
:param torch.Tensor value: scalar or tensor value to be scored.
"""
if self._validate_args:
if value.shape != self.batch_shape + self.event_shape:
raise ValueError(
"``value.shape`` must be {}".format(
self.batch_shape + self.event_shape
)
)
if self.batch_shape:
value = value.unsqueeze(self._aggregation_dim)
selection_mask = self._samples.eq(value)
# Get a mask for all entries in the ``weights`` tensor
# that correspond to ``value``.
for _ in range(len(self.event_shape)):
selection_mask = selection_mask.min(dim=-1)[0]
selection_mask = selection_mask.type(self._categorical.probs.type())
return (self._categorical.probs * selection_mask).sum(dim=-1).log()
def _weighted_mean(self, value, keepdim=False):
weights = self._log_weights.reshape(
self._log_weights.size()
+ torch.Size([1] * (value.dim() - self._log_weights.dim()))
)
dim = self._aggregation_dim
max_weight = weights.max(dim=dim, keepdim=True)[0]
relative_probs = (weights - max_weight).exp()
return (value * relative_probs).sum(
dim=dim, keepdim=keepdim
) / relative_probs.sum(dim=dim, keepdim=keepdim)
@property
def event_shape(self):
return self._event_shape
@property
def mean(self):
if self._samples.dtype in (torch.int32, torch.int64):
raise ValueError(
"Mean for discrete empirical distribution undefined. "
+ "Consider converting samples to ``torch.float32`` "
+ "or ``torch.float64``. If these are samples from a "
+ "`Categorical` distribution, consider converting to a "
+ "`OneHotCategorical` distribution."
)
return self._weighted_mean(self._samples)
@property
def variance(self):
if self._samples.dtype in (torch.int32, torch.int64):
raise ValueError(
"Variance for discrete empirical distribution undefined. "
+ "Consider converting samples to ``torch.float32`` "
+ "or ``torch.float64``. If these are samples from a "
+ "`Categorical` distribution, consider converting to a "
+ "`OneHotCategorical` distribution."
)
mean = self.mean.unsqueeze(self._aggregation_dim)
deviation_squared = torch.pow(self._samples - mean, 2)
return self._weighted_mean(deviation_squared)
@property
def log_weights(self):
return self._log_weights
def enumerate_support(self, expand=True):
# Empirical does not support batching, so expanding is a no-op.
return self._samples
class Empirical(TorchDistribution):
r"""
Empirical distribution associated with the sampled data.
"""
arg_constraints = {}
support = constraints.real
has_enumerate_support = True
def __init__(self, validate_args=None):
self._samples = None
self._log_weights = None
self._categorical = None
self._samples_buffer = []
self._weights_buffer = []
super(TorchDistribution, self).__init__(batch_shape=torch.Size(),
validate_args=validate_args)
@staticmethod
def _append_from_buffer(tensor, buffer):
"""
Append values from the buffer to the finalized tensor, along the
leftmost dimension.
:param torch.Tensor tensor: tensor containing existing values.
:param list buffer: list of new values.
:return: tensor with new values appended at the bottom.
"""
buffer_tensor = torch.stack(buffer, dim=0)
return torch.cat([tensor, buffer_tensor], dim=0)
def _finalize(self):
"""
Appends values collected in the samples/weights buffers to their
corresponding tensors.
"""
if not self._samples_buffer:
return
self._samples = self._append_from_buffer(self._samples,
self._samples_buffer)
self._log_weights = self._append_from_buffer(self._log_weights,
self._weights_buffer)
self._categorical = Categorical(logits=self._log_weights)
# Reset buffers.
self._samples_buffer, self._weights_buffer = [], []
@property
def sample_size(self):
"""
Number of samples that constitute the empirical distribution.
:return int: number of samples collected.
"""
self._finalize()
if self._samples is None:
return 0
return self._samples.size(0)
def add(self, value, weight=None, log_weight=None):
"""
Adds a new data point to the sample. The values in successive calls to
``add`` must have the same tensor shape and size. Optionally, an
importance weight can be specified via ``log_weight`` or ``weight``
(default value of `1` is used if not specified).
:param torch.Tensor value: tensor to add to the sample.
:param torch.Tensor weight: log weight (optional) corresponding
to the sample.
:param torch.Tensor log_weight: weight (optional) corresponding
to the sample.
"""
if self._validate_args:
if weight is not None and log_weight is not None:
raise ValueError(
"Only one of ```weight`` or ``log_weight`` should be specified."
)
weight_type = value.new_empty(1).float().type() if value.dtype in (torch.int32, torch.int64) \
else value.type()
# Apply default weight of 1.0.
if log_weight is None and weight is None:
log_weight = torch.tensor(0.0).type(weight_type)
elif weight is not None and log_weight is None:
log_weight = math.log(weight)
if isinstance(log_weight, numbers.Number):
log_weight = torch.tensor(log_weight).type(weight_type)
if self._validate_args and log_weight.dim() > 0:
raise ValueError(
"``weight.dim() > 0``, but weight should be a scalar.")
# Seed the container tensors with the correct tensor types
if self._samples is None:
self._samples = value.new()
self._log_weights = log_weight.new()
# Append to the buffer list
self._samples_buffer.append(value)
self._weights_buffer.append(log_weight)
def sample(self, sample_shape=torch.Size()):
self._finalize()
idxs = self._categorical.sample(sample_shape=sample_shape)
return self._samples[idxs]
def log_prob(self, value):
"""
Returns the log of the probability mass function evaluated at ``value``.
Note that this currently only supports scoring values with empty
``sample_shape``, i.e. an arbitrary batched sample is not allowed.
:param torch.Tensor value: scalar or tensor value to be scored.
"""
if self._validate_args:
if value.size() != self.event_shape:
raise ValueError("``value.size()`` must be {}".format(
self.event_shape))
self._finalize()
selection_mask = self._samples.eq(value).contiguous().view(
self.sample_size, -1)
# Return -Inf if value is outside the support.
if not selection_mask.any():
return self._log_weights.new_zeros(torch.Size()).log()
idxs = torch.arange(self.sample_size)[selection_mask.min(dim=-1)[0]]
log_probs = self._categorical.log_prob(idxs)
return log_sum_exp(log_probs)
def _weighted_mean(self, value, dim=0):
weights = self._log_weights
for _ in range(value.dim() - 1):
weights = weights.unsqueeze(-1)
max_val = weights.max(dim)[0]
return max_val.exp() * (value *
(weights - max_val.unsqueeze(-1)).exp()).sum(
dim=dim)
@property
def event_shape(self):
self._finalize()
if self._samples is None:
return None
return self._samples.size()[1:]
@property
def mean(self):
self._finalize()
if self._samples.dtype in (torch.int32, torch.int64):
raise ValueError(
"Mean for discrete empirical distribution undefined. " +
"Consider converting samples to ``torch.float32`` " +
"or ``torch.float64``. If these are samples from a " +
"`Categorical` distribution, consider converting to a " +
"`OneHotCategorical` distribution.")
return self._weighted_mean(self._samples) / self._weighted_mean(
self._samples.new_tensor([1.]))
@property
def variance(self):
self._finalize()
if self._samples.dtype in (torch.int32, torch.int64):
raise ValueError(
"Variance for discrete empirical distribution undefined. " +
"Consider converting samples to ``torch.float32`` " +
"or ``torch.float64``. If these are samples from a " +
"`Categorical` distribution, consider converting to a " +
"`OneHotCategorical` distribution.")
deviation_squared = torch.pow(self._samples - self.mean, 2)
return self._weighted_mean(deviation_squared) / self._weighted_mean(
self._samples.new_tensor([1.]))
def get_samples_and_weights(self):
self._finalize()
return self._samples, self._log_weights
def enumerate_support(self):
self._finalize()
return self._samples
class Empirical(TorchDistribution):
r"""
Empirical distribution associated with the sampled data.
"""
arg_constraints = {}
support = constraints.real
has_enumerate_support = True
def __init__(self, validate_args=None):
self._samples = None
self._log_weights = None
self._categorical = None
self._samples_buffer = []
self._weights_buffer = []
super(TorchDistribution, self).__init__(batch_shape=torch.Size(), validate_args=validate_args)
@staticmethod
def _append_from_buffer(tensor, buffer):
"""
Append values from the buffer to the finalized tensor, along the
leftmost dimension.
:param torch.Tensor tensor: tensor containing existing values.
:param list buffer: list of new values.
:return: tensor with new values appended at the bottom.
"""
buffer_tensor = torch.stack(buffer, dim=0)
return torch.cat([tensor, buffer_tensor], dim=0)
def _finalize(self):
"""
Appends values collected in the samples/weights buffers to their
corresponding tensors.
"""
if not self._samples_buffer:
return
self._samples = self._append_from_buffer(self._samples, self._samples_buffer)
self._log_weights = self._append_from_buffer(self._log_weights, self._weights_buffer)
self._categorical = Categorical(logits=self._log_weights)
# Reset buffers.
self._samples_buffer, self._weights_buffer = [], []
@property
def sample_size(self):
"""
Number of samples that constitute the empirical distribution.
:return int: number of samples collected.
"""
self._finalize()
if self._samples is None:
return 0
return self._samples.size(0)
def add(self, value, weight=None, log_weight=None):
"""
Adds a new data point to the sample. The values in successive calls to
``add`` must have the same tensor shape and size. Optionally, an
importance weight can be specified via ``log_weight`` or ``weight``
(default value of `1` is used if not specified).
:param torch.Tensor value: tensor to add to the sample.
:param torch.Tensor weight: log weight (optional) corresponding
to the sample.
:param torch.Tensor log_weight: weight (optional) corresponding
to the sample.
"""
if self._validate_args:
if weight is not None and log_weight is not None:
raise ValueError("Only one of ```weight`` or ``log_weight`` should be specified.")
weight_type = value.new_empty(1).float().type() if value.dtype in (torch.int32, torch.int64) \
else value.type()
# Apply default weight of 1.0.
if log_weight is None and weight is None:
log_weight = torch.tensor(0.0).type(weight_type)
elif weight is not None and log_weight is None:
log_weight = math.log(weight)
if isinstance(log_weight, numbers.Number):
log_weight = torch.tensor(log_weight).type(weight_type)
if self._validate_args and log_weight.dim() > 0:
raise ValueError("``weight.dim() > 0``, but weight should be a scalar.")
# Seed the container tensors with the correct tensor types
if self._samples is None:
self._samples = value.new()
self._log_weights = log_weight.new()
# Append to the buffer list
self._samples_buffer.append(value)
self._weights_buffer.append(log_weight)
def sample(self, sample_shape=torch.Size()):
self._finalize()
idxs = self._categorical.sample(sample_shape=sample_shape)
return self._samples[idxs]
def log_prob(self, value):
"""
Returns the log of the probability mass function evaluated at ``value``.
Note that this currently only supports scoring values with empty
``sample_shape``, i.e. an arbitrary batched sample is not allowed.
:param torch.Tensor value: scalar or tensor value to be scored.
"""
if self._validate_args:
if value.size() != self.event_shape:
raise ValueError("``value.size()`` must be {}".format(self.event_shape))
self._finalize()
selection_mask = self._samples.eq(value).contiguous().view(self.sample_size, -1)
# Return -Inf if value is outside the support.
if not selection_mask.any():
return self._log_weights.new_zeros(torch.Size()).log()
idxs = torch.arange(self.sample_size)[selection_mask.min(dim=-1)[0]]
log_probs = self._categorical.log_prob(idxs)
return log_sum_exp(log_probs)
def _weighted_mean(self, value, dim=0):
weights = self._log_weights
for _ in range(value.dim() - 1):
weights = weights.unsqueeze(-1)
max_val = weights.max(dim)[0]
return max_val.exp() * (value * (weights - max_val.unsqueeze(-1)).exp()).sum(dim=dim)
@property
def event_shape(self):
self._finalize()
if self._samples is None:
return None
return self._samples.size()[1:]
@property
def mean(self):
self._finalize()
if self._samples.dtype in (torch.int32, torch.int64):
raise ValueError("Mean for discrete empirical distribution undefined. " +
"Consider converting samples to ``torch.float32`` " +
"or ``torch.float64``. If these are samples from a " +
"`Categorical` distribution, consider converting to a " +
"`OneHotCategorical` distribution.")
return self._weighted_mean(self._samples) / self._weighted_mean(self._samples.new_tensor([1.]))
@property
def variance(self):
self._finalize()
if self._samples.dtype in (torch.int32, torch.int64):
raise ValueError("Variance for discrete empirical distribution undefined. " +
"Consider converting samples to ``torch.float32`` " +
"or ``torch.float64``. If these are samples from a " +
"`Categorical` distribution, consider converting to a " +
"`OneHotCategorical` distribution.")
deviation_squared = torch.pow(self._samples - self.mean, 2)
return self._weighted_mean(deviation_squared) / self._weighted_mean(self._samples.new_tensor([1.]))
def get_samples_and_weights(self):
self._finalize()
return self._samples, self._log_weights
def enumerate_support(self):
self._finalize()
return self._samples