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()))