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