def imq_kernel(x, y, score_x, score_y, g=1, beta=0.5, return_kernel=False):
    """Compute the IMQ 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
    g : float
        Bandwidth
    beta : float
        Power of the kernel
    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
    """
    _, p = x.shape
    d = x[:, None, :] - y[None, :, :]
    dists = (d ** 2).sum(axis=-1)
    res = 1 + g * dists
    kxy = res ** (-beta)
    scores_d = score_x[:, None, :] - score_y[None, :, :]
    temp = d * scores_d
    dkxy = 2 * beta * g * (res) ** (-beta - 1) * temp.sum(axis=-1)
    d2kxy = 2 * (
        beta * g * (res) ** (-beta - 1) * p
        - 2 * beta * (beta + 1) * g ** 2 * dists * res ** (-beta - 2)
    )
    k_pi = score_x.mm(score_y.T) * kxy + dkxy + d2kxy
    if return_kernel:
        return k_pi, kxy
    return k_pi