Caffe2 - Python API A deep learning, cross platform ML framework
binomial.py
1 from numbers import Number
2 import torch
3 from torch.distributions import constraints
4 from torch.distributions.distribution import Distribution
5 from torch.distributions.utils import broadcast_all, probs_to_logits, lazy_property, logits_to_probs
6
7
9  r"""
10  Creates a Binomial distribution parameterized by :attr:`total_count` and
11  either :attr:`probs` or :attr:`logits` (but not both). :attr:`total_count` must be
13
14  Example::
15
16  >>> m = Binomial(100, torch.tensor([0 , .2, .8, 1]))
17  >>> x = m.sample()
18  tensor([ 0., 22., 71., 100.])
19
20  >>> m = Binomial(torch.tensor([[5.], [10.]]), torch.tensor([0.5, 0.8]))
21  >>> x = m.sample()
22  tensor([[ 4., 5.],
23  [ 7., 6.]])
24
25  Args:
26  total_count (int or Tensor): number of Bernoulli trials
27  probs (Tensor): Event probabilities
28  logits (Tensor): Event log-odds
29  """
30  arg_constraints = {'total_count': constraints.nonnegative_integer,
31  'probs': constraints.unit_interval,
32  'logits': constraints.real}
33  has_enumerate_support = True
34
35  def __init__(self, total_count=1, probs=None, logits=None, validate_args=None):
36  if (probs is None) == (logits is None):
37  raise ValueError("Either `probs` or `logits` must be specified, but not both.")
38  if probs is not None:
39  self.total_count, self.probs, = broadcast_all(total_count, probs)
40  self.total_count = self.total_count.type_as(self.logits)
41  is_scalar = isinstance(self.probs, Number)
42  else:
43  self.total_count, self.logits, = broadcast_all(total_count, logits)
44  self.total_count = self.total_count.type_as(self.logits)
45  is_scalar = isinstance(self.logits, Number)
46
47  self._param = self.probs if probs is not None else self.logits
48  if is_scalar:
49  batch_shape = torch.Size()
50  else:
51  batch_shape = self._param.size()
52  super(Binomial, self).__init__(batch_shape, validate_args=validate_args)
53
54  def expand(self, batch_shape, _instance=None):
55  new = self._get_checked_instance(Binomial, _instance)
56  batch_shape = torch.Size(batch_shape)
57  new.total_count = self.total_count.expand(batch_shape)
58  if 'probs' in self.__dict__:
59  new.probs = self.probs.expand(batch_shape)
60  new._param = new.probs
61  else:
62  new.logits = self.logits.expand(batch_shape)
63  new._param = new.logits
64  super(Binomial, new).__init__(batch_shape, validate_args=False)
65  new._validate_args = self._validate_args
66  return new
67
68  def _new(self, *args, **kwargs):
69  return self._param.new(*args, **kwargs)
70
71  @constraints.dependent_property
72  def support(self):
73  return constraints.integer_interval(0, self.total_count)
74
75  @property
76  def mean(self):
77  return self.total_count * self.probs
78
79  @property
80  def variance(self):
81  return self.total_count * self.probs * (1 - self.probs)
82
83  @lazy_property
84  def logits(self):
85  return probs_to_logits(self.probs, is_binary=True)
86
87  @lazy_property
88  def probs(self):
89  return logits_to_probs(self.logits, is_binary=True)
90
91  @property
92  def param_shape(self):
93  return self._param.size()
94
95  def sample(self, sample_shape=torch.Size()):
97  max_count = max(int(self.total_count.max()), 1)
98  shape = self._extended_shape(sample_shape) + (max_count,)
99  bernoullis = torch.bernoulli(self.probs.unsqueeze(-1).expand(shape))
100  if self.total_count.min() != max_count:
101  arange = torch.arange(max_count, dtype=self._param.dtype, device=self._param.device)
102  mask = arange >= self.total_count.unsqueeze(-1)
103  if torch._C._get_tracing_state():
104  # [JIT WORKAROUND] lack of support for .masked_fill_()
106  else:
108  return bernoullis.sum(dim=-1)
109
110  def log_prob(self, value):
111  if self._validate_args:
112  self._validate_sample(value)
113  log_factorial_n = torch.lgamma(self.total_count + 1)
114  log_factorial_k = torch.lgamma(value + 1)
115  log_factorial_nmk = torch.lgamma(self.total_count - value + 1)
116  # Note that: torch.log1p(-self.probs)) = - torch.log1p(self.logits.exp()))
117  return (log_factorial_n - log_factorial_k - log_factorial_nmk +
118  value * self.logits - self.total_count * torch.log1p(self.logits.exp()))
119
120  def enumerate_support(self, expand=True):
121  total_count = int(self.total_count.max())
122  if not self.total_count.min() == total_count:
123  raise NotImplementedError("Inhomogeneous total count not supported by `enumerate_support`.")
124  values = torch.arange(1 + total_count, dtype=self._param.dtype, device=self._param.device)
125  values = values.view((-1,) + (1,) * len(self._batch_shape))
126  if expand:
127  values = values.expand((-1,) + self._batch_shape)
128  return values
def _get_checked_instance(self, cls, _instance=None)
def _extended_shape(self, sample_shape=torch.Size())