code: def bert_sample_tokens(
      reserved_length: int,
      seq_length: int,
      batch_size: int,
      layers: int = 1,
      device: str = "cpu",
      attn_mask: torch.Tensor = None,
  ):
      assert attn_mask is not None
      prob_dist = torch.ones(
          (layers * batch_size, seq_length), device=device
      )
      sampled_indices = torch.multinomial(prob_dist, reserved_length)
  
      sampled_indices = sampled_indices.reshape(
          layers, batch_size, reserved_length
      ).to(torch.int32)
      global random_ltd_module
      if random_ltd_module is None:
          random_ltd_module = RandomLTDBuilder().load()
  
      sampled_indices = random_ltd_module.token_sort_(
          sampled_indices, seq_length
      )
      dtype = sampled_indices.dtype
  
      sampled_indices = sampled_indices.to(torch.long)
      new_mask = []
      for l in range(layers):
          tmp_mask_list = []
          for i in range(batch_size):
              mask_tmp = attn_mask[
                  i : i + 1, :, sampled_indices[l][i], :
              ]
              tmp_mask_list.append(
                  mask_tmp[:, :, :, sampled_indices[l][i]]
              )
          new_mask.append(torch.cat(tmp_mask_list, dim=0))
  
      return sampled_indices.to(dtype), new_mask
  
file_path: deepspeed/deepspeed/ops/random_ltd/dropping_utils.py
imports: import torch
  from deepspeed.ops.op_builder import RandomLTDBuilder
label_desc: Use image node_id 2 for calling a global function with example usage: bert_sample_tokens(reserved_length, seq_length, batch_size, layers, device, attn_mask) and returns: sampled_indices,  new_mask