import torch


def sim(query, data, l2_normalize=False, aligned=True):
    """
    Batch dot-product similarity computed using matrix multiplication the
    hidden shapes must be broadcastable (numpy style)

    :param query: the input data to be compared  (batch_size, h, p) p = memory_size if aligned is True and p = content_size if aligned is False
    :param data: Input state  (batch_size, content_size, memory_size]
    :param l2_normalize: boolean, determines where to normalize the query and the data before the dot product
    :param aligned: boolean, determines whether to transpose data along the last two dimensions
    :return:  out[...,i,j] = sum_k q[...,i,k] * data_gen[...,j,k] for the default options

    """

    # data_gen.shape = hidden_shape_1 x M x N
    # query.shape = hidden_shape_2 x h x p, where:
    #        p = N if aligned is True and p = M if aligned is False
    # out[...,i,j] = sum_k q[...,i,k] * data_gen[...,j,k] for the default
    # options

    if aligned:  # transpose last 2 dims to enable matrix multiplication
        data = torch.transpose(data, -1, -2)

    assert query.size()[-1] == data.size()[-2]

    if l2_normalize:
        query = torch.nn.functional.normalize(query, dim=-1)
        data = torch.nn.functional.normalize(data, dim=-2)

    return torch.matmul(query, data)