code: def gumbel_rsample(shape: Tuple, device: torch.device) -> Tensor:
      gumbel = gumbel_map.get(device)
      if gumbel is None:
          one = torch.tensor(1.0, device=device)
          zero = torch.tensor(0.0, device=device)
          gumbel = torch.distributions.gumbel.Gumbel(zero, one).rsample  # type: ignore
          gumbel_map[device] = gumbel
      return gumbel(shape)
  
file_path: fairscale/fairscale/nn/moe/top2gate.py
imports: from typing import Callable, Dict, Tuple
  import torch
  from torch import Tensor
  import torch.nn.functional
label_desc: Use image node_id 1 for calling a global function with example usage: gumbel_rsample(shape, device) and returns: gumbel