/
gumbel.py
108 lines (83 loc) · 2.83 KB
/
gumbel.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
import numpy
import chainer
from chainer.backends import cuda
from chainer import distribution
from chainer.functions.math import exponential
from chainer.functions.math import lgamma
from chainer.utils import cache
EULER = 0.57721566490153286060651209008240243104215933593992
class Gumbel(distribution.Distribution):
"""Gumbel Distribution.
The probability density function of the distribution is expressed as
.. math::
f(x) = \\frac{1}{\\eta} \
\\exp\\left\\{ - \\frac{x - \\mu}{\\eta} \\right\\} \
\\exp\\left[-\\exp\\left\\{-\\frac{x - \\mu}{\\eta} \
\\right\\}\\right],
Args:
loc(:class:`~chainer.Variable` or :ref:`ndarray`): Parameter of
distribution :math:`\\mu`.
scale(:class:`~chainer.Variable` or :ref:`ndarray`): Parameter of
distribution :math:`\\eta`.
"""
def __init__(self, loc, scale):
super(Gumbel, self).__init__()
self.__loc = loc
self.__scale = scale
@cache.cached_property
def loc(self):
return chainer.as_variable(self.__loc)
@cache.cached_property
def scale(self):
return chainer.as_variable(self.__scale)
@cache.cached_property
def _log_scale(self):
return exponential.log(self.scale)
@property
def batch_shape(self):
return self.loc.shape
@cache.cached_property
def entropy(self):
return self._log_scale + (EULER + 1)
@property
def event_shape(self):
return ()
@property
def _is_gpu(self):
return isinstance(self.loc.data, cuda.ndarray)
def log_prob(self, x):
y = (x - self.loc) / self.scale
return - self._log_scale - y - exponential.exp(-y)
@cache.cached_property
def mean(self):
return self.loc + EULER * self.scale
@property
def params(self):
return {'loc': self.loc, 'scale': self.scale}
def sample_n(self, n):
xp = chainer.backend.get_array_module(self.loc)
if xp is cuda.cupy:
eps = xp.random.gumbel(
size=(n,)+self.batch_shape, dtype=self.loc.dtype)
else:
eps = xp.random.gumbel(
size=(n,)+self.batch_shape).astype(self.loc.dtype)
noise = self.scale * eps + self.loc
return noise
@property
def support(self):
return 'real'
@cache.cached_property
def variance(self):
return (numpy.pi ** 2 / 6) * self.scale ** 2
@distribution.register_kl(Gumbel, Gumbel)
def _kl_gumbel_gumbel(dist1, dist2):
scale_1d2 = dist1.scale / dist2.scale
return (
dist2._log_scale
- dist1._log_scale
+ EULER * (scale_1d2 - 1.)
+ exponential.exp((dist2.loc - dist1.loc) / dist2.scale
+ lgamma.lgamma(scale_1d2 + 1.))
- 1
+ (dist1.loc - dist2.loc) / dist2.scale)