def sample_z()

in qlearn/commun/local_mnf_layer.py [0:0]

• 30 lines of code
• 7 McCabe index (conditional complexity)

    def sample_z(self, batch_size, kl=True, same_noise=False):
        if self.training:
            if batch_size > 1:
                assert kl == False
                qzero_std = torch.exp(0.5 * self.qzero_logvar)
                qzero_std = qzero_std.expand(batch_size, self.in_features)
                z_mu = self.qzero_mu.expand(batch_size, self.in_features)
                if same_noise:
                    epsilon_z = self.epsilon_z.expand(batch_size, self.in_features)
                else:
                    epsilon_z = Variable(torch.randn(batch_size, self.in_features))
                    if self.use_cuda:
                        epsilon_z = epsilon_z.cuda()

                z = z_mu + qzero_std * epsilon_z
                z  = self.flow_q(z, kl=False)
                return z
            if batch_size ==  1:
                qzero_std = torch.exp(0.5 * self.qzero_logvar)
                z =  self.qzero_mu + qzero_std * self.epsilon_z
                if kl:
                    z, logdets = self.flow_q(z, kl=True)
                    return z, logdets
                else:
                    z = self.flow_q(z, kl=False)
                    return z
        else:
            assert kl == False
            z = self.qzero_mu
            z = self.flow_q(z, kl=False)
            return z