code: def get_mask_subset_with_prob(patched_input, prob):
      batch, seq_len, _, device = (
          *patched_input.shape,
          patched_input.device,
      )
      max_masked = math.ceil(prob * seq_len)
  
      rand = torch.rand((batch, seq_len), device=device)
      _, sampled_indices = rand.topk(max_masked, dim=-1)
  
      new_mask = torch.zeros((batch, seq_len), device=device)
      new_mask.scatter_(1, sampled_indices, 1)
      return new_mask.bool()
  
file_path: vit-pytorch/vit_pytorch/mpp.py
imports: import math
  import torch
  from torch import nn
  import torch.nn.functional
  from einops import rearrange, repeat, reduce
label_desc: Use image node_id 3 for calling a global function with example usage: get_mask_subset_with_prob(patched_input, prob) and returns: new_mask