def _log_cdf(self, y):
low = self._low
high = self._high
# Recall the promise:
# cdf(y) := P[Y <= y]
# = 1, if y >= high,
# = 0, if y < low,
# = P[X <= y], otherwise.
# P[Y <= j] = P[floor(Y) <= j] since mass is only at integers, not in
# between.
j = tf.floor(y)
result_so_far = self.distribution.log_cdf(j)
# Re-define values at the cutoffs.
if low is not None:
result_so_far = tf.where(
j < low,
dtype_util.as_numpy_dtype(self.dtype)(-np.inf), result_so_far)
if high is not None:
result_so_far = tf.where(j >= high, tf.zeros_like(result_so_far),
result_so_far)
return result_so_far
def _cdf(self, y):
low = self._low
high = self._high
# Recall the promise:
# cdf(y) := P[Y <= y]
# = 1, if y >= high,
# = 0, if y < low,
# = P[X <= y], otherwise.
# P[Y <= j] = P[floor(Y) <= j] since mass is only at integers, not in
# between.
j = tf.floor(y)
# P[X <= j], used when low < X < high.
result_so_far = self.distribution.cdf(j)
# Re-define values at the cutoffs.
if low is not None:
result_so_far = tf.where(j < low, tf.zeros_like(result_so_far),
result_so_far)
if high is not None:
result_so_far = tf.where(j >= high, tf.ones_like(result_so_far),
result_so_far)
return result_so_far
def loop_body(should_continue, k):
"""Resample the non-accepted points."""
# The range of U is chosen so that the resulting sample K lies in
# [0, tf.int64.max). The final sample, if accepted, is K + 1.
u = tf.random.uniform(
shape,
minval=minval_u,
maxval=maxval_u,
dtype=power.dtype,
seed=seed())
# Sample the point X from the continuous density h(x) \propto x^(-power).
x = self._hat_integral_inverse(u, power=power)
# Rejection-inversion requires a `hat` function, h(x) such that
# \int_{k - .5}^{k + .5} h(x) dx >= pmf(k + 1) for points k in the
# support. A natural hat function for us is h(x) = x^(-power).
#
# After sampling X from h(x), suppose it lies in the interval
# (K - .5, K + .5) for integer K. Then the corresponding K is accepted if
# if lies to the left of x_K, where x_K is defined by:
# \int_{x_k}^{K + .5} h(x) dx = H(x_K) - H(K + .5) = pmf(K + 1),
# where H(x) = \int_x^inf h(x) dx.
# Solving for x_K, we find that x_K = H_inverse(H(K + .5) + pmf(K + 1)).
# Or, the acceptance condition is X <= H_inverse(H(K + .5) + pmf(K + 1)).
# Since X = H_inverse(U), this simplifies to U <= H(K + .5) + pmf(K + 1).
# Update the non-accepted points.
# Since X \in (K - .5, K + .5), the sample K is chosen as floor(X + 0.5).
k = tf.where(should_continue, tf.floor(x + 0.5), k)
accept = (u <= self._hat_integral(k + .5, power=power) + tf.exp(
self._log_prob(k + 1, power=power)))
return [should_continue & (~accept), k]
def _log_prob(self, x, power=None):
# The log probability at positive integer points x is log(x^(-power) / Z)
# where Z is the normalization constant. For x < 1 and non-integer points,
# the log-probability is -inf.
#
# However, if interpolate_nondiscrete is True, we return the natural
# continuous relaxation for x >= 1 which agrees with the log probability at
# positive integer points.
#
# If interpolate_nondiscrete is False and validate_args is True, we check
# that the sample point x is in the support. That is, x is equivalent to a
# positive integer.
power = power if power is not None else tf.convert_to_tensor(self.power)
x = tf.cast(x, power.dtype)
if self.validate_args and not self.interpolate_nondiscrete:
x = distribution_util.embed_check_integer_casting_closed(
x, target_dtype=self.dtype, assert_positive=True)
log_normalization = tf.math.log(tf.math.zeta(power, 1.))
safe_x = tf.maximum(x if self.interpolate_nondiscrete else tf.floor(x), 1.)
y = -power * tf.math.log(safe_x)
log_unnormalized_prob = tf.where(
tf.equal(x, safe_x), y, dtype_util.as_numpy_dtype(y.dtype)(-np.inf))
return log_unnormalized_prob - log_normalization
def _log_unnormalized_prob(self, x, log_rate):
# The log-probability at negative points is always -inf.
# Catch such x's and set the output value accordingly.
safe_x = tf.maximum(x if self.interpolate_nondiscrete else tf.floor(x),
0.)
y = safe_x * log_rate - tf.math.lgamma(1. + safe_x)
return tf.where(tf.equal(x, safe_x), y,
dtype_util.as_numpy_dtype(y.dtype)(-np.inf))
def _log_prob(self, x):
with tf.control_dependencies(self._maybe_assert_valid_sample(x)):
probs = self._probs_parameter_no_checks()
if not self.validate_args:
# For consistency with cdf, we take the floor.
x = tf.floor(x)
safe_domain = tf.where(tf.equal(x, 0.), tf.zeros_like(probs),
probs)
return x * tf.math.log1p(-safe_domain) + tf.math.log(probs)
def _cdf(self, x):
with tf.control_dependencies(self._maybe_assert_valid_sample(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., tf.zeros_like(x), -tf.math.expm1(
(1. + x) * tf.math.log1p(-probs)))
def _cdf(self, x):
# CDF is the probability that the Poisson variable is less or equal to x.
# For fractional x, the CDF is equal to the CDF at n = floor(x).
# For negative x, the CDF is zero, but tf.igammac gives NaNs, so we impute
# the values and handle this case explicitly.
safe_x = tf.maximum(x if self.interpolate_nondiscrete else tf.floor(x),
0.)
cdf = tf.math.igammac(1. + safe_x, self._rate_parameter_no_checks())
return tf.where(x < 0., tf.zeros_like(cdf), cdf)
def _cdf(self, x):
# CDF(x) at positive integer x is the probability that the Zipf variable is
# less than or equal to x; given by the formula:
# CDF(x) = 1 - (zeta(power, x + 1) / Z)
# For fractional x, the CDF is equal to the CDF at n = floor(x).
# For x < 1, the CDF is zero.
# If interpolate_nondiscrete is True, we return a continuous relaxation
# which agrees with the CDF at integer points.
power = tf.convert_to_tensor(self.power)
x = tf.cast(x, power.dtype)
safe_x = tf.maximum(x if self.interpolate_nondiscrete else tf.floor(x), 0.)
cdf = 1. - (
tf.math.zeta(power, safe_x + 1.) / tf.math.zeta(power, 1.))
return tf.where(x < 1., tf.zeros_like(cdf), cdf)
def _sample_n(self, n, seed=None):
# Uniform variates must be sampled from the open-interval `(0, 1)` rather
# than `[0, 1)`. To do so, we use
# `np.finfo(dtype_util.as_numpy_dtype(self.dtype)).tiny`
# because it is the smallest, positive, 'normal' number. A 'normal' number
# is such that the mantissa has an implicit leading 1. Normal, positive
# numbers x, y have the reasonable property that, `x + y >= max(x, y)`. In
# this case, a subnormal number (i.e., np.nextafter) can cause us to sample
# 0.
probs = self._probs_parameter_no_checks()
sampled = tf.random.uniform(
tf.concat([[n], tf.shape(probs)], 0),
minval=np.finfo(dtype_util.as_numpy_dtype(self.dtype)).tiny,
maxval=1.,
seed=seed,
dtype=self.dtype)
return tf.floor(tf.math.log(sampled) / tf.math.log1p(-probs))
def _mode(self):
return tf.floor(self._rate_parameter_no_checks())
def _mode(self):
return tf.floor(
(1. + self._total_count) * self._probs_parameter_no_checks())
def _mode(self):
total_count = tf.convert_to_tensor(self.total_count)
adjusted_count = tf.where(1. < total_count, total_count - 1.,
tf.zeros_like(total_count))
return tf.floor(adjusted_count *
tf.exp(self._logits_parameter_no_checks()))
