import torch


def reweight(msa1hot, cutoff):
    """Reweigh 1-hot MSA based on a `cutoff`."""
    id_min = msa1hot.shape[1] * cutoff
    id_mtx = torch.einsum("ikl,jkl->ij", msa1hot, msa1hot)
    id_mask = id_mtx > id_min
    w = 1.0 / id_mask.float().sum(dim=-1)
    return w