|
|
|
|
|
import warnings |
|
|
from typing import Optional |
|
|
from typing_extensions import deprecated |
|
|
|
|
|
import torch |
|
|
from torch import Tensor |
|
|
from torch.distributions import constraints |
|
|
from torch.distributions.utils import lazy_property |
|
|
from torch.types import _size |
|
|
|
|
|
|
|
|
__all__ = ["Distribution"] |
|
|
|
|
|
|
|
|
class Distribution: |
|
|
r""" |
|
|
Distribution is the abstract base class for probability distributions. |
|
|
""" |
|
|
|
|
|
has_rsample = False |
|
|
has_enumerate_support = False |
|
|
_validate_args = __debug__ |
|
|
|
|
|
@staticmethod |
|
|
def set_default_validate_args(value: bool) -> None: |
|
|
""" |
|
|
Sets whether validation is enabled or disabled. |
|
|
|
|
|
The default behavior mimics Python's ``assert`` statement: validation |
|
|
is on by default, but is disabled if Python is run in optimized mode |
|
|
(via ``python -O``). Validation may be expensive, so you may want to |
|
|
disable it once a model is working. |
|
|
|
|
|
Args: |
|
|
value (bool): Whether to enable validation. |
|
|
""" |
|
|
if value not in [True, False]: |
|
|
raise ValueError |
|
|
Distribution._validate_args = value |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
batch_shape: torch.Size = torch.Size(), |
|
|
event_shape: torch.Size = torch.Size(), |
|
|
validate_args: Optional[bool] = None, |
|
|
): |
|
|
self._batch_shape = batch_shape |
|
|
self._event_shape = event_shape |
|
|
if validate_args is not None: |
|
|
self._validate_args = validate_args |
|
|
if self._validate_args: |
|
|
try: |
|
|
arg_constraints = self.arg_constraints |
|
|
except NotImplementedError: |
|
|
arg_constraints = {} |
|
|
warnings.warn( |
|
|
f"{self.__class__} does not define `arg_constraints`. " |
|
|
+ "Please set `arg_constraints = {}` or initialize the distribution " |
|
|
+ "with `validate_args=False` to turn off validation." |
|
|
) |
|
|
for param, constraint in arg_constraints.items(): |
|
|
if constraints.is_dependent(constraint): |
|
|
continue |
|
|
if param not in self.__dict__ and isinstance( |
|
|
getattr(type(self), param), lazy_property |
|
|
): |
|
|
continue |
|
|
value = getattr(self, param) |
|
|
valid = constraint.check(value) |
|
|
if not torch._is_all_true(valid): |
|
|
raise ValueError( |
|
|
f"Expected parameter {param} " |
|
|
f"({type(value).__name__} of shape {tuple(value.shape)}) " |
|
|
f"of distribution {repr(self)} " |
|
|
f"to satisfy the constraint {repr(constraint)}, " |
|
|
f"but found invalid values:\n{value}" |
|
|
) |
|
|
super().__init__() |
|
|
|
|
|
def expand(self, batch_shape: _size, _instance=None): |
|
|
""" |
|
|
Returns a new distribution instance (or populates an existing instance |
|
|
provided by a derived class) with batch dimensions expanded to |
|
|
`batch_shape`. This method calls :class:`~torch.Tensor.expand` on |
|
|
the distribution's parameters. As such, this does not allocate new |
|
|
memory for the expanded distribution instance. Additionally, |
|
|
this does not repeat any args checking or parameter broadcasting in |
|
|
`__init__.py`, when an instance is first created. |
|
|
|
|
|
Args: |
|
|
batch_shape (torch.Size): the desired expanded size. |
|
|
_instance: new instance provided by subclasses that |
|
|
need to override `.expand`. |
|
|
|
|
|
Returns: |
|
|
New distribution instance with batch dimensions expanded to |
|
|
`batch_size`. |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
@property |
|
|
def batch_shape(self) -> torch.Size: |
|
|
""" |
|
|
Returns the shape over which parameters are batched. |
|
|
""" |
|
|
return self._batch_shape |
|
|
|
|
|
@property |
|
|
def event_shape(self) -> torch.Size: |
|
|
""" |
|
|
Returns the shape of a single sample (without batching). |
|
|
""" |
|
|
return self._event_shape |
|
|
|
|
|
@property |
|
|
def arg_constraints(self) -> dict[str, constraints.Constraint]: |
|
|
""" |
|
|
Returns a dictionary from argument names to |
|
|
:class:`~torch.distributions.constraints.Constraint` objects that |
|
|
should be satisfied by each argument of this distribution. Args that |
|
|
are not tensors need not appear in this dict. |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
@property |
|
|
def support(self) -> Optional[constraints.Constraint]: |
|
|
""" |
|
|
Returns a :class:`~torch.distributions.constraints.Constraint` object |
|
|
representing this distribution's support. |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
@property |
|
|
def mean(self) -> Tensor: |
|
|
""" |
|
|
Returns the mean of the distribution. |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
@property |
|
|
def mode(self) -> Tensor: |
|
|
""" |
|
|
Returns the mode of the distribution. |
|
|
""" |
|
|
raise NotImplementedError(f"{self.__class__} does not implement mode") |
|
|
|
|
|
@property |
|
|
def variance(self) -> Tensor: |
|
|
""" |
|
|
Returns the variance of the distribution. |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
@property |
|
|
def stddev(self) -> Tensor: |
|
|
""" |
|
|
Returns the standard deviation of the distribution. |
|
|
""" |
|
|
return self.variance.sqrt() |
|
|
|
|
|
def sample(self, sample_shape: _size = torch.Size()) -> Tensor: |
|
|
""" |
|
|
Generates a sample_shape shaped sample or sample_shape shaped batch of |
|
|
samples if the distribution parameters are batched. |
|
|
""" |
|
|
with torch.no_grad(): |
|
|
return self.rsample(sample_shape) |
|
|
|
|
|
def rsample(self, sample_shape: _size = torch.Size()) -> Tensor: |
|
|
""" |
|
|
Generates a sample_shape shaped reparameterized sample or sample_shape |
|
|
shaped batch of reparameterized samples if the distribution parameters |
|
|
are batched. |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
@deprecated( |
|
|
"`sample_n(n)` will be deprecated. Use `sample((n,))` instead.", |
|
|
category=FutureWarning, |
|
|
) |
|
|
def sample_n(self, n: int) -> Tensor: |
|
|
""" |
|
|
Generates n samples or n batches of samples if the distribution |
|
|
parameters are batched. |
|
|
""" |
|
|
return self.sample(torch.Size((n,))) |
|
|
|
|
|
def log_prob(self, value: Tensor) -> Tensor: |
|
|
""" |
|
|
Returns the log of the probability density/mass function evaluated at |
|
|
`value`. |
|
|
|
|
|
Args: |
|
|
value (Tensor): |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
def cdf(self, value: Tensor) -> Tensor: |
|
|
""" |
|
|
Returns the cumulative density/mass function evaluated at |
|
|
`value`. |
|
|
|
|
|
Args: |
|
|
value (Tensor): |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
def icdf(self, value: Tensor) -> Tensor: |
|
|
""" |
|
|
Returns the inverse cumulative density/mass function evaluated at |
|
|
`value`. |
|
|
|
|
|
Args: |
|
|
value (Tensor): |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
def enumerate_support(self, expand: bool = True) -> Tensor: |
|
|
""" |
|
|
Returns tensor containing all values supported by a discrete |
|
|
distribution. The result will enumerate over dimension 0, so the shape |
|
|
of the result will be `(cardinality,) + batch_shape + event_shape` |
|
|
(where `event_shape = ()` for univariate distributions). |
|
|
|
|
|
Note that this enumerates over all batched tensors in lock-step |
|
|
`[[0, 0], [1, 1], ...]`. With `expand=False`, enumeration happens |
|
|
along dim 0, but with the remaining batch dimensions being |
|
|
singleton dimensions, `[[0], [1], ..`. |
|
|
|
|
|
To iterate over the full Cartesian product use |
|
|
`itertools.product(m.enumerate_support())`. |
|
|
|
|
|
Args: |
|
|
expand (bool): whether to expand the support over the |
|
|
batch dims to match the distribution's `batch_shape`. |
|
|
|
|
|
Returns: |
|
|
Tensor iterating over dimension 0. |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
def entropy(self) -> Tensor: |
|
|
""" |
|
|
Returns entropy of distribution, batched over batch_shape. |
|
|
|
|
|
Returns: |
|
|
Tensor of shape batch_shape. |
|
|
""" |
|
|
raise NotImplementedError |
|
|
|
|
|
def perplexity(self) -> Tensor: |
|
|
""" |
|
|
Returns perplexity of distribution, batched over batch_shape. |
|
|
|
|
|
Returns: |
|
|
Tensor of shape batch_shape. |
|
|
""" |
|
|
return torch.exp(self.entropy()) |
|
|
|
|
|
def _extended_shape(self, sample_shape: _size = torch.Size()) -> torch.Size: |
|
|
""" |
|
|
Returns the size of the sample returned by the distribution, given |
|
|
a `sample_shape`. Note, that the batch and event shapes of a distribution |
|
|
instance are fixed at the time of construction. If this is empty, the |
|
|
returned shape is upcast to (1,). |
|
|
|
|
|
Args: |
|
|
sample_shape (torch.Size): the size of the sample to be drawn. |
|
|
""" |
|
|
if not isinstance(sample_shape, torch.Size): |
|
|
sample_shape = torch.Size(sample_shape) |
|
|
return torch.Size(sample_shape + self._batch_shape + self._event_shape) |
|
|
|
|
|
def _validate_sample(self, value: Tensor) -> None: |
|
|
""" |
|
|
Argument validation for distribution methods such as `log_prob`, |
|
|
`cdf` and `icdf`. The rightmost dimensions of a value to be |
|
|
scored via these methods must agree with the distribution's batch |
|
|
and event shapes. |
|
|
|
|
|
Args: |
|
|
value (Tensor): the tensor whose log probability is to be |
|
|
computed by the `log_prob` method. |
|
|
Raises |
|
|
ValueError: when the rightmost dimensions of `value` do not match the |
|
|
distribution's batch and event shapes. |
|
|
""" |
|
|
if not isinstance(value, torch.Tensor): |
|
|
raise ValueError("The value argument to log_prob must be a Tensor") |
|
|
|
|
|
event_dim_start = len(value.size()) - len(self._event_shape) |
|
|
if value.size()[event_dim_start:] != self._event_shape: |
|
|
raise ValueError( |
|
|
f"The right-most size of value must match event_shape: {value.size()} vs {self._event_shape}." |
|
|
) |
|
|
|
|
|
actual_shape = value.size() |
|
|
expected_shape = self._batch_shape + self._event_shape |
|
|
for i, j in zip(reversed(actual_shape), reversed(expected_shape)): |
|
|
if i != 1 and j != 1 and i != j: |
|
|
raise ValueError( |
|
|
f"Value is not broadcastable with batch_shape+event_shape: {actual_shape} vs {expected_shape}." |
|
|
) |
|
|
try: |
|
|
support = self.support |
|
|
except NotImplementedError: |
|
|
warnings.warn( |
|
|
f"{self.__class__} does not define `support` to enable " |
|
|
+ "sample validation. Please initialize the distribution with " |
|
|
+ "`validate_args=False` to turn off validation." |
|
|
) |
|
|
return |
|
|
assert support is not None |
|
|
valid = support.check(value) |
|
|
if not torch._is_all_true(valid): |
|
|
raise ValueError( |
|
|
"Expected value argument " |
|
|
f"({type(value).__name__} of shape {tuple(value.shape)}) " |
|
|
f"to be within the support ({repr(support)}) " |
|
|
f"of the distribution {repr(self)}, " |
|
|
f"but found invalid values:\n{value}" |
|
|
) |
|
|
|
|
|
def _get_checked_instance(self, cls, _instance=None): |
|
|
if _instance is None and type(self).__init__ != cls.__init__: |
|
|
raise NotImplementedError( |
|
|
f"Subclass {self.__class__.__name__} of {cls.__name__} that defines a custom __init__ method " |
|
|
"must also define a custom .expand() method." |
|
|
) |
|
|
return self.__new__(type(self)) if _instance is None else _instance |
|
|
|
|
|
def __repr__(self) -> str: |
|
|
param_names = [k for k, _ in self.arg_constraints.items() if k in self.__dict__] |
|
|
args_string = ", ".join( |
|
|
[ |
|
|
f"{p}: {self.__dict__[p] if self.__dict__[p].numel() == 1 else self.__dict__[p].size()}" |
|
|
for p in param_names |
|
|
] |
|
|
) |
|
|
return self.__class__.__name__ + "(" + args_string + ")" |
|
|
|
