code: def weighted_index(self, dim=None):
      """
      Returns a tensor with entries that are one-hot along dimension `dim`.
      These one-hot entries are set at random with weights given by the input
      `self`.
  
      Examples::
  
          >>> encrypted_tensor = MPCTensor(torch.tensor([1., 6.]))
          >>> index = encrypted_tensor.weighted_index().get_plain_text()
          # With 1 / 7 probability
          torch.tensor([1., 0.])
  
          # With 6 / 7 probability
          torch.tensor([0., 1.])
      """
      if dim is None:
          return self.flatten().weighted_index(dim=0).view(self.size())
  
      x = self.cumsum(dim)
      max_weight = x.index_select(
          dim, torch.tensor(x.size(dim) - 1, device=self.device)
      )
      r = (
          crypten.rand(max_weight.size(), device=self.device)
          * max_weight
      )
  
      gt = x.gt(r)
      shifted = gt.roll(1, dims=dim)
      shifted.data.index_fill_(
          dim, torch.tensor(0, device=self.device), 0
      )
  
      return gt - shifted
  
imports: import crypten
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: weighted_index(dim) and returns: unknown, self