def linear_stein_kernel(x, y, score_x, score_y, return_kernel=False):
    """Compute the linear Stein kernel between x and y


    Parameters
    ----------
    x : torch.tensor, shape (n, p)
        Input particles
    y : torch.tensor, shape (n, p)
        Input particles
    score_x : torch.tensor, shape (n, p)
        The score of x
    score_y : torch.tensor, shape (n, p)
        The score of y
    return_kernel : bool
        whether the original kernel k(xi, yj) should also be returned

    Return
    ------
    stein_kernel : torch.tensor, shape (n, n)
        The linear Stein kernel
    kernel : torch.tensor, shape (n, n)
        The base kernel, only returned id return_kernel is True
    """
    n, d = x.shape
    kernel = x @ y.t()
    stein_kernel = (
        score_x @ score_y.t() * kernel + score_x @ x.t() + score_y @ y.t() + d
    )
    if return_kernel:
        return stein_kernel, kernel
    return stein_kernel