We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Implement random_gamma() sampler (with gradients)
Many probability distributions require sampling from the [Gamma distribution](https://en.wikipedia.org/wiki/Gamma_distribution), including: `Gamma`, `Beta`, and `Dirichlet`.

Since Gamma samplers have complex control flow (for rejection sampling) and are seldom a bottleneck in probabilistic algorithms, it should suffice to implement a CPU-only implementation at first. What is more important than a CUDA implementation is a reparameterized sampler so that stochastic gradients can be propagated through the sampler (see [paper](http://proceedings.mlr.press/v54/naesseth17a/naesseth17a.pdf) and [reference implementation](https://github.com/blei-lab/ars-reparameterization) by @naesseth).

## Tasks
- [x] #3841 CPU implementation of basic `random_gamma(requires_grad=False)`
- [ ] #3978 Support reparameterized `random_gamma(requires_grad=True)`

## Map of modifications
- `aten/src/TH/THRandom.c/h` random single numbers
- `aten/src/TH/generic/THTensorRandom.c/h` random tensors
- `aten/src/ATen/Declarations.cwrap` bindings for ATen
- `torch/csrc/generic/methods/TensorRandom.cwrap` bindings for torch.Tensor
- `torch/autograd/variable.py` - Variable
- `torch/distributions.py` - Distributions

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch/distributions.py`
Content:
```
1 r"""
2 The ``distributions`` package contains parameterizable probability distributions
3 and sampling functions.
4 
5 Policy gradient methods can be implemented using the
6 :meth:`~torch.distributions.Distribution.log_prob` method, when the probability
7 density function is differentiable with respect to its parameters. A basic
8 method is the REINFORCE rule:
9 
10 .. math::
11 
12     \Delta\theta  = \alpha r \frac{\partial\log p(a|\pi^\theta(s))}{\partial\theta}
13 
14 where :math:`\theta` are the parameters, :math:`\alpha` is the learning rate,
15 :math:`r` is the reward and :math:`p(a|\pi^\theta(s))` is the probability of
16 taking action :math:`a` in state :math:`s` given policy :math:`\pi^\theta`.
17 
18 In practice we would sample an action from the output of a network, apply this
19 action in an environment, and then use ``log_prob`` to construct an equivalent
20 loss function. Note that we use a negative because optimisers use gradient
21 descent, whilst the rule above assumes gradient ascent. With a categorical
22 policy, the code for implementing REINFORCE would be as follows::
23 
24     probs = policy_network(state)
25     # NOTE: this is equivalent to what used to be called multinomial
26     m = Categorical(probs)
27     action = m.sample()
28     next_state, reward = env.step(action)
29     loss = -m.log_prob(action) * reward
30     loss.backward()
31 """
32 import math
33 from numbers import Number
34 import torch
35 from torch.autograd import Variable
36 
37 
38 __all__ = ['Distribution', 'Bernoulli', 'Categorical', 'Normal', 'Gamma']
39 
40 
41 def _expand_n(v, n):
42     r"""
43     Cleanly expand float or Tensor or Variable parameters.
44     """
45     if isinstance(v, Number):
46         return torch.Tensor([v]).expand(n, 1)
47     else:
48         return v.expand(n, *v.size())
49 
50 
51 class Distribution(object):
52     r"""
53     Distribution is the abstract base class for probability distributions.
54     """
55 
56     def sample(self):
57         """
58         Generates a single sample or single batch of samples if the distribution
59         parameters are batched.
60         """
61         raise NotImplementedError
62 
63     def sample_n(self, n):
64         """
65         Generates n samples or n batches of samples if the distribution parameters
66         are batched.
67         """
68         raise NotImplementedError
69 
70     def log_prob(self, value):
71         """
72         Returns the log of the probability density/mass function evaluated at
73         `value`.
74 
75         Args:
76             value (Tensor or Variable):
77         """
78         raise NotImplementedError
79 
80 
81 class Bernoulli(Distribution):
82     r"""
83     Creates a Bernoulli distribution parameterized by `probs`.
84 
85     Samples are binary (0 or 1). They take the value `1` with probability `p`
86     and `0` with probability `1 - p`.
87 
88     Example::
89 
90         >>> m = Bernoulli(torch.Tensor([0.3]))
91         >>> m.sample()  # 30% chance 1; 70% chance 0
92          0.0
93         [torch.FloatTensor of size 1]
94 
95     Args:
96         probs (Tensor or Variable): the probabilty of sampling `1`
97     """
98 
99     def __init__(self, probs):
100         self.probs = probs
101 
102     def sample(self):
103         return torch.bernoulli(self.probs)
104 
105     def sample_n(self, n):
106         return torch.bernoulli(self.probs.expand(n, *self.probs.size()))
107 
108     def log_prob(self, value):
109         # compute the log probabilities for 0 and 1
110         log_pmf = (torch.stack([1 - self.probs, self.probs])).log()
111 
112         # evaluate using the values
113         return log_pmf.gather(0, value.unsqueeze(0).long()).squeeze(0)
114 
115 
116 class Categorical(Distribution):
117     r"""
118     Creates a categorical distribution parameterized by `probs`.
119 
120     .. note::
121         It is equivalent to the distribution that ``multinomial()`` samples from.
122 
123     Samples are integers from `0 ... K-1` where `K` is probs.size(-1).
124 
125     If `probs` is 1D with length-`K`, each element is the relative probability
126     of sampling the class at that index.
127 
128     If `probs` is 2D, it is treated as a batch of probability vectors.
129 
130     See also: :func:`torch.multinomial`
131 
132     Example::
133 
134         >>> m = Categorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))
135         >>> m.sample()  # equal probability of 0, 1, 2, 3
136          3
137         [torch.LongTensor of size 1]
138 
139     Args:
140         probs (Tensor or Variable): event probabilities
141     """
142 
143     def __init__(self, probs):
144         if probs.dim() != 1 and probs.dim() != 2:
145             # TODO: treat higher dimensions as part of the batch
146             raise ValueError("probs must be 1D or 2D")
147         self.probs = probs
148 
149     def sample(self):
150         return torch.multinomial(self.probs, 1, True).squeeze(-1)
151 
152     def sample_n(self, n):
153         if n == 1:
154             return self.sample().expand(1, 1)
155         else:
156             return torch.multinomial(self.probs, n, True).t()
157 
158     def log_prob(self, value):
159         p = self.probs / self.probs.sum(-1, keepdim=True)
160         if value.dim() == 1 and self.probs.dim() == 1:
161             # special handling until we have 0-dim tensor support
162             return p.gather(-1, value).log()
163 
164         return p.gather(-1, value.unsqueeze(-1)).squeeze(-1).log()
165 
166 
167 class Normal(Distribution):
168     r"""
169     Creates a normal (also called Gaussian) distribution parameterized by
170     `mean` and `std`.
171 
172     Example::
173 
174         >>> m = Normal(torch.Tensor([0.0]), torch.Tensor([1.0]))
175         >>> m.sample()  # normally distributed with mean=0 and stddev=1
176          0.1046
177         [torch.FloatTensor of size 1]
178 
179     Args:
180         mean (float or Tensor or Variable): mean of the distribution
181         std (float or Tensor or Variable): standard deviation of the distribution
182     """
183 
184     def __init__(self, mean, std):
185         self.mean = mean
186         self.std = std
187 
188     def sample(self):
189         return torch.normal(self.mean, self.std)
190 
191     def sample_n(self, n):
192         return torch.normal(_expand_n(self.mean, n), _expand_n(self.std, n))
193 
194     def log_prob(self, value):
195         # compute the variance
196         var = (self.std ** 2)
197         log_std = math.log(self.std) if isinstance(self.std, Number) else self.std.log()
198         return -((value - self.mean) ** 2) / (2 * var) - log_std - math.log(math.sqrt(2 * math.pi))
199 
200 
201 def _standard_gamma(alpha):
202     if not isinstance(alpha, Variable):
203         return torch._C._standard_gamma(alpha)
204     return Variable(torch._C._standard_gamma(alpha.data))
205 
206 
207 class Gamma(Distribution):
208     r"""
209     Creates a Gamma distribution parameterized by shape `alpha` and rate `beta`.
210 
211     Example::
212 
213         >>> m = Gamma(torch.Tensor([1.0]), torch.Tensor([1.0]))
214         >>> m.sample()  # Gamma distributed with shape alpha=1 and rate beta=1
215          0.1046
216         [torch.FloatTensor of size 1]
217 
218     Args:
219         alpha (float or Tensor or Variable): shape parameter of the distribution
220         beta (float or Tensor or Variable): rate = 1 / scale of the distribution
221     """
222 
223     def __init__(self, alpha, beta):
224         # TODO handle (Variable, Number) cases
225         alpha_num = isinstance(alpha, Number)
226         beta_num = isinstance(beta, Number)
227         if alpha_num and not beta_num:
228             alpha = beta.new(beta.size()).fill_(alpha)
229         elif not alpha_num and beta_num:
230             beta = alpha.new(alpha.size()).fill_(beta)
231         elif alpha_num and beta_num:
232             alpha, beta = torch.Tensor([alpha]), torch.Tensor([beta])
233         elif alpha.size() != beta.size():
234             raise ValueError('Expected alpha.size() == beta.size(), actual {} vs {}'.format(
235                 alpha.size(), beta.size()))
236         self.alpha = alpha
237         self.beta = beta
238 
239     def sample(self):
240         return _standard_gamma(self.alpha) / self.beta
241 
242     def sample_n(self, n):
243         return _standard_gamma(_expand_n(self.alpha, n)) / self.beta
244 
245     def log_prob(self, value):
246         return (self.alpha * torch.log(self.beta) +
247                 (self.alpha - 1) * torch.log(value) -
248                 self.beta * value - torch.lgamma(self.alpha))
249
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torch/distributions.py b/torch/distributions.py
--- a/torch/distributions.py
+++ b/torch/distributions.py
@@ -32,8 +32,8 @@
 import math
 from numbers import Number
 import torch
-from torch.autograd import Variable
-
+from torch.autograd import Function, Variable
+from torch.autograd.function import once_differentiable
 
 __all__ = ['Distribution', 'Bernoulli', 'Categorical', 'Normal', 'Gamma']
 
@@ -198,10 +198,25 @@
         return -((value - self.mean) ** 2) / (2 * var) - log_std - math.log(math.sqrt(2 * math.pi))
 
 
+class _StandardGamma(Function):
+    @staticmethod
+    def forward(ctx, alpha):
+        x = torch._C._standard_gamma(alpha)
+        ctx.save_for_backward(x, alpha)
+        return x
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        x, alpha = ctx.saved_tensors
+        grad = torch._C._standard_gamma_grad(x, alpha)
+        return grad_output * grad
+
+
 def _standard_gamma(alpha):
     if not isinstance(alpha, Variable):
         return torch._C._standard_gamma(alpha)
-    return Variable(torch._C._standard_gamma(alpha.data))
+    return _StandardGamma.apply(alpha)
 
 
 class Gamma(Distribution):
@@ -219,6 +234,7 @@
         alpha (float or Tensor or Variable): shape parameter of the distribution
         beta (float or Tensor or Variable): rate = 1 / scale of the distribution
     """
+    has_rsample = True
 
     def __init__(self, alpha, beta):
         # TODO handle (Variable, Number) cases

{"golden_diff": "diff --git a/torch/distributions.py b/torch/distributions.py\n--- a/torch/distributions.py\n+++ b/torch/distributions.py\n@@ -32,8 +32,8 @@\n import math\n from numbers import Number\n import torch\n-from torch.autograd import Variable\n-\n+from torch.autograd import Function, Variable\n+from torch.autograd.function import once_differentiable\n \n __all__ = ['Distribution', 'Bernoulli', 'Categorical', 'Normal', 'Gamma']\n \n@@ -198,10 +198,25 @@\n         return -((value - self.mean) ** 2) / (2 * var) - log_std - math.log(math.sqrt(2 * math.pi))\n \n \n+class _StandardGamma(Function):\n+    @staticmethod\n+    def forward(ctx, alpha):\n+        x = torch._C._standard_gamma(alpha)\n+        ctx.save_for_backward(x, alpha)\n+        return x\n+\n+    @staticmethod\n+    @once_differentiable\n+    def backward(ctx, grad_output):\n+        x, alpha = ctx.saved_tensors\n+        grad = torch._C._standard_gamma_grad(x, alpha)\n+        return grad_output * grad\n+\n+\n def _standard_gamma(alpha):\n     if not isinstance(alpha, Variable):\n         return torch._C._standard_gamma(alpha)\n-    return Variable(torch._C._standard_gamma(alpha.data))\n+    return _StandardGamma.apply(alpha)\n \n \n class Gamma(Distribution):\n@@ -219,6 +234,7 @@\n         alpha (float or Tensor or Variable): shape parameter of the distribution\n         beta (float or Tensor or Variable): rate = 1 / scale of the distribution\n     \"\"\"\n+    has_rsample = True\n \n     def __init__(self, alpha, beta):\n         # TODO handle (Variable, Number) cases\n", "issue": "Implement random_gamma() sampler (with gradients)\nMany probability distributions require sampling from the [Gamma distribution](https://en.wikipedia.org/wiki/Gamma_distribution), including: `Gamma`, `Beta`, and `Dirichlet`.\r\n\r\nSince Gamma samplers have complex control flow (for rejection sampling) and are seldom a bottleneck in probabilistic algorithms, it should suffice to implement a CPU-only implementation at first. What is more important than a CUDA implementation is a reparameterized sampler so that stochastic gradients can be propagated through the sampler (see [paper](http://proceedings.mlr.press/v54/naesseth17a/naesseth17a.pdf) and [reference implementation](https://github.com/blei-lab/ars-reparameterization) by @naesseth).\r\n\r\n## Tasks\r\n- [x] #3841 CPU implementation of basic `random_gamma(requires_grad=False)`\r\n- [ ] #3978 Support reparameterized `random_gamma(requires_grad=True)`\r\n\r\n## Map of modifications\r\n- `aten/src/TH/THRandom.c/h` random single numbers\r\n- `aten/src/TH/generic/THTensorRandom.c/h` random tensors\r\n- `aten/src/ATen/Declarations.cwrap` bindings for ATen\r\n- `torch/csrc/generic/methods/TensorRandom.cwrap` bindings for torch.Tensor\r\n- `torch/autograd/variable.py` - Variable\r\n- `torch/distributions.py` - Distributions\n", "before_files": [{"content": "r\"\"\"\nThe ``distributions`` package contains parameterizable probability distributions\nand sampling functions.\n\nPolicy gradient methods can be implemented using the\n:meth:`~torch.distributions.Distribution.log_prob` method, when the probability\ndensity function is differentiable with respect to its parameters. A basic\nmethod is the REINFORCE rule:\n\n.. math::\n\n    \\Delta\\theta  = \\alpha r \\frac{\\partial\\log p(a|\\pi^\\theta(s))}{\\partial\\theta}\n\nwhere :math:`\\theta` are the parameters, :math:`\\alpha` is the learning rate,\n:math:`r` is the reward and :math:`p(a|\\pi^\\theta(s))` is the probability of\ntaking action :math:`a` in state :math:`s` given policy :math:`\\pi^\\theta`.\n\nIn practice we would sample an action from the output of a network, apply this\naction in an environment, and then use ``log_prob`` to construct an equivalent\nloss function. Note that we use a negative because optimisers use gradient\ndescent, whilst the rule above assumes gradient ascent. With a categorical\npolicy, the code for implementing REINFORCE would be as follows::\n\n    probs = policy_network(state)\n    # NOTE: this is equivalent to what used to be called multinomial\n    m = Categorical(probs)\n    action = m.sample()\n    next_state, reward = env.step(action)\n    loss = -m.log_prob(action) * reward\n    loss.backward()\n\"\"\"\nimport math\nfrom numbers import Number\nimport torch\nfrom torch.autograd import Variable\n\n\n__all__ = ['Distribution', 'Bernoulli', 'Categorical', 'Normal', 'Gamma']\n\n\ndef _expand_n(v, n):\n    r\"\"\"\n    Cleanly expand float or Tensor or Variable parameters.\n    \"\"\"\n    if isinstance(v, Number):\n        return torch.Tensor([v]).expand(n, 1)\n    else:\n        return v.expand(n, *v.size())\n\n\nclass Distribution(object):\n    r\"\"\"\n    Distribution is the abstract base class for probability distributions.\n    \"\"\"\n\n    def sample(self):\n        \"\"\"\n        Generates a single sample or single batch of samples if the distribution\n        parameters are batched.\n        \"\"\"\n        raise NotImplementedError\n\n    def sample_n(self, n):\n        \"\"\"\n        Generates n samples or n batches of samples if the distribution parameters\n        are batched.\n        \"\"\"\n        raise NotImplementedError\n\n    def log_prob(self, value):\n        \"\"\"\n        Returns the log of the probability density/mass function evaluated at\n        `value`.\n\n        Args:\n            value (Tensor or Variable):\n        \"\"\"\n        raise NotImplementedError\n\n\nclass Bernoulli(Distribution):\n    r\"\"\"\n    Creates a Bernoulli distribution parameterized by `probs`.\n\n    Samples are binary (0 or 1). They take the value `1` with probability `p`\n    and `0` with probability `1 - p`.\n\n    Example::\n\n        >>> m = Bernoulli(torch.Tensor([0.3]))\n        >>> m.sample()  # 30% chance 1; 70% chance 0\n         0.0\n        [torch.FloatTensor of size 1]\n\n    Args:\n        probs (Tensor or Variable): the probabilty of sampling `1`\n    \"\"\"\n\n    def __init__(self, probs):\n        self.probs = probs\n\n    def sample(self):\n        return torch.bernoulli(self.probs)\n\n    def sample_n(self, n):\n        return torch.bernoulli(self.probs.expand(n, *self.probs.size()))\n\n    def log_prob(self, value):\n        # compute the log probabilities for 0 and 1\n        log_pmf = (torch.stack([1 - self.probs, self.probs])).log()\n\n        # evaluate using the values\n        return log_pmf.gather(0, value.unsqueeze(0).long()).squeeze(0)\n\n\nclass Categorical(Distribution):\n    r\"\"\"\n    Creates a categorical distribution parameterized by `probs`.\n\n    .. note::\n        It is equivalent to the distribution that ``multinomial()`` samples from.\n\n    Samples are integers from `0 ... K-1` where `K` is probs.size(-1).\n\n    If `probs` is 1D with length-`K`, each element is the relative probability\n    of sampling the class at that index.\n\n    If `probs` is 2D, it is treated as a batch of probability vectors.\n\n    See also: :func:`torch.multinomial`\n\n    Example::\n\n        >>> m = Categorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))\n        >>> m.sample()  # equal probability of 0, 1, 2, 3\n         3\n        [torch.LongTensor of size 1]\n\n    Args:\n        probs (Tensor or Variable): event probabilities\n    \"\"\"\n\n    def __init__(self, probs):\n        if probs.dim() != 1 and probs.dim() != 2:\n            # TODO: treat higher dimensions as part of the batch\n            raise ValueError(\"probs must be 1D or 2D\")\n        self.probs = probs\n\n    def sample(self):\n        return torch.multinomial(self.probs, 1, True).squeeze(-1)\n\n    def sample_n(self, n):\n        if n == 1:\n            return self.sample().expand(1, 1)\n        else:\n            return torch.multinomial(self.probs, n, True).t()\n\n    def log_prob(self, value):\n        p = self.probs / self.probs.sum(-1, keepdim=True)\n        if value.dim() == 1 and self.probs.dim() == 1:\n            # special handling until we have 0-dim tensor support\n            return p.gather(-1, value).log()\n\n        return p.gather(-1, value.unsqueeze(-1)).squeeze(-1).log()\n\n\nclass Normal(Distribution):\n    r\"\"\"\n    Creates a normal (also called Gaussian) distribution parameterized by\n    `mean` and `std`.\n\n    Example::\n\n        >>> m = Normal(torch.Tensor([0.0]), torch.Tensor([1.0]))\n        >>> m.sample()  # normally distributed with mean=0 and stddev=1\n         0.1046\n        [torch.FloatTensor of size 1]\n\n    Args:\n        mean (float or Tensor or Variable): mean of the distribution\n        std (float or Tensor or Variable): standard deviation of the distribution\n    \"\"\"\n\n    def __init__(self, mean, std):\n        self.mean = mean\n        self.std = std\n\n    def sample(self):\n        return torch.normal(self.mean, self.std)\n\n    def sample_n(self, n):\n        return torch.normal(_expand_n(self.mean, n), _expand_n(self.std, n))\n\n    def log_prob(self, value):\n        # compute the variance\n        var = (self.std ** 2)\n        log_std = math.log(self.std) if isinstance(self.std, Number) else self.std.log()\n        return -((value - self.mean) ** 2) / (2 * var) - log_std - math.log(math.sqrt(2 * math.pi))\n\n\ndef _standard_gamma(alpha):\n    if not isinstance(alpha, Variable):\n        return torch._C._standard_gamma(alpha)\n    return Variable(torch._C._standard_gamma(alpha.data))\n\n\nclass Gamma(Distribution):\n    r\"\"\"\n    Creates a Gamma distribution parameterized by shape `alpha` and rate `beta`.\n\n    Example::\n\n        >>> m = Gamma(torch.Tensor([1.0]), torch.Tensor([1.0]))\n        >>> m.sample()  # Gamma distributed with shape alpha=1 and rate beta=1\n         0.1046\n        [torch.FloatTensor of size 1]\n\n    Args:\n        alpha (float or Tensor or Variable): shape parameter of the distribution\n        beta (float or Tensor or Variable): rate = 1 / scale of the distribution\n    \"\"\"\n\n    def __init__(self, alpha, beta):\n        # TODO handle (Variable, Number) cases\n        alpha_num = isinstance(alpha, Number)\n        beta_num = isinstance(beta, Number)\n        if alpha_num and not beta_num:\n            alpha = beta.new(beta.size()).fill_(alpha)\n        elif not alpha_num and beta_num:\n            beta = alpha.new(alpha.size()).fill_(beta)\n        elif alpha_num and beta_num:\n            alpha, beta = torch.Tensor([alpha]), torch.Tensor([beta])\n        elif alpha.size() != beta.size():\n            raise ValueError('Expected alpha.size() == beta.size(), actual {} vs {}'.format(\n                alpha.size(), beta.size()))\n        self.alpha = alpha\n        self.beta = beta\n\n    def sample(self):\n        return _standard_gamma(self.alpha) / self.beta\n\n    def sample_n(self, n):\n        return _standard_gamma(_expand_n(self.alpha, n)) / self.beta\n\n    def log_prob(self, value):\n        return (self.alpha * torch.log(self.beta) +\n                (self.alpha - 1) * torch.log(value) -\n                self.beta * value - torch.lgamma(self.alpha))\n", "path": "torch/distributions.py"}], "after_files": [{"content": "r\"\"\"\nThe ``distributions`` package contains parameterizable probability distributions\nand sampling functions.\n\nPolicy gradient methods can be implemented using the\n:meth:`~torch.distributions.Distribution.log_prob` method, when the probability\ndensity function is differentiable with respect to its parameters. A basic\nmethod is the REINFORCE rule:\n\n.. math::\n\n    \\Delta\\theta  = \\alpha r \\frac{\\partial\\log p(a|\\pi^\\theta(s))}{\\partial\\theta}\n\nwhere :math:`\\theta` are the parameters, :math:`\\alpha` is the learning rate,\n:math:`r` is the reward and :math:`p(a|\\pi^\\theta(s))` is the probability of\ntaking action :math:`a` in state :math:`s` given policy :math:`\\pi^\\theta`.\n\nIn practice we would sample an action from the output of a network, apply this\naction in an environment, and then use ``log_prob`` to construct an equivalent\nloss function. Note that we use a negative because optimisers use gradient\ndescent, whilst the rule above assumes gradient ascent. With a categorical\npolicy, the code for implementing REINFORCE would be as follows::\n\n    probs = policy_network(state)\n    # NOTE: this is equivalent to what used to be called multinomial\n    m = Categorical(probs)\n    action = m.sample()\n    next_state, reward = env.step(action)\n    loss = -m.log_prob(action) * reward\n    loss.backward()\n\"\"\"\nimport math\nfrom numbers import Number\nimport torch\nfrom torch.autograd import Function, Variable\nfrom torch.autograd.function import once_differentiable\n\n__all__ = ['Distribution', 'Bernoulli', 'Categorical', 'Normal', 'Gamma']\n\n\ndef _expand_n(v, n):\n    r\"\"\"\n    Cleanly expand float or Tensor or Variable parameters.\n    \"\"\"\n    if isinstance(v, Number):\n        return torch.Tensor([v]).expand(n, 1)\n    else:\n        return v.expand(n, *v.size())\n\n\nclass Distribution(object):\n    r\"\"\"\n    Distribution is the abstract base class for probability distributions.\n    \"\"\"\n\n    def sample(self):\n        \"\"\"\n        Generates a single sample or single batch of samples if the distribution\n        parameters are batched.\n        \"\"\"\n        raise NotImplementedError\n\n    def sample_n(self, n):\n        \"\"\"\n        Generates n samples or n batches of samples if the distribution parameters\n        are batched.\n        \"\"\"\n        raise NotImplementedError\n\n    def log_prob(self, value):\n        \"\"\"\n        Returns the log of the probability density/mass function evaluated at\n        `value`.\n\n        Args:\n            value (Tensor or Variable):\n        \"\"\"\n        raise NotImplementedError\n\n\nclass Bernoulli(Distribution):\n    r\"\"\"\n    Creates a Bernoulli distribution parameterized by `probs`.\n\n    Samples are binary (0 or 1). They take the value `1` with probability `p`\n    and `0` with probability `1 - p`.\n\n    Example::\n\n        >>> m = Bernoulli(torch.Tensor([0.3]))\n        >>> m.sample()  # 30% chance 1; 70% chance 0\n         0.0\n        [torch.FloatTensor of size 1]\n\n    Args:\n        probs (Tensor or Variable): the probabilty of sampling `1`\n    \"\"\"\n\n    def __init__(self, probs):\n        self.probs = probs\n\n    def sample(self):\n        return torch.bernoulli(self.probs)\n\n    def sample_n(self, n):\n        return torch.bernoulli(self.probs.expand(n, *self.probs.size()))\n\n    def log_prob(self, value):\n        # compute the log probabilities for 0 and 1\n        log_pmf = (torch.stack([1 - self.probs, self.probs])).log()\n\n        # evaluate using the values\n        return log_pmf.gather(0, value.unsqueeze(0).long()).squeeze(0)\n\n\nclass Categorical(Distribution):\n    r\"\"\"\n    Creates a categorical distribution parameterized by `probs`.\n\n    .. note::\n        It is equivalent to the distribution that ``multinomial()`` samples from.\n\n    Samples are integers from `0 ... K-1` where `K` is probs.size(-1).\n\n    If `probs` is 1D with length-`K`, each element is the relative probability\n    of sampling the class at that index.\n\n    If `probs` is 2D, it is treated as a batch of probability vectors.\n\n    See also: :func:`torch.multinomial`\n\n    Example::\n\n        >>> m = Categorical(torch.Tensor([ 0.25, 0.25, 0.25, 0.25 ]))\n        >>> m.sample()  # equal probability of 0, 1, 2, 3\n         3\n        [torch.LongTensor of size 1]\n\n    Args:\n        probs (Tensor or Variable): event probabilities\n    \"\"\"\n\n    def __init__(self, probs):\n        if probs.dim() != 1 and probs.dim() != 2:\n            # TODO: treat higher dimensions as part of the batch\n            raise ValueError(\"probs must be 1D or 2D\")\n        self.probs = probs\n\n    def sample(self):\n        return torch.multinomial(self.probs, 1, True).squeeze(-1)\n\n    def sample_n(self, n):\n        if n == 1:\n            return self.sample().expand(1, 1)\n        else:\n            return torch.multinomial(self.probs, n, True).t()\n\n    def log_prob(self, value):\n        p = self.probs / self.probs.sum(-1, keepdim=True)\n        if value.dim() == 1 and self.probs.dim() == 1:\n            # special handling until we have 0-dim tensor support\n            return p.gather(-1, value).log()\n\n        return p.gather(-1, value.unsqueeze(-1)).squeeze(-1).log()\n\n\nclass Normal(Distribution):\n    r\"\"\"\n    Creates a normal (also called Gaussian) distribution parameterized by\n    `mean` and `std`.\n\n    Example::\n\n        >>> m = Normal(torch.Tensor([0.0]), torch.Tensor([1.0]))\n        >>> m.sample()  # normally distributed with mean=0 and stddev=1\n         0.1046\n        [torch.FloatTensor of size 1]\n\n    Args:\n        mean (float or Tensor or Variable): mean of the distribution\n        std (float or Tensor or Variable): standard deviation of the distribution\n    \"\"\"\n\n    def __init__(self, mean, std):\n        self.mean = mean\n        self.std = std\n\n    def sample(self):\n        return torch.normal(self.mean, self.std)\n\n    def sample_n(self, n):\n        return torch.normal(_expand_n(self.mean, n), _expand_n(self.std, n))\n\n    def log_prob(self, value):\n        # compute the variance\n        var = (self.std ** 2)\n        log_std = math.log(self.std) if isinstance(self.std, Number) else self.std.log()\n        return -((value - self.mean) ** 2) / (2 * var) - log_std - math.log(math.sqrt(2 * math.pi))\n\n\nclass _StandardGamma(Function):\n    @staticmethod\n    def forward(ctx, alpha):\n        x = torch._C._standard_gamma(alpha)\n        ctx.save_for_backward(x, alpha)\n        return x\n\n    @staticmethod\n    @once_differentiable\n    def backward(ctx, grad_output):\n        x, alpha = ctx.saved_tensors\n        grad = torch._C._standard_gamma_grad(x, alpha)\n        return grad_output * grad\n\n\ndef _standard_gamma(alpha):\n    if not isinstance(alpha, Variable):\n        return torch._C._standard_gamma(alpha)\n    return _StandardGamma.apply(alpha)\n\n\nclass Gamma(Distribution):\n    r\"\"\"\n    Creates a Gamma distribution parameterized by shape `alpha` and rate `beta`.\n\n    Example::\n\n        >>> m = Gamma(torch.Tensor([1.0]), torch.Tensor([1.0]))\n        >>> m.sample()  # Gamma distributed with shape alpha=1 and rate beta=1\n         0.1046\n        [torch.FloatTensor of size 1]\n\n    Args:\n        alpha (float or Tensor or Variable): shape parameter of the distribution\n        beta (float or Tensor or Variable): rate = 1 / scale of the distribution\n    \"\"\"\n    has_rsample = True\n\n    def __init__(self, alpha, beta):\n        # TODO handle (Variable, Number) cases\n        alpha_num = isinstance(alpha, Number)\n        beta_num = isinstance(beta, Number)\n        if alpha_num and not beta_num:\n            alpha = beta.new(beta.size()).fill_(alpha)\n        elif not alpha_num and beta_num:\n            beta = alpha.new(alpha.size()).fill_(beta)\n        elif alpha_num and beta_num:\n            alpha, beta = torch.Tensor([alpha]), torch.Tensor([beta])\n        elif alpha.size() != beta.size():\n            raise ValueError('Expected alpha.size() == beta.size(), actual {} vs {}'.format(\n                alpha.size(), beta.size()))\n        self.alpha = alpha\n        self.beta = beta\n\n    def sample(self):\n        return _standard_gamma(self.alpha) / self.beta\n\n    def sample_n(self, n):\n        return _standard_gamma(_expand_n(self.alpha, n)) / self.beta\n\n    def log_prob(self, value):\n        return (self.alpha * torch.log(self.beta) +\n                (self.alpha - 1) * torch.log(value) -\n                self.beta * value - torch.lgamma(self.alpha))\n", "path": "torch/distributions.py"}]}