def _inverse(self, y): with tf.control_dependencies(self._maybe_assert_valid_y(y)): concentration = tf.convert_to_tensor(self.concentration) reciprocal_concentration = tf.math.reciprocal(concentration) z = -tfp_math.lambertw(reciprocal_concentration * tf.math.exp( reciprocal_concentration + tf.math.log(y))) * concentration # Due to numerical instability, when y approaches 1, this expression # can be less than -1. We clip the value to prevent that. z = tf.clip_by_value(z, -1., np.inf) return -tf.math.log1p(z) / self.rate
def _inverse_log_det_jacobian(self, y): """Returns log of the Jacobian determinant of the inverse function.""" # Jacobian is log(W_delta(y, delta) / y) - log(1 + W(delta * y^2)). # Note that W_delta(y, delta) / y >= 0, since they share the same sign. # For numerical stability use log(abs()) - log(abs()). # See also Eq (11) and (31) of # https://www.hindawi.com/journals/tswj/2015/909231/ log_jacobian_term_nonzero = ( tf.math.log(tf.abs(_w_delta_squared(y, delta=self._tailweight))) - tf.math.log(tf.abs(y)) - tf.math.log(1. + tfp_math.lambertw(self._tailweight * y**2))) # If y = 0 the expression becomes log(0/0) - log(1 + 0), and the first term # equals log(1) = 0. Hence, for y = 0 the whole expression equals 0. return tf.where(tf.equal(y, 0.0), tf.zeros_like(y), log_jacobian_term_nonzero)
def _w_delta_squared(z, delta): """Applies W_delta transformation to the input. For a given z, `W_delta(z) = sign(z) * (W(delta * z^2)/delta)^0.5`. This transformation is defined in Equation (9) of [1]. Args: z: Input of the transformation. delta: Parameter delta of the transformation. Returns: The transformed Tensor with same shape and same dtype as `z`. """ delta = tf.convert_to_tensor(delta, dtype=z.dtype) z = tf.broadcast_to(z, ps.broadcast_shape(ps.shape(z), ps.shape(delta))) wd = tf.sign(z) * tf.sqrt(tfp_math.lambertw(delta * z**2) / delta) return tf.where(tf.equal(delta, 0.0), z, wd)