@wraps(Prior.log_p_z)
    def log_p_z(self, z):
        return -0.5 * (
            T.log(2 * pi * T.exp(2 * self.log_prior_sigma)) +
            ((z - self.prior_mu) / T.exp(self.log_prior_sigma)) ** 2
        ).sum(axis=2)
dataset/ETHPy150Open lisa-lab/pylearn2/pylearn2/models/vae/prior.py/DiagonalGaussianPrior.log_p_z/original