import tensorflow as tf

        x2, ldj = log_gaussianize(z2, mus, log_sigmas, inverse=tf.constant(True))
        x2 = tf.where(z2 > self.epsilon, x2, z2)
        ldj = tf.where(z2 > self.epsilon, ldj, tf.zeros_like(ldj))
        return x2, tf.math.reduce_sum(ldj, axis=[1,2,3])
