import torch


def gradient_detailed_overlap(grad, V, num_classes):
    """Similar to `gradient_overlap`, but here, we compute the overlap of the gradient
    with the individual eigenvectors. The gradiejt overlap can be written as

    overlap(grad, S_V) = ||P_V grad||^2 / ||grad||^2
                       = (sum_{c=1}^C (v_c.T grad)^2) / (grad.T grad)

    Here, we consider the individual projections (v_c.T 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_coeffs = torch.square(torch.matmul(V.T, grad))
    return proj_coeffs / (grad.T @ grad)