import torch


def cos_difference_mem(a, b, adjacency_matrix):
    """
    The difference of coupled states in an oscillator model.
    Vectorized, faster but costs more memory.
    :param a: the vector a
    :param b: the vector b
    :param adjacency_matrix: The adjacency matrix
    :return: The vector y with elements calculated as: \sum_j cos(a_j - b_i).
    """
    if len(a.shape) > 1:
        adjacency_matrix = adjacency_matrix.unsqueeze(0)

    res = (torch.cos((a.unsqueeze(-2) - b.unsqueeze(-1))) * adjacency_matrix).sum(-1)
    return res.squeeze(-1)