def _log_cdf(self, x):
# The CDF is (p**x * (1 - p)**(1 - x) + p - 1) / (2 * p - 1).
# We do this computation in logit space to be more numerically stable.
# p**x * (1- p)**(1 - x) becomes
# 1 / (1 + exp(-logits))**x *
# exp(-logits * (1 - x)) / (1 + exp(-logits)) ** (1 - x) =
# exp(-logits * (1 - x)) / (1 + exp(-logits))
# p - 1 becomes -exp(-logits) / (1 + exp(-logits))
# Thus the whole numerator is
# (exp(-logits * (1 - x)) - exp(-logits)) / (1 + exp(-logits))
# The denominator is (1 - exp(-logits)) / (1 + exp(-logits))
# Putting it all together, this gives:
# (exp(-logits * (1 - x)) - exp(-logits)) / (1 - exp(-logits)) =
# (exp(logits * x) - 1) / (exp(logits) - 1)
logits = self._logits_parameter_no_checks()
# For logits < 0, we can directly use the expression.
safe_logits = tf.where(logits < 0., logits, -1.)
result_negative_logits = (
tfp_math.log1mexp(tf.math.multiply_no_nan(safe_logits, x)) -
tfp_math.log1mexp(safe_logits))
# For logits > 0, to avoid infs with large arguments we rewrite the
# expression. Let z = log(exp(logits) - 1)
# log_cdf = log((exp(logits * x) - 1) / (exp(logits) - 1))
# = log(exp(logits * x) - 1) - log(exp(logits) - 1)
# = log(exp(logits * x) - 1) - log(exp(z))
# = log(exp(logits * x - z) - exp(-z))
# Because logits > 0, logits * x - z > -z, so we can pull it out to get
# = log(exp(logits * x - z) * (1 - exp(-logits * x)))
# = logits * x - z + tf.math.log(1 - exp(-logits * x))
dtype = dtype_util.as_numpy_dtype(x.dtype)
eps = np.finfo(dtype).eps
# log(exp(logits) - 1)
safe_logits = tf.where(logits > 0., logits, 1.)
z = tf.where(safe_logits > -np.log(eps), safe_logits,
tf.math.log(tf.math.expm1(safe_logits)))
result_positive_logits = tf.math.multiply_no_nan(
safe_logits, x) - z + tfp_math.log1mexp(
-tf.math.multiply_no_nan(safe_logits, x))
result = tf.where(
logits < 0., result_negative_logits,
tf.where(logits > 0., result_positive_logits, tf.math.log(x)))
# Finally, handle the case where `logits` and `p` are on the boundary,
# as the above expressions can result in ratio of `infs` in that case as
# well.
result = tf.where(tf.math.equal(logits, np.inf), dtype(-np.inf),
result)
result = tf.where(
(tf.math.equal(logits, -np.inf) & tf.math.not_equal(x, 0.)) |
(tf.math.equal(logits, np.inf) & tf.math.equal(x, 1.)),
tf.zeros_like(logits), result)
result = tf.where(x < 0., dtype(-np.inf),
tf.where(x > 1., tf.zeros_like(x), result))
return result
def _quantile(self, p, logits=None):
if logits is None:
logits = self._logits_parameter_no_checks()
logp = tf.math.log(p)
# The expression for the quantile function is:
# log(1 + (e^s - 1) * p) / s, where s is `logits`. When s is large,
# the e^s sub-term becomes increasingly ill-conditioned. However,
# since the numerator tends to s, we can reformulate the s > 0 case
# as a offset from 1, which is more accurate. Coincidentally,
# this eliminates a ratio of infinities problem when `s == +inf`.
safe_negative_logits = tf.where(logits < 0., logits, -1.)
safe_positive_logits = tf.where(logits > 0., logits, 1.)
result = tf.where(
logits > 0.,
1. + tfp_math.log_add_exp(
logp + tfp_math.log1mexp(safe_positive_logits),
tf.math.negative(safe_positive_logits)) / safe_positive_logits,
tf.math.log1p(
tf.math.expm1(safe_negative_logits) * p) / safe_negative_logits)
# When logits is zero, we can simplify
# log(1 + (e^s - 1) * p) / s ~= log(1 + s * p) / s ~= s * p / s = p
# Specifically, when logits is zero, the naive computation produces a NaN.
result = tf.where(tf.math.equal(logits, 0.), p, result)
# Finally, handle the case where `logits` and `p` are on the boundary,
# as the above expressions can result in ratio of `infs` in that case as
# well.
return tf.where(
(tf.math.equal(logits, -np.inf) & tf.math.equal(logp, 0.)) |
(tf.math.equal(logits, np.inf) & tf.math.is_inf(logp)),
tf.ones_like(logits),
result)
def _quantile(self, p, probs=None):
if probs is None:
probs = self._probs_parameter_no_checks()
cut_probs = self._cut_probs(probs)
cut_logits = tf.math.log(cut_probs) - tf.math.log1p(-cut_probs)
logp = tf.math.log(p)
# The expression for the quantile function is:
# log(1 + (e^s - 1) * p) / s, where s is `cut_logits`. When s is large,
# the e^s sub-term becomes increasingly ill-conditioned. However,
# since the numerator tends to s, we can reformulate the s > 0 case
# as a offset from 1, which is more accurate. Coincidentally,
# this eliminates a ratio of infinities problem when `s == +inf`.
result = tf.where(
cut_logits > 0.,
1. + tfp_math.log_add_exp(logp + tfp_math.log1mexp(cut_logits),
-cut_logits) / cut_logits,
tf.math.log1p(tf.math.expm1(cut_logits) * p) / cut_logits)
# Finally, handle the case where `cut_logits` and `p` are on the boundary,
# as the above expressions can result in ratio of `infs` in that case as
# well.
result = tf.where(
(tf.math.equal(cut_probs, 0.) & tf.math.equal(logp, 0.)) |
(tf.math.equal(cut_probs, 1.) & tf.math.is_inf(logp)),
tf.ones_like(cut_probs), result)
return tf.where((probs < self._lims[0]) | (probs > self._lims[1]),
result, p)
def _log_cdf(self, x):
# Going through the survival function is more accurate when conc is near
# zero, because it amounts to computing the (1 + conc * z)**(-1 / conc)
# term in log-space with log1p.
# tfp_math.log1mexp(a) accurately computes log(1 - exp(-|a|)). The negation
# and the absolute value are fine here because the log survival function is
# always non-positive.
return tfp_math.log1mexp(self._log_survival_function(x))
def _log_cdf(self, x):
probs = self._probs_parameter_no_checks()
if not self.validate_args:
# Whether or not x is integer-form, the following is well-defined.
# However, scipy takes the floor, so we do too.
x = tf.floor(x)
return tf.where(x < 0.,
dtype_util.as_numpy_dtype(x.dtype)(-np.inf),
tfp_math.log1mexp((1. + x) * tf.math.log1p(-probs)))
def categorical_log_probs(self):
"""Log probabilities for the `K + 1` sequential categories."""
cutpoints = tf.convert_to_tensor(self.cutpoints)
loc = tf.convert_to_tensor(self.loc)
num_cat = self._num_categories()
# For the StoppingRatioLogistic, we have:
# P(X = c; X >= c, cutpoints, loc) = sigmoid(cutpoints[c] - loc)
# Given these conditional probabilities, we would like to retrieve
# P(X = c; cutpoints, loc).
# Let F(c) = P(X = c; X >= c, cutpoints, loc) and
# G(c) = P(X = c; cutpoints, loc)
# Conditional probabilities. These are log(F(k)) and log(1 - F(k))
conditional_log_probs = tf.math.log_sigmoid(cutpoints -
loc[..., tf.newaxis])
conditional_log_probs_complement = tfp_math.log1mexp(
conditional_log_probs)
# Note that F(0) = G(0).
# G(1) = P(X = 1; cutpoints, loc) =
# P(X = 1; X >= 1, cutpoints, loc) * P(X >= 1) = F(1) * (1 - G(0))
# G(2) = P(X = 2; cutpoints, loc) =
# P(X = 2; X >= 2, cutpoints, loc) * P(X >= 2) = F(2) * (1 - G(0) - G(1))
# In general, G(k) = F(k) * (1 - \sum_{k-1} G(i))
# We rewrite this recurrence in terms of F(k)
# G(1) = F(1) * (1 - G(0)) = F(1) * (1 - F(0))
# G(2) = F(2) * (1 - G(0) - G(1)) = (1 - F(0) - F(1) * (1 - F(0))
# = F(2) * (1 - F(0)) * (1 - F(1))
# G(k) = F(k) * \prod_{k-1} (1 - F(i))
# log(F(k)) + log(\prod (1 - F(i)))
categorical_log_probs = conditional_log_probs + tf.math.cumsum(
conditional_log_probs_complement[..., :(num_cat - 1)],
axis=-1,
exclusive=True)
# Finally we need to handle the last category.
return tf.concat([
categorical_log_probs,
tf.math.reduce_sum(conditional_log_probs_complement[..., :num_cat],
axis=-1,
keepdims=True)
],
axis=-1)
def _log_cdf(self, x):
return tfp_math.log1mexp(self.concentration0 *
tf.math.log1p(-x**self.concentration1))
def _forward(self, x):
with tf.control_dependencies(self._maybe_assert_valid_x(x)):
rate = tf.convert_to_tensor(self.rate)
log1mexpx = tfp_math.log1mexp(-rate * x)
return tf.math.exp(log1mexpx -
tf.math.exp(-rate * x) / self.concentration)
def _log_cdf(self, x):
return tfp_math.log1mexp(self._log_survival_function(x))