jakobamb

import torch


def rankMe(z: torch.Tensor, epsilon=1e-7):
    """
    Implementation of RankMe by Garrido et al. 2023 (https://proceedings.mlr.press/v202/garrido23a.html)
    
    param z: a tensor of shape (B, N, K), where B is batch size.
    """