class_docstr: Tutel optimized MOELayer
      
class_docstr_tok: ["Tutel","optimized","MOELayer"]
code: def __init__(
      self,
      gate_type,
      model_dim: int,
      experts=None,
      scan_expert_func=None,
      result_func=None,
      group=None,
      seeds=None,
      a2a_ffn_overlap_degree=1,
      is_postscore=True,
      batch_prioritized_routing=False,
      normalize_gate=True,
      is_gshard_loss=True,
      parallel_type="adaptive:1",
      use_2dh=False,
      **kwargs,
  ):
      super().__init__()
      assert model_dim % 2 == 0, (
          "Model_dim (%s) must be even value, while this Model_dim mod 2 > 0."
          % model_dim
      )
      group = group or dist.group.WORLD
  
      if "pad_samples" in kwargs:
          logging.warning(
              f"`pad_samples` option in Tutel Moe-layer has been deprecated, as Tutel always assumes `pad_samples=False` for better efficiency."
          )
          kwargs.pop("pad_samples")
      for k in kwargs:
          raise Exception(
              "Unrecognized argument provided to Tutel Moe-layer: %s"
              % k
          )
  
      self.group = group
      self.result_func = result_func
      self.skip_moe = int(os.environ.get("SKIP_MOE", "0")) != 0
  
      self.num_local_experts = experts.pop("count_per_node", 1)
      self.register_buffer(
          "_num_global_experts",
          torch.tensor(
              MOELayer.global_expert_count(
                  self.num_local_experts, self.group
              )
          ),
      )
  
      self.world_size = C.get_world_size(self.group)
      if self.num_global_experts < self.world_size:
          self.sharded_count = (
              self.world_size // self.num_global_experts
          )
          self.num_local_experts = 1
      else:
          self.sharded_count = 1
  
      (
          self.auto_parallel,
          self.adaptive_degree,
          self.use_model_parallel,
      ) = (False, self.sharded_count, True)
      self.valid_rs = [0] + [
          i
          for i in range(1, self.sharded_count + 1)
          if self.sharded_count % i == 0
      ]
  
      if parallel_type.startswith("adaptive:"):
          self.adaptive_degree = int(
              parallel_type[parallel_type.index(":") + 1 :]
          )
          self.adaptive_degree = min(
              max(self.adaptive_degree, 0), self.sharded_count
          )
          if self.adaptive_degree not in self.valid_rs:
              raise Exception(
                  "Unexpected value of adaptive_degree: %d, expecting a candidate within %s."
                  % (self.adaptive_degree, self.valid_rs)
              )
      elif self.sharded_count == 1:
          pass
      elif parallel_type in ("data", "model"):
          self.adaptive_degree = (
              1 if parallel_type == "data" else self.sharded_count
          )
      elif parallel_type == "auto":
          self.adaptive_degree = 1
      else:
          raise Exception(
              "Unrecognized parallel type specified: %s" % parallel_type
          )
  
      self.model_dim = model_dim
  
      self.is_postscore = is_postscore
      self.batch_prioritized_routing = batch_prioritized_routing
      if int(os.environ.get("BATCH_PRIO", 0)) != 0:
          self.batch_prioritized_routing = True
      self.normalize_gate = normalize_gate
      self.is_gshard_loss = is_gshard_loss
  
      self.a2a_ffn_overlap_degree = a2a_ffn_overlap_degree
      self.use_2dh = use_2dh
  
      if seeds is not None and seeds[1] is not None:
          torch.manual_seed(seeds[1])
  
      experts_type = experts.pop("type")
      if experts_type == "custom":
          self.experts = cast(ModuleList, experts["module"])
      else:
          assert re.match(
              r"[a-zA-Z0-9\_]+", experts_type
          ), "Expert type must only include digits, letters and underline characters."
          try:
              fused_experts = importlib.import_module(
                  f"...experts.{experts_type}", __name__
              )
          except ModuleNotFoundError:
              raise Exception(
                  "Builtin expert type is not recognized: %s"
                  % experts_type
              )
  
          if experts_type == "ffn":
              assert (
                  "fused_custom_fn" not in experts
              ), "`fused_custom_fn` option for Tutel Moe-layer has been deprecated, please follows helloworld_from_scratch.py for custom construction instead."
              assert (
                  "implicit_dropout_p" not in experts
              ), "`implicit_dropout_p` option for Tutel Moe-layer has been deprecated, please use torch.nn.Dropout(p=implicit_dropout_p) on custom activation_fn (for fc1_dropout) and after Tutel Moe-layer (for fc2_dropout) instead."
  
          experts["model_dim"] = self.model_dim
          experts["local_experts"] = self.num_local_experts
          experts["sharded_count"] = self.sharded_count
          self.experts = fused_experts.ExpertModule(**experts)
  
      if scan_expert_func is not None:
          for n, p in self.experts.named_parameters():
              scan_expert_func(n, p)
      for n, p in self.experts.named_parameters():
          setattr(p, "_tutel_expert", True)
  
      if isinstance(gate_type, str):
          assert re.match(r"^Top[0-9]+Gate$", gate_type), (
              "Unrecognized gate_type: %s" % gate_type
          )
          top_k = int(gate_type[3:-4])
          logging.warning(
              f"gate_type value `{gate_type}` in Tutel Moe-layer has been deprecated, please use gate_type = {{'type': 'top', 'k': {top_k}}} instead."
          )
          gate_type = {"type": "top", "k": top_k}
  
      if not isinstance(gate_type, list):
          gate_type = [gate_type]
  
      self.gates = []
      for gi, single_gate_type in enumerate(gate_type):
          gate_type = single_gate_type["type"]
          single_gate_type.pop("type")
          assert re.match(
              r"[a-zA-Z0-9\_]+", gate_type
          ), "Gate type must only include digits, letters and underline characters."
  
          if seeds is not None and seeds[0] is not None:
              torch.manual_seed(seeds[0] + gi)
          try:
              single_gate = importlib.import_module(
                  f"...gates.{gate_type}", __name__
              )
          except ModuleNotFoundError:
              raise Exception("Unrecognized gate_type: %s" % gate_type)
  
          gate_module = single_gate.Gate(
              model_dim=self.model_dim,
              num_global_experts=self.num_global_experts,
              **single_gate_type,
          )
          if not hasattr(gate_module, "gate_noise"):
              gate_module.gate_noise = single_gate_type.get(
                  "gate_noise", 0.0
              )
          if not hasattr(gate_module, "capacity_factor"):
              gate_module.capacity_factor = single_gate_type.get(
                  "capacity_factor",
                  float(os.environ.get("CAP_FACTOR", 1.0)),
              )
  
          self.gates += [gate_module]
  
      self.gates = ModuleList(self.gates)
  
      if seeds is not None and len(seeds) > 2 and seeds[2] is not None:
          torch.manual_seed(seeds[2])
  
imports: from typing import TYPE_CHECKING, Any, Optional, Tuple, Union, cast
  import copy
  import os
  import re
  import time
  import logging
  import collections
  import importlib
  import torch
  from torch import Tensor
  import torch.distributed
  from torch.nn import ModuleList
  import torch.nn.functional
  from ..impls import communicate
  from ..impls.fast_dispatch import fast_encode, fast_decode, extract_critical, get_dispatch_count
  from ..impls.overlap import a2a_ffn_overlap_forward
  from .None import losses
label_desc: Use image node_id 5 to create a new MOELayer object from inherited base classes: torch.nn with example: obj = MOELayer(gate_type, model_dim, experts, scan_expert_func, result_func, group, seeds, a2a_ffn_overlap_degree, is_postscore, batch_prioritized_routing, normalize_gate, is_gshard_loss, parallel_type, use_2dh)