import torch


def loss_loglik(y_mean, y_logvar, x):
    """
    Args:
      y_mean of shape (batch_size, 1, 28, 28): Predictive mean of the VAE reconstruction of x.
      y_logvar of shape (batch_size, 1, 28, 28): Predictive log-variance of the VAE reconstruction of x.
      x of shape (batch_size, 1, 28, 28): Training samples.
    
    Returns:
      loss (torch scalar): Expected log-likelihood loss.
    """
    # constant = (-n*K/2) * torch.log(2*np.pi)
    n = 1
    batch_size = y_mean.shape[0]
    loss = []

    for i in range(batch_size):        
        a= (x[i] - y_mean[i]) ** 2
        img_loss = 0.5 * (y_logvar[i].sum() + (a / y_logvar[i].exp()).sum())
        loss.append(img_loss)    
    loss = torch.stack(loss).mean()
    
    return loss