Cacher is the code snippet organizer for pro developers

class Encoder(torch.nn.Module):
    def __init__(self, in_dim, h_dim, z_dim):
        super(Encoder, self).__init__()
        self.h_net = nn.Sequential(
            nn.Linear(in_dim, h_dim),
            nn.Tanh()
        )
        self.mu_linear = nn.Linear(h_dim, z_dim)
        self.log_sigma_square_linear = nn.Linear(h_dim, z_dim)

    def forward(self, x):
        # x = batch x in_dim
        h = self.h_net(x)  # h = batch x h_dim
        mu = self.mu_linear(h)  # mu = batch x z_dim
        log_sigma_square = self.log_sigma_square_linear(h)  # log(sigma^2) = batch x z_dim

        return mu, log_sigma_square
        
        
class Decoder(torch.nn.Module):
    def __init__(self, z_dim, h_dim, out_dim):
        super(Decoder, self).__init__()
        self.net = self.sequential = nn.Sequential(
            nn.Linear(z_dim, h_dim),
            nn.Tanh(),
            nn.Linear(h_dim, out_dim),
            nn.Sigmoid() # so the input need to within (0-1)?
        )

    def forward(self, z):
        # z = batch x z_dim
        return self.net(z)  # batch x out_dim
        
class VAE(torch.nn.Module):
    latent_dim = 8

    def __init__(self, input_dim, enc_h_dim, z_dim, dec_h_dim, out_dim):
        super(VAE, self).__init__()
        assert input_dim == out_dim
        self.encoder = Encoder(input_dim, enc_h_dim, z_dim)
        self.decoder = Decoder(z_dim, dec_h_dim, out_dim)

    def _sample_latent(self, z_mu, log_sigma_square):
        """
        Return the latent normal sample z ~ N(mu, sigma^2)
        """
        z_sigma = torch.exp(0.5 * log_sigma_square)
        gaussian = torch.randn(z_sigma.size())
        return z_mu + z_sigma * Variable(gaussian, requires_grad=False)  # Reparameterization trick

    def forward(self, input):
        # input = batch x in_dim
        z_mu, log_sigma_square = self.encoder(input)  # batch x z_dim
        z = self._sample_latent(z_mu, log_sigma_square)  # batch x z_dim
        return self.decoder(z), z_mu, log_sigma_square  # batch x out_dim; batch x z_dim

    def decode(self, z):

        return self.decoder(z)
        
def vae_loss(out_x, x, z_mu, log_sigma_square):
    # batch x dim
    # loss = kl + cross_entropy
    ## batch
    batch = x.size(0)

    #BCE = F.binary_cross_entropy(out_x, x.view(-1, 784))  # False: sum of all the element in batch x dim
    BCE = F.mse_loss(out_x, x.view(-1, 784))
    # BCE = BCE / batch
    ''' 
    When the prior p(z)=N(0,1) and posterior approximation q(z|x) is also Gaussian,
    the KLD(q|p) is as below.
    '''
    KLD = -0.5 * torch.sum(1 + log_sigma_square - torch.mul(z_mu, z_mu) - torch.exp(log_sigma_square))
    KLD = KLD / (batch * 28 * 28)
    # KLD per picture

    return KLD + BCE