import torch


def recons_probability(x_hat):
    """
    Returns the sum of the probabilities of each pixel in the image

    INPUT:
    x_hat: Tensor (B, C, H, W) -> reconstruction

    OUTPUT:
    recons_probability: Tensor (B, 1) -> sum of of the reconstruction
                                         probabilities in the image
    """

    probs = -torch.log(x_hat)
    sum_probs = torch.sum(probs, dim=(1, 2, 3))
    return sum_probs