class_docstr: Gate module which implements Top2Gating as described in Gshard_.
  ::
  
      gate = Top2Gate(model_dim, num_experts)
      l_aux, combine_weights, dispatch_mask = gate(input)
  
  .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf
  
  Args:
      model_dim (int):
          size of model embedding dimension
      num_experts (ints):
          number of experts in model
class_docstr_tok: ["Gate","module","which","implements","Top2Gating","as","described","in","Gshard_",".",":",":","gate","=","Top2Gate","(","model_dim",",","num_experts",")","l_aux",",","combine_weights",",","dispatch_mask","=","gate","(","input",")","..","Gshard_",":","https",":","\/\/arxiv.org\/pdf\/2006.16668.pdf","Args",":","model_dim","(","int",")",":","size","of","model","embedding","dimension","num_experts","(","ints",")",":","number","of","experts","in","model"]
code: def forward(self, input: torch.Tensor) -> Tuple[Tensor, Tensor, Tensor]:  # type: ignore
      logits = self.wg(input)
      return top2gating(logits)
  
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 2 for calling the Top2Gate obj's underlying member method code with example usage: obj.forward(input) and returns: top2gating