import torch


def compute_reconstruction_error(vae, data_loader, **kwargs):
    """
    Computes log p(x/z), which is the reconstruction error.

    Differs from the marginal log likelihood, but still gives good
    insights on the modeling of the data, and is fast to compute.
    """
    # Iterate once over the data and computes the reconstruction error
    log_lkl = {}
    for tensors in data_loader:
        loss_kwargs = dict(kl_weight=1)
        _, _, losses = vae(tensors, loss_kwargs=loss_kwargs)
        for key, value in losses._reconstruction_loss.items():
            if key in log_lkl:
                log_lkl[key] += torch.sum(value).item()
            else:
                log_lkl[key] = 0.0

    n_samples = len(data_loader.indices)
    for key, value in log_lkl.items():
        log_lkl[key] = log_lkl[key] / n_samples
        log_lkl[key] = -log_lkl[key]
    return log_lkl