def _log_normalization(self, concentration=None):
"""Computes the log-normalizer of the distribution."""
if concentration is None:
concentration = tf.convert_to_tensor(self.concentration)
event_dim = tf.cast(self._event_shape_tensor()[0], self.dtype)
safe_conc = tf.where(concentration > 0, concentration,
tf.ones_like(concentration))
safe_lognorm = ((event_dim / 2 - 1) * tf.math.log(safe_conc) -
(event_dim / 2) * np.log(2 * np.pi) -
tfp_math.log_bessel_ive(event_dim / 2 - 1, safe_conc) -
tf.abs(safe_conc))
log_nsphere_surface_area = (
np.log(2.) + (event_dim / 2) * np.log(np.pi) -
tf.math.lgamma(tf.cast(event_dim / 2, self.dtype)))
return tf.where(concentration > 0, -safe_lognorm,
log_nsphere_surface_area)
def _log_prob(self, x):
# The log-probability at negative points is always -inf.
# Catch such x's and set the output value accordingly.
lr1, r1, lr2, r2 = self._all_rate_parameters()
safe_x = tf.floor(x) if self.force_probs_to_zero_outside_support else x
y = tf.math.multiply_no_nan(0.5 * (lr1 - lr2), safe_x)
numpy_dtype = dtype_util.as_numpy_dtype(y.dtype)
# When both rates are zero, the above computation gives a NaN, whereas
# it should give zero.
y = tf.where(
tf.math.equal(r1, 0.) & tf.math.equal(r2, 0.), numpy_dtype(0.), y)
y = y + tfp_math.log_bessel_ive(safe_x, 2. * tf.math.sqrt(
r1 * r2)) - tf.math.square(tf.math.sqrt(r1) - tf.math.sqrt(r2))
y = tf.where(tf.math.equal(x, safe_x), y, numpy_dtype(-np.inf))
if self.force_probs_to_zero_outside_support:
# Ensure the gradient wrt `rate` is zero at non-integer points.
y = tf.where((y < 0.) & tf.math.is_inf(y), numpy_dtype(-np.inf), y)
return y