class_docstr: MOELayer module which implements MixtureOfExperts as described in Gshard_.
  ::
  
      gate = TopKGate(model_dim, num_experts)
      moe = MOELayer(gate, expert)
      output = moe(input)
      l_aux = moe.l_aux
  
  .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf
  
  Args:
      gate (torch.nn.Module):
          gate network
      expert (torch.nn.Module):
          expert network
class_docstr_tok: ["MOELayer","module","which","implements","MixtureOfExperts","as","described","in","Gshard_",".",":",":","gate","=","TopKGate","(","model_dim",",","num_experts",")","moe","=","MOELayer","(","gate",",","expert",")","output","=","moe","(","input",")","l_aux","=","moe.l_aux","..","Gshard_",":","https",":","\/\/arxiv.org\/pdf\/2006.16668.pdf","Args",":","gate","(","torch.nn.Module",")",":","gate","network","expert","(","torch.nn.Module",")",":","expert","network"]
code: def forward(self, *input: Tensor, **kwargs: Any) -> Tensor:
      if self.wall_clock_breakdown:
          self.timers(MOE_TIMER).start()
  
      # Implement Algorithm 2 from GShard paper.
      d_model = input[0].shape[-1]
  
      # Initial implementation -> Reshape into S tokens by dropping sequence dimension.
      # Reshape into G groups so that each group can distribute tokens equally
      # group_size = kwargs['group_size'] if 'group_size' in kwargs.keys() else 1
      reshaped_input = input[0].reshape(-1, d_model)
  
      if self.use_tutel:
          (
              self.l_aux,
              C,
              E,
              indices_,
              locations_,
              gates_,
              self.exp_counts,
          ) = self.gate(reshaped_input, input[1], True)
          S, M = reshaped_input.size(0), reshaped_input.size(1)
  
          if not hasattr(self, "_tutel_dispatcher"):
              self._tutel_dispatcher = tutel_moe.fast_dispatcher(
                  E, C, M, dispatch_dtype=reshaped_input.dtype
              )
          self._tutel_dispatcher.update(
              indices_, locations_, gates_, capacity=C
          )
          dispatched_input = self._tutel_dispatcher.encode(
              reshaped_input
          )
      else:
          (
              self.l_aux,
              combine_weights,
              dispatch_mask,
              self.exp_counts,
          ) = self.gate(reshaped_input, input[1])
          dispatched_input = einsum(
              "sec,sm->ecm",
              dispatch_mask.type_as(input[0]),
              reshaped_input,
          )
  
      if self.wall_clock_breakdown:
          self.timers(FIRST_ALLTOALL_TIMER).start()
  
      if groups._get_expert_model_parallel_world_size() == 1:
          # If the non-expert is tensor-parallel, it will create
          # duplicate tokens on the tensor-parallel ranks.
          # Since our experts are not tensor-parallel, these duplicates
          # need to be dropped to ensure correctness.
          # this also doubles up as a communication optimization as we are
          # reducing the all-to-all communication volume.
          dispatched_input = drop_tokens(dispatched_input, dim=1)
  
      dispatched_input = _AllToAll.apply(
          self.ep_group, dispatched_input
      )
  
      if self.wall_clock_breakdown:
          self.timers(FIRST_ALLTOALL_TIMER).stop()
          self.time_falltoall = self.timers(
              FIRST_ALLTOALL_TIMER
          ).elapsed(reset=False)
  
      # Re-shape after all-to-all: ecm -> gecm
      dispatched_input = dispatched_input.reshape(
          self.ep_size, self.num_local_experts, -1, d_model
      )
  
      expert_output = self.experts(dispatched_input)
  
      if self.wall_clock_breakdown:
          self.timers(SECOND_ALLTOALL_TIMER).start()
  
      expert_output = _AllToAll.apply(self.ep_group, expert_output)
  
      if self.wall_clock_breakdown:
          self.timers(SECOND_ALLTOALL_TIMER).stop()
          self.time_salltoall = self.timers(
              SECOND_ALLTOALL_TIMER
          ).elapsed(reset=False)
  
      # Re-shape back: gecm -> ecm
      expert_output = expert_output.reshape(
          self.ep_size * self.num_local_experts, -1, d_model
      )
  
      if groups._get_expert_model_parallel_world_size() == 1:
          # the dropped duplicate tokens need to be gathered on each
          # tensor parallel rank again for the tensor-parallel
          # non-expert of the next layer.
          expert_output = gather_tokens(expert_output, dim=1)
  
      if self.use_tutel:
          combined_output = self._tutel_dispatcher.decode(
              expert_output.view(E * C, M)
          )
      else:
          combined_output = einsum(
              "sec,ecm->sm",
              combine_weights.type_as(input[0]),
              expert_output,
          )
  
      a = combined_output.reshape(input[0].shape)
  
      if self.wall_clock_breakdown:
          self.timers(MOE_TIMER).stop()
          self.time_moe = self.timers(MOE_TIMER).elapsed(reset=False)
  
      return a
  
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 3 for calling the MOELayer obj's underlying member method code with example usage: obj.forward() and returns: a