Caffe2 - Python API A deep learning, cross platform ML framework
multinomial.py
1 import torch
2 from torch._six import inf
3 from torch.distributions.distribution import Distribution
4 from torch.distributions import Categorical
5 from numbers import Number
6 from torch.distributions import constraints
8
9
11  r"""
12  Creates a Multinomial distribution parameterized by :attr:`total_count` and
13  either :attr:`probs` or :attr:`logits` (but not both). The innermost dimension of
14  :attr:`probs` indexes over categories. All other dimensions index over batches.
15
16  Note that :attr:`total_count` need not be specified if only :meth:`log_prob` is
17  called (see example below)
18
19  .. note:: :attr:`probs` must be non-negative, finite and have a non-zero sum,
20  and it will be normalized to sum to 1.
21
22  - :meth:`sample` requires a single shared `total_count` for all
23  parameters and samples.
24  - :meth:`log_prob` allows different `total_count` for each parameter and
25  sample.
26
27  Example::
28
29  >>> m = Multinomial(100, torch.tensor([ 1., 1., 1., 1.]))
30  >>> x = m.sample() # equal probability of 0, 1, 2, 3
31  tensor([ 21., 24., 30., 25.])
32
33  >>> Multinomial(probs=torch.tensor([1., 1., 1., 1.])).log_prob(x)
34  tensor([-4.1338])
35
36  Args:
37  total_count (int): number of trials
38  probs (Tensor): event probabilities
39  logits (Tensor): event log probabilities
40  """
41  arg_constraints = {'probs': constraints.simplex,
42  'logits': constraints.real}
43
44  @property
45  def mean(self):
46  return self.probs * self.total_count
47
48  @property
49  def variance(self):
50  return self.total_count * self.probs * (1 - self.probs)
51
52  def __init__(self, total_count=1, probs=None, logits=None, validate_args=None):
53  if not isinstance(total_count, Number):
54  raise NotImplementedError('inhomogeneous total_count is not supported')
55  self.total_count = total_count
56  self._categorical = Categorical(probs=probs, logits=logits)
57  batch_shape = self._categorical.batch_shape
58  event_shape = self._categorical.param_shape[-1:]
59  super(Multinomial, self).__init__(batch_shape, event_shape, validate_args=validate_args)
60
61  def expand(self, batch_shape, _instance=None):
62  new = self._get_checked_instance(Multinomial, _instance)
63  batch_shape = torch.Size(batch_shape)
64  new.total_count = self.total_count
65  new._categorical = self._categorical.expand(batch_shape)
66  super(Multinomial, new).__init__(batch_shape, self.event_shape, validate_args=False)
67  new._validate_args = self._validate_args
68  return new
69
70  def _new(self, *args, **kwargs):
71  return self._categorical._new(*args, **kwargs)
72
73  @constraints.dependent_property
74  def support(self):
75  return constraints.integer_interval(0, self.total_count)
76
77  @property
78  def logits(self):
79  return self._categorical.logits
80
81  @property
82  def probs(self):
83  return self._categorical.probs
84
85  @property
86  def param_shape(self):
87  return self._categorical.param_shape
88
89  def sample(self, sample_shape=torch.Size()):
90  sample_shape = torch.Size(sample_shape)
91  samples = self._categorical.sample(torch.Size((self.total_count,)) + sample_shape)
92  # samples.shape is (total_count, sample_shape, batch_shape), need to change it to
93  # (sample_shape, batch_shape, total_count)
94  shifted_idx = list(range(samples.dim()))
95  shifted_idx.append(shifted_idx.pop(0))
96  samples = samples.permute(*shifted_idx)
97  counts = samples.new(self._extended_shape(sample_shape)).zero_()