import torch


def gradient_overlap(grad, V, num_classes):
    """Compute the overlap of the gradient and an eigenspace.

    Mathematical background: We consider a space, spanned by a set of orthonormal
    vectors S_V = span(v_1, ..., v_C). This spaces is represented by the (D x C)-matrix
    V = (v_1, ..., v_C). The respective projection matrix are P_V = V @ V.T.

    As in the paper "GRADIENT DESCENT HAPPENS IN A TINY SUBSPACE" (see
    https://arxiv.org/pdf/1812.04754.pdf), we define the overlap (a number in [0, 1])

    overlap(grad, S_V) = ||P_V grad||^2 / ||grad||^2
    """

    # Check that grad, V have the correct shape
    D, C = V.shape
    assert C == num_classes, "V doesn't have `num_classes` columns"
    assert grad.numel() == D, f"grad does not have {D} entries"

    proj_grad = torch.matmul(V, torch.matmul(V.T, grad))
    overlap = torch.dot(proj_grad, proj_grad) / torch.dot(grad, grad)
    return overlap.item()