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)