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 forward(
      self,
      input: torch.Tensor,
      used_token: torch.Tensor = None,
      use_tutel: bool = False,
  ) -> Tuple[Tensor, Tensor, Tensor]:  # type: ignore
      if self.wall_clock_breakdown:
          self.timers(TOPK_GATE_TIMER).start()
  
      if self.wg.weight.dtype != torch.float32:
          self.wg = self.wg.float()
      input_fp32 = input.float()
      # input jittering
      if self.noisy_gate_policy == "Jitter" and self.training:
          input_fp32 = multiplicative_jitter(
              input_fp32, device=input.device
          )
      logits = self.wg(input_fp32)
  
      if self.k == 1:
          gate_output = top1gating(
              logits,
              self.capacity_factor
              if self.training
              else self.eval_capacity_factor,
              self.min_capacity,
              used_token,
              self.noisy_gate_policy if self.training else None,
              self.drop_tokens,
              self.use_rts,
              use_tutel,
          )
  
      else:
          gate_output = top2gating(
              logits,
              self.capacity_factor
              if self.training
              else self.eval_capacity_factor,
              self.min_capacity,
          )
  
      if self.wall_clock_breakdown:
          self.timers(TOPK_GATE_TIMER).stop()
          self.gate_time = self.timers(TOPK_GATE_TIMER).elapsed(
              reset=False
          )
  
      return gate_output
  
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 2 for calling the TopKGate obj's underlying member method code with example usage: obj.forward(input, used_token, use_tutel) and returns: gate_output