import tensorflow as tf

        mus, log_sigmas = params[:,:,:,0::2], params[:,:,:,1::2]
        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])
    
def half_gaussianize(x, log_sigmas, inverse=tf.constant(False)):
    if inverse:
        z = tf.math.exp(log_sigmas)*x
        ldj = tf.math.reduce_sum(log_sigmas, axis=[1,2,3])
    else:
        z = x*tf.math.exp(-log_sigmas)
        ldj = -tf.math.reduce_sum(log_sigmas, axis=[1,2,3])
    return z, ldj

class HalfGaussianize(Parameterize):
