import tensorflow as tf

    """
    Implementation of parameterize for a half-Gaussian prior.
    """
    def __init__(self, input_shape=None, name='gaussianize', *args, **kwargs):
        super().__init__(*args, num_parameters=1, input_shape=input_shape, name=name, **kwargs)
        
    def _forward(self, x1, x2, **kwargs):
        log_sigmas = self.parameterizer(x1)
        z2, fldj = half_gaussianize(x2, log_sigmas)
        return z2, fldj
    
    def _inverse(self, x1, z2, **kwargs):
        log_sigmas = self.parameterizer(x1)
        x2, ildj = half_gaussianize(z2, log_sigmas, inverse=tf.constant(True))
        return x2, ildj
    
def exponentiate(x, log_lambdas, inverse=tf.constant(False)):
    if not inverse:
        z = tf.math.exp(log_lambdas)*x
        ldj = tf.math.reduce_sum(log_lambdas, axis=[1,2,3])
    else:
        z = x*tf.math.exp(-log_lambdas)
        ldj = -tf.math.reduce_sum(log_lambdas, axis=[1,2,3])
    return z, ldj

class Exponentiate(Parameterize):
