diff --git a/tensorflow_probability/python/distributions/exponential.py b/tensorflow_probability/python/distributions/exponential.py index a59ac6c32b..d8bfed874c 100644 --- a/tensorflow_probability/python/distributions/exponential.py +++ b/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) diff --git a/tensorflow_probability/python/distributions/geometric.py b/tensorflow_probability/python/distributions/geometric.py index e303463522..06318a2c8c 100644 --- a/tensorflow_probability/python/distributions/geometric.py +++ b/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()