import tensorflow as tf

        mus, log_sigmas = params[:,:,:,0::2], params[:,:,:,1::2]
        z2, fldj = gaussianize(x2, mus, log_sigmas)
        return z2, fldj
    
    def _inverse(self, x1, z2, **kwargs):
        params = self.parameterizer(x1)
        mus, log_sigmas = params[:,:,:,0::2], params[:,:,:,1::2]
        x2, ildj = gaussianize(z2, mus, log_sigmas, inverse=tf.constant(True))
        return x2, ildj
    
def log_gaussianize(x, mus, log_sigmas, inverse=tf.constant(False)):
    """
    Standardize log normal random variable x using mus and log_sigmas.
    """
    if inverse:
        scales = tf.math.exp(log_sigmas)
        log_x = tf.math.log(x)
        ldj = log_x
        log_y = log_x*scales + mus
        ldj += log_sigmas
        z = tf.math.exp(log_y)
        return z, ldj
    else:
        scales = tf.math.exp(-log_sigmas)
        log_x = tf.math.log(x)
        ldj = -log_x
        log_y = (log_x - mus)*scales
        ldj -= log_sigmas
        z = tf.math.exp(log_y)
        return z, ldj
