fix the case that value < 0 for exp and geo dist

tensorflow_probability/python/distributions/exponential.py

@@ -71,6 +71,7 @@ class Exponential(gamma.Gamma):
 
   def __init__(self,
                rate,
+               force_probs_to_zero_outside_support=False,
                validate_args=False,
                allow_nan_stats=True,
                name="Exponential"):
@@ -79,6 +80,17 @@ def __init__(self,
     Args:
       rate: Floating point tensor, equivalent to `1 / mean`. Must contain only
         positive values.
+      force_probs_to_zero_outside_support: Python `bool`. When `True`, negative
+        and non-integer values are evaluated "strictly": `cdf` returns
+        `0`, `sf` returns `1`, and `log_cdf` and `log_sf` correspond.  When
+        `False`, the implementation is free to save computation (and TF graph
+        size) by evaluating something that matches the Exponential cdf at
+        non-negative values `x` but produces an unrestricted result on
+        other inputs. In the case of Exponential distribution, the `cdf`
+        formula in this case happens to be the continuous function
+        `1 - exp(rate * value)`.
+        Note that this function is not itself a cdf function.
+        Default value: `False`.
       validate_args: Python `bool`, default `False`. When `True` distribution
         parameters are checked for validity despite possibly degrading runtime
         performance. When `False` invalid inputs may silently render incorrect
@@ -99,6 +111,8 @@ def __init__(self,
           rate,
           name="rate",
           dtype=dtype_util.common_dtype([rate], dtype_hint=tf.float32))
+      self._force_probs_to_zero_outside_support = (
+          force_probs_to_zero_outside_support)
       super(Exponential, self).__init__(
           concentration=1.,
           rate=self._rate,
@@ -120,12 +134,29 @@ def _parameter_properties(cls, dtype, num_classes=None):
   def rate(self):
     return self._rate
 
+  @property
+  def force_probs_to_zero_outside_support(self):
+    """Return 0 probabilities on non-integer inputs."""
+    return self._force_probs_to_zero_outside_support
+
   def _cdf(self, value):
-    return -tf.math.expm1(-self.rate * value)
+    cdf = -tf.math.expm1(-self.rate * value)
+
+    if self.force_probs_to_zero_outside_support:
+      # Set cdf = 0 when value is less than 0.
+      cdf = tf.where(value < 0., tf.zeros_like(cdf), cdf)
+
+    return cdf
 
   def _log_survival_function(self, value):
     rate = tf.convert_to_tensor(self._rate)
-    return self._log_prob(value, rate=rate) - tf.math.log(rate)
+    log_sf = self._log_prob(value, rate=rate) - tf.math.log(rate)
+
+    if self.force_probs_to_zero_outside_support:
+      # Set log_survival_function = 0 when value is less than 0.
+      log_sf = tf.where(value < 0., tf.zeros_like(log_sf), log_sf)
+
+    return log_sf
 
   def _sample_n(self, n, seed=None):
     rate = tf.convert_to_tensor(self.rate)

tensorflow_probability/python/distributions/geometric.py

@@ -59,6 +59,7 @@ class Geometric(distribution.Distribution):
   def __init__(self,
                logits=None,
                probs=None,
+               force_probs_to_zero_outside_support=False,
                validate_args=False,
                allow_nan_stats=True,
                name='Geometric'):
@@ -75,6 +76,16 @@ def __init__(self,
         represents the probability of success for independent Geometric
         distributions and must be in the range `(0, 1]`. Only one of `logits`
         or `probs` should be specified.
+      force_probs_to_zero_outside_support: Python `bool`. When `True`, negative
+        and non-integer values are evaluated "strictly": `log_prob` returns
+        `-inf`, `prob` returns `0`, and `cdf` and `sf` correspond.  When
+        `False`, the implementation is free to save computation (and TF graph
+        size) by evaluating something that matches the Geometric pmf at integer
+        values `k` but produces an unrestricted result on other inputs. In the
+        case of Geometric distribution, the `log_prob` formula in this case
+        happens to be the continuous function `k * log(1-probs) + log(probs)`.
+        Note that this function is not a normalized probability log-density.
+        Default value: `False`.
       validate_args: Python `bool`, default `False`. When `True` distribution
         parameters are checked for validity despite possibly degrading runtime
         performance. When `False` invalid inputs may silently render incorrect
@@ -95,6 +106,8 @@ def __init__(self,
           probs, dtype=dtype, name='probs')
       self._logits = tensor_util.convert_nonref_to_tensor(
           logits, dtype=dtype, name='logits')
+      self._force_probs_to_zero_outside_support = (
+          force_probs_to_zero_outside_support)
       super(Geometric, self).__init__(
           dtype=dtype,
           reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
@@ -122,6 +135,11 @@ def probs(self):
     """Input argument `probs`."""
     return self._probs
 
+  @property
+  def force_probs_to_zero_outside_support(self):
+    """Return 0 probabilities on non-integer inputs."""
+    return self._force_probs_to_zero_outside_support
+
   def _batch_shape_tensor(self):
     x = self._probs if self._logits is None else self._logits
     return ps.shape(x)
@@ -182,7 +200,16 @@ def _log_prob(self, x):
     if not self.validate_args:
       # For consistency with cdf, we take the floor.
       x = tf.floor(x)
-    return tf.math.xlog1py(x, -probs) + tf.math.log(probs)
+
+    log_probs = tf.math.xlog1py(x, -probs) + tf.math.log(probs)
+
+    if self.force_probs_to_zero_outside_support:
+      # Set log_prob = -inf when value is less than 0, ie prob = 0.
+      log_probs = tf.where(
+          x < 0.,
+          dtype_util.as_numpy_dtype(x.dtype)(-np.inf),
+          log_probs)
+    return log_probs
 
   def _entropy(self):
     logits, probs = self._logits_and_probs_no_checks()