import torch


def attention_consistency_loss(attention_weights, reattention_weights):
    """Returns Attention consistency loss.

    Attention consistency loss is defined as the distance between the visual
    attention maps learned in the attention layer (learned from the fusion of
    the question and image features) and the attention map learned in the
    re-attention layer(learned from the fusion of the gleaned answer
    representation and image features). It is implemented here as defined in
    'Re-attention for Visual Question Answering' by Guo et al
    (https://ojs.aaai.org//index.php/AAAI/article/view/5338). Attention
    consistency loss is used to re-learn the importance of the visual features
    guided by which objects correspond to the learned answer,

    Args:
        attention_weights: visual attention weights learned in the attention
            layer, prior to obtaining the answer representation.
        reattention_weights: visual attention weights learned in the
            re-attention layer.

    Returns:
        attention consistency loss.
    """
    return torch.sum(
        torch.square(torch.subtract(attention_weights, reattention_weights))
    )