def random_masking(sequence, noise, len_keep, attention_masks=None):
    """
    Perform random masking by per-sample shuffling on frame-level. Per-sample shuffling is done by argsort random
    noise. sequence: [batch_size, seq_len, hidden_dim], sequence
    """

    batch_size, seq_len, hidden_dim = sequence.shape

    # sort noise for each sample
    ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
    ids_restore = torch.argsort(ids_shuffle, dim=1)

    # keep the first subset
    ids_keep = ids_shuffle[:, :len_keep]
    sequence_masked = torch.gather(sequence, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, hidden_dim))

    # generate the binary mask: 0 is keep, 1 is remove
    label_masks = torch.ones([batch_size, seq_len], device=sequence.device)
    label_masks[:, :len_keep] = 0
    # unshuffle to get the binary mask
    label_masks = torch.gather(label_masks, dim=1, index=ids_restore)

    if attention_masks is not None:
        label_masks *= attention_masks
        attention_masks = torch.gather(attention_masks, dim=1, index=ids_keep)

    return sequence_masked, attention_masks, label_masks, ids_restore

    Perform random masking by per-sample shuffling on frame-level. Per-sample shuffling is done by argsort random
    noise. sequence: [batch_size, seq_len, hidden_dim], sequence
    
dpkd/transformers/src/transformers/models/tvlt/modeling_tvlt.py