
from __future__ import absolute_import
import math

#overrides
import torch

from allennlp.modules.similarity_functions.similarity_function import SimilarityFunction


class DotProductSimilarity(SimilarityFunction):
    u"""
    This similarity function simply computes the dot product between each pair of vectors, with an
    optional scaling to reduce the variance of the output elements.

    Parameters
    ----------
    scale_output : ``bool``, optional
        If ``True``, we will scale the output by ``math.sqrt(tensor.size(-1))``, to reduce the
        variance in the result.
    """
    def __init__(self, scale_output       = False)        :
        super(DotProductSimilarity, self).__init__()
        self._scale_output = scale_output

    #overrides
    def forward(self, tensor_1              , tensor_2              )                :
        result = (tensor_1 * tensor_2).sum(dim=-1)
        if self._scale_output:
            result *= math.sqrt(tensor_1.size(-1))
        return result

DotProductSimilarity = SimilarityFunction.register(u"dot_product")(DotProductSimilarity)
