mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
4618371da5
This is a new version of #15648 based on the latest master branch. Unlike the previous PR where I fixed a lot of the doctests in addition to integrating xdoctest, I'm going to reduce the scope here. I'm simply going to integrate xdoctest, and then I'm going to mark all of the failing tests as "SKIP". This will let xdoctest run on the dashboards, provide some value, and still let the dashboards pass. I'll leave fixing the doctests themselves to another PR. In my initial commit, I do the bare minimum to get something running with failing dashboards. The few tests that I marked as skip are causing segfaults. Running xdoctest results in 293 failed, 201 passed tests. The next commits will be to disable those tests. (unfortunately I don't have a tool that will insert the `#xdoctest: +SKIP` directive over every failing test, so I'm going to do this mostly manually.) Fixes https://github.com/pytorch/pytorch/issues/71105 @ezyang Pull Request resolved: https://github.com/pytorch/pytorch/pull/82797 Approved by: https://github.com/ezyang
122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
from numbers import Number
|
|
|
|
import torch
|
|
from torch._six import nan
|
|
from torch.distributions import constraints
|
|
from torch.distributions.exp_family import ExponentialFamily
|
|
from torch.distributions.utils import broadcast_all, probs_to_logits, logits_to_probs, lazy_property
|
|
from torch.nn.functional import binary_cross_entropy_with_logits
|
|
|
|
__all__ = ['Bernoulli']
|
|
|
|
class Bernoulli(ExponentialFamily):
|
|
r"""
|
|
Creates a Bernoulli distribution parameterized by :attr:`probs`
|
|
or :attr:`logits` (but not both).
|
|
|
|
Samples are binary (0 or 1). They take the value `1` with probability `p`
|
|
and `0` with probability `1 - p`.
|
|
|
|
Example::
|
|
|
|
>>> # xdoctest: +IGNORE_WANT("non-deterinistic")
|
|
>>> m = Bernoulli(torch.tensor([0.3]))
|
|
>>> m.sample() # 30% chance 1; 70% chance 0
|
|
tensor([ 0.])
|
|
|
|
Args:
|
|
probs (Number, Tensor): the probability of sampling `1`
|
|
logits (Number, Tensor): the log-odds of sampling `1`
|
|
"""
|
|
arg_constraints = {'probs': constraints.unit_interval,
|
|
'logits': constraints.real}
|
|
support = constraints.boolean
|
|
has_enumerate_support = True
|
|
_mean_carrier_measure = 0
|
|
|
|
def __init__(self, probs=None, logits=None, validate_args=None):
|
|
if (probs is None) == (logits is None):
|
|
raise ValueError("Either `probs` or `logits` must be specified, but not both.")
|
|
if probs is not None:
|
|
is_scalar = isinstance(probs, Number)
|
|
self.probs, = broadcast_all(probs)
|
|
else:
|
|
is_scalar = isinstance(logits, Number)
|
|
self.logits, = broadcast_all(logits)
|
|
self._param = self.probs if probs is not None else self.logits
|
|
if is_scalar:
|
|
batch_shape = torch.Size()
|
|
else:
|
|
batch_shape = self._param.size()
|
|
super(Bernoulli, self).__init__(batch_shape, validate_args=validate_args)
|
|
|
|
def expand(self, batch_shape, _instance=None):
|
|
new = self._get_checked_instance(Bernoulli, _instance)
|
|
batch_shape = torch.Size(batch_shape)
|
|
if 'probs' in self.__dict__:
|
|
new.probs = self.probs.expand(batch_shape)
|
|
new._param = new.probs
|
|
if 'logits' in self.__dict__:
|
|
new.logits = self.logits.expand(batch_shape)
|
|
new._param = new.logits
|
|
super(Bernoulli, new).__init__(batch_shape, validate_args=False)
|
|
new._validate_args = self._validate_args
|
|
return new
|
|
|
|
def _new(self, *args, **kwargs):
|
|
return self._param.new(*args, **kwargs)
|
|
|
|
@property
|
|
def mean(self):
|
|
return self.probs
|
|
|
|
@property
|
|
def mode(self):
|
|
mode = (self.probs >= 0.5).to(self.probs)
|
|
mode[self.probs == 0.5] = nan
|
|
return mode
|
|
|
|
@property
|
|
def variance(self):
|
|
return self.probs * (1 - self.probs)
|
|
|
|
@lazy_property
|
|
def logits(self):
|
|
return probs_to_logits(self.probs, is_binary=True)
|
|
|
|
@lazy_property
|
|
def probs(self):
|
|
return logits_to_probs(self.logits, is_binary=True)
|
|
|
|
@property
|
|
def param_shape(self):
|
|
return self._param.size()
|
|
|
|
def sample(self, sample_shape=torch.Size()):
|
|
shape = self._extended_shape(sample_shape)
|
|
with torch.no_grad():
|
|
return torch.bernoulli(self.probs.expand(shape))
|
|
|
|
def log_prob(self, value):
|
|
if self._validate_args:
|
|
self._validate_sample(value)
|
|
logits, value = broadcast_all(self.logits, value)
|
|
return -binary_cross_entropy_with_logits(logits, value, reduction='none')
|
|
|
|
def entropy(self):
|
|
return binary_cross_entropy_with_logits(self.logits, self.probs, reduction='none')
|
|
|
|
def enumerate_support(self, expand=True):
|
|
values = torch.arange(2, dtype=self._param.dtype, device=self._param.device)
|
|
values = values.view((-1,) + (1,) * len(self._batch_shape))
|
|
if expand:
|
|
values = values.expand((-1,) + self._batch_shape)
|
|
return values
|
|
|
|
@property
|
|
def _natural_params(self):
|
|
return (torch.log(self.probs / (1 - self.probs)), )
|
|
|
|
def _log_normalizer(self, x):
|
|
return torch.log(1 + torch.exp(x))
|