add gev distribution and test

tensorflow_probability/python/bijectors/gev_cdf.py

@@ -120,39 +120,53 @@ def _is_increasing(cls):
   def _forward(self, x):
     loc = tf.convert_to_tensor(self.loc)
     scale = tf.convert_to_tensor(self.scale)
-    concentration = tf.convert_to_tensor(self.concentration)
+    conc = tf.convert_to_tensor(self.concentration)
     with tf.control_dependencies(
         self._maybe_assert_valid_x(
-            x, loc=loc, scale=scale, concentration=concentration)):
+            x, loc=loc, scale=scale, concentration=conc)):
       z = (x - loc) / scale
+
+      equal_zero = tf.equal(conc, 0.)
+      # deal with case that gradient is N/A when conc = 0
+      safe_conc = tf.where(equal_zero, 1., conc)
       t = tf.where(
-          tf.equal(concentration, 0.), tf.math.exp(-z),
-          tf.math.exp(-tf.math.log1p(z * concentration) / concentration))
+          equal_zero, tf.math.exp(-z),
+          tf.math.exp(-tf.math.log1p(z * safe_conc) / safe_conc))
       return tf.exp(-t)
 
   def _inverse(self, y):
     with tf.control_dependencies(self._maybe_assert_valid_y(y)):
       t = -tf.math.log(y)
 
       conc = tf.convert_to_tensor(self.concentration)
+
+      equal_zero = tf.equal(conc, 0.)
+      # deal with case that gradient is N/A when conc = 0
+      safe_conc = tf.where(equal_zero, 1., conc)
+
       z = tf.where(
-          tf.equal(conc, 0.), -tf.math.log(t),
-          tf.math.expm1(-tf.math.log(t) * conc) / conc)
+          equal_zero, -tf.math.log(t),
+          tf.math.expm1(-tf.math.log(t) * safe_conc) / safe_conc)
 
       return self.loc + self.scale * z
 
   def _forward_log_det_jacobian(self, x):
     loc = tf.convert_to_tensor(self.loc)
     scale = tf.convert_to_tensor(self.scale)
-    concentration = tf.convert_to_tensor(self.concentration)
+    conc = tf.convert_to_tensor(self.concentration)
     with tf.control_dependencies(
         self._maybe_assert_valid_x(
-            x, loc=loc, scale=scale, concentration=concentration)):
+            x, loc=loc, scale=scale, concentration=conc)):
       z = (x - loc) / scale
+
+      equal_zero = tf.equal(conc, 0.)
+      # deal with case that gradient is N/A when conc = 0
+      safe_conc = tf.where(equal_zero, 1., conc)
+
       log_t = tf.where(
-          tf.equal(concentration, 0.), -z,
-          -tf.math.log1p(z * concentration) / concentration)
-      return (tf.math.multiply_no_nan(concentration + 1., log_t) -
+          equal_zero, -z,
+          -tf.math.log1p(z * safe_conc) / safe_conc)
+      return (tf.math.multiply_no_nan(conc + 1., log_t) -
               tf.math.exp(log_t) - tf.math.log(scale))
 
   def _inverse_log_det_jacobian(self, y):

tensorflow_probability/python/distributions/BUILD

@@ -72,6 +72,7 @@ multi_substrate_py_library(
         ":generalized_pareto",
         ":geometric",
         ":gumbel",
+        ":gev",
         ":half_cauchy",
         ":half_normal",
         ":half_student_t",
@@ -740,6 +741,24 @@ multi_substrate_py_library(
     ],
 )
 
+multi_substrate_py_library(
+    name = "gev",
+    srcs = ["gev.py"],
+    deps = [
+        ":kullback_leibler",
+        ":transformed_distribution",
+        ":uniform",
+        # numpy dep,
+        # tensorflow dep,
+        "//tensorflow_probability/python/bijectors:gev_cdf",
+        "//tensorflow_probability/python/bijectors:identity",
+        "//tensorflow_probability/python/bijectors:invert",
+        "//tensorflow_probability/python/internal:distribution_util",
+        "//tensorflow_probability/python/internal:dtype_util",
+        "//tensorflow_probability/python/internal:tensor_util",
+    ],
+)
+
 multi_substrate_py_library(
     name = "half_cauchy",
     srcs = ["half_cauchy.py"],
@@ -2508,6 +2527,19 @@ multi_substrate_py_test(
     ],
 )
 
+multi_substrate_py_test(
+    name = "gev_test",
+    srcs = ["gev_test.py"],
+    jax_size = "medium",
+    deps = [
+        # numpy dep,
+        # scipy dep,
+        # tensorflow dep,
+        "//tensorflow_probability",
+        "//tensorflow_probability/python/internal:test_util",
+    ],
+)
+
 multi_substrate_py_test(
     name = "half_cauchy_test",
     srcs = ["half_cauchy_test.py"],

tensorflow_probability/python/distributions/__init__.py

@@ -53,6 +53,7 @@
 from tensorflow_probability.python.distributions.generalized_pareto import GeneralizedPareto
 from tensorflow_probability.python.distributions.geometric import Geometric
 from tensorflow_probability.python.distributions.gumbel import Gumbel
+from tensorflow_probability.python.distributions.gev import GeneralizedExtremeValue
 from tensorflow_probability.python.distributions.half_cauchy import HalfCauchy
 from tensorflow_probability.python.distributions.half_normal import HalfNormal
 from tensorflow_probability.python.distributions.half_student_t import HalfStudentT
@@ -184,6 +185,7 @@
     'GaussianProcessRegressionModel',
     'VariationalGaussianProcess',
     'Gumbel',
+    'GeneralizedExtremeValue',
     'HalfCauchy',
     'HalfNormal',
     'HalfStudentT',