def _forward_log_det_jacobian(self, x):
x -= 0.5
fractional_part = x - tf.math.floor(x)
inner_part = (fractional_part - 0.5) / self.temperature
offset = (tf.math.log(self.temperature) -
tf.nn.softplus(0.5 / self.temperature) +
distribution_util.softplus_inverse(0.5 / self.temperature))
return -tf.nn.softplus(-inner_part) - tf.nn.softplus(
inner_part) - offset
def _inverse(self, y):
if self.hinge_softness is None:
return distribution_util.softplus_inverse(y)
hinge_softness = tf.cast(self.hinge_softness, y.dtype)
return hinge_softness * distribution_util.softplus_inverse(
y / hinge_softness)