def estimate_hsic(a_matrix, b_matrix, mode='biased'):
    """
    Estimates HSIC (if mode='biased') or pHSIC (if mode='plausible') between variables A and B.
    :param a_matrix:    torch.Tensor, a_matrix_ij = k(a_i,a_j), symmetric
    :param b_matrix:    torch.Tensor, b_matrix_ij = k(b_i,b_j), symmetric, must be the same size as a_matrix
    :param mode:        str, 'biased' (HSIC) or 'plausible' (pHSIC)
    :return: float, HSIC or pHSIC estimate
    """
    if mode == 'biased':
        a_vec = a_matrix.mean(dim=0)
        b_vec = b_matrix.mean(dim=0)
        # same as tr(HAHB)/m^2 for A=a_matrix, B=b_matrix, H=I - 11^T/m (centering matrix)
        return (a_matrix * b_matrix).mean() - 2 * (a_vec * b_vec).mean() + a_vec.mean() * b_vec.mean()
    if mode == 'plausible':
        # same as tr((A - mean(A))(B - mean(B)))/m^2
        return ((a_matrix - a_matrix.mean()) * b_matrix).mean()

    raise NotImplementedError('mode must be either biased or plausible, but %s was given' % mode)