import torch


def subspaces_overlap(U, V, num_classes):
    """Compute the overlap of two eigenspaces of the same dimensionality.

    Mathematical background: We consider two spaces, spanned by a set of C orthonormal
    vectors S_U = span(u_1, ..., u_C) and S_V = span(v_1, ..., v_C). These spaces are
    represented by the (D x C)-matrices U = (u_1, ..., u_C) and V = (v_1, ..., v_C). The
    respective projection matrices are P_U = U @ U.T and 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(S_U, S_V) = tr(P_U @ P_V) / sqrt(tr(P_U) tr(P_V))

    The numerator and denominator can be computed efficiently by using the trace's
    cyclic property. It holds tr(P_U @ P_V) = tr(W.T @ W) with W := U.T @ V (note that
    this is a "small" C x C matrix). This is normalized by sqrt(tr(P_U) tr(P_V)) = C.
    """

    # Check that U, V have the correct shape
    assert U.shape == V.shape, "U and V don't have the same shape"
    _, C = U.shape
    assert C == num_classes, "U doesn't have `num_classes` columns"

    W = U.T @ V
    overlap = torch.trace(W.T @ W) / num_classes
    return overlap.item()