class_docstr: Gate module which implements Top2Gating as described in Gshard_.
  ::
  
      gate = TopKGate(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","=","TopKGate","(","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 __init__(
      self,
      model_dim: int,
      num_experts: int,
      k: int = 1,
      capacity_factor: float = 1.0,
      eval_capacity_factor: float = 1.0,
      min_capacity: int = 8,
      noisy_gate_policy: Optional[str] = None,
      drop_tokens: bool = True,
      use_rts: bool = True,
  ) -> None:
      super().__init__()
  
      # Only top-1 and top-2 are supported at the moment.
      if k != 1 and k != 2:
          raise ValueError(
              "Only top-1 and top-2 gatings are supported."
          )
      self.wg = torch.nn.Linear(
          model_dim, num_experts, bias=False
      ).float()
      self.k = k
      self.capacity_factor = capacity_factor
      self.eval_capacity_factor = eval_capacity_factor
      self.min_capacity = min_capacity
      self.noisy_gate_policy = noisy_gate_policy
      self.timers = SynchronizedWallClockTimer()
      self.wall_clock_breakdown = False
      self.gate_time = 0.0
      self.drop_tokens = drop_tokens
      self.use_rts = use_rts
  
imports: from deepspeed.utils.timer import SynchronizedWallClockTimer
  from deepspeed.utils import logger
  from typing import Callable, Dict, TYPE_CHECKING, Any, Optional, Tuple
  import torch
  from torch import Tensor
  from torch.nn import Module
  import torch.nn.functional
  from deepspeed.utils import groups
  from .mappings import drop_tokens, gather_tokens
  from deepspeed import comm
label_desc: Use image node_id 1 to create a new TopKGate object from inherited base classes: Module with example: obj = TopKGate(model_dim, num_experts, k, capacity_factor, eval_capacity_factor, min_capacity, noisy_gate_policy, drop_tokens, use_rts)