Add Gumbel distribution. (#550)

* Add Gumbel distribution.

* Fix algebra

* Add missing - sign

* Add autogenerated documentation for Gumbel

* Add Kolmogorov-Smirnov test of samples from distributions for those we do not directly get from upstream

Author: daydreamt

docs/source/distributions.rst

@@ -85,6 +85,14 @@ Gamma
     :show-inheritance:
     :member-order: bysource
 
+Gumbel
+-----
+.. autoclass:: numpyro.distributions.continuous.Gumbel
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 GaussianRandomWalk
 ------------------
 .. autoclass:: numpyro.distributions.continuous.GaussianRandomWalk

numpyro/distributions/__init__.py

@@ -12,6 +12,7 @@
     Exponential,
     Gamma,
     GaussianRandomWalk,
+    Gumbel,
     HalfCauchy,
     HalfNormal,
     InverseGamma,
@@ -73,6 +74,7 @@
     'Gamma',
     'GammaPoisson',
     'GaussianRandomWalk',
+    'Gumbel',
     'HalfCauchy',
     'HalfNormal',
     'Independent',

numpyro/distributions/continuous.py

@@ -48,6 +48,9 @@
 from numpyro.util import copy_docs_from
 
 
+EULER_MASCHERONI = 0.5772156649015328606065120900824024310421
+
+
 @copy_docs_from(Distribution)
 class Beta(Distribution):
     arg_constraints = {'concentration1': constraints.positive, 'concentration0': constraints.positive}
@@ -326,6 +329,39 @@ def variance(self):
         return np.where(self.concentration <= 2, np.inf, a)
 
 
+@copy_docs_from(Distribution)
+class Gumbel(Distribution):
+    arg_constraints = {'loc': constraints.real, 'scale': constraints.positive}
+    support = constraints.real
+    reparametrized_params = ['loc', 'scale']
+
+    def __init__(self, loc=0., scale=1., validate_args=None):
+        self.loc, self.scale = promote_shapes(loc, scale)
+        batch_shape = lax.broadcast_shapes(np.shape(loc), np.shape(scale))
+
+        super(Gumbel, self).__init__(batch_shape=batch_shape,
+                                     validate_args=validate_args)
+
+    def sample(self, key, sample_shape=()):
+        standard_gumbel_sample = random.gumbel(key, shape=sample_shape + self.batch_shape + self.event_shape)
+        return self.loc + self.scale * standard_gumbel_sample
+
+    @validate_sample
+    def log_prob(self, value):
+        z = (value - self.loc) / self.scale
+        return -(z + np.exp(-z)) - np.log(self.scale)
+
+    @property
+    def mean(self):
+        return np.broadcast_to(self.loc + self.scale * EULER_MASCHERONI,
+                               self.batch_shape)
+
+    @property
+    def variance(self):
+        return np.broadcast_to(np.pi**2 / 6. * self.scale**2,
+                               self.batch_shape)
+
+
 @copy_docs_from(Distribution)
 class LKJ(TransformedDistribution):
     r"""

test/test_distributions.py

@@ -68,6 +68,7 @@ def _lowrank_mvn_to_scipy(loc, cov_fac, cov_diag):
     dist.Dirichlet: lambda conc: osp.dirichlet(conc),
     dist.Exponential: lambda rate: osp.expon(scale=np.reciprocal(rate)),
     dist.Gamma: lambda conc, rate: osp.gamma(conc, scale=1./rate),
+    dist.Gumbel: lambda loc, scale: osp.gumbel_r(loc=loc, scale=scale),
     dist.HalfCauchy: lambda scale: osp.halfcauchy(scale=scale),
     dist.HalfNormal: lambda scale: osp.halfnorm(scale=scale),
     dist.InverseGamma: lambda conc, rate: osp.invgamma(conc, scale=rate),
@@ -103,6 +104,9 @@ def _lowrank_mvn_to_scipy(loc, cov_fac, cov_diag):
     T(dist.Gamma, np.array([0.5, 1.3]), np.array([[1.], [3.]])),
     T(dist.GaussianRandomWalk, 0.1, 10),
     T(dist.GaussianRandomWalk, np.array([0.1, 0.3, 0.25]), 10),
+    T(dist.Gumbel, 0., 1.),
+    T(dist.Gumbel, 0.5, 2.),
+    T(dist.Gumbel, np.array([0., 0.5]), np.array([1., 2.])),
     T(dist.HalfCauchy, 1.),
     T(dist.HalfCauchy, np.array([1., 2.])),
     T(dist.HalfNormal, 1.),
@@ -937,3 +941,24 @@ def test_unpack_transform():
     z = transform.inv(y)
     assert_allclose(y['key'], x)
     assert_allclose(z, x)
+
+
+@pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS)
+def test_generated_sample_distribution(jax_dist, sp_dist, params,
+                                       N_sample=100_000,
+                                       key=random.PRNGKey(11)):
+    """ On samplers that we do not get directly from JAX, (e.g. we only get
+    Gumbel(0,1) but also provide samplers for Gumbel(loc, scale)), also test
+    agreement in the empirical distribution of generated samples between our
+    samplers and those from SciPy.
+    """
+
+    if jax_dist not in [dist.Gumbel]:
+        pytest.skip("{} sampling method taken from upstream, no need to"
+                    "test generated samples.".format(jax_dist.__name__))
+
+    jax_dist = jax_dist(*params)
+    if sp_dist and not jax_dist.event_shape and not jax_dist.batch_shape:
+        our_samples = jax_dist.sample(key, (N_sample,))
+        ks_result = osp.kstest(our_samples, sp_dist(*params).cdf)
+        assert ks_result.pvalue > 0.05