class_docstr: MOELayer module which implements MixtureOfExperts as described in Gshard_.
  ::
  
      gate = Top2Gate(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: gate network
      expert: expert network
      group: group to use for all-to-all communication
class_docstr_tok: ["MOELayer","module","which","implements","MixtureOfExperts","as","described","in","Gshard_",".",":",":","gate","=","Top2Gate","(","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",":","gate","network","expert",":","expert","network","group",":","group","to","use","for","all-to-all","communication"]
code: def __init__(
      self,
      gate: Module,
      experts: Union[Module, ModuleList],
      group: Optional[Any] = None,
  ) -> None:
      super().__init__()
      self.gate = gate
      if type(experts) == ModuleList:
          self.experts = cast(ModuleList, experts)
      else:
          self.experts = ModuleList([experts])
      self.group = group if group is not None else dist.group.WORLD
      for expert in self.experts:
          for p in experts.parameters():
              p.expert = True  # type: ignore
      self.world_size = dist.get_world_size(self.group)
      self.num_local_experts = len(self.experts)
  
file_path: fairscale/fairscale/nn/moe/moe_layer.py
imports: from typing import TYPE_CHECKING, Any, Optional, Tuple, Union, cast
  import torch
  from torch import Tensor
  import torch.distributed
  from torch.nn import Module, ModuleList
label_desc: Use image node_id 1 to create a new MOELayer object from inherited base classes: Base with example: obj = MOELayer(gate, experts, group)