import tensorflow as tf

        The forward pass scales to a standard log normal with mu=0, sigma=1 by computing:
        exp(Z) = (X / exp(mu))^(1/sigma)
        """
        params = self.parameterizer(x1)
        mus, log_sigmas = params[:,:,:,0::2], params[:,:,:,1::2]
        # compute softplus activation
        z2, ldj = log_gaussianize(x2, mus, log_sigmas)
        z2 = tf.where(x2 > self.epsilon, z2, x2)
        ldj = tf.where(x2 > self.epsilon, ldj, tf.zeros_like(ldj))
        return z2, tf.math.reduce_sum(ldj, axis=[1,2,3])
    
    def _inverse(self, x1, z2, **kwargs):
        params = self.parameterizer(x1)
        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])
