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)