class_docstr: A class for users to create MoE modules in their models.
  
  Args:
      dim_model (int): Hidden dimension of training model
      num_experts (int): The number experts
      top_k (int, optional): The number of experts for dispatchment of each token
      capacity_factor_train (float, optional): Capacity factor in routing during training
      capacity_factor_eval (float, optional): Capacity factor in routing during evaluation
      min_capacity (int, optional): The minimum number of the capacity of each expert
      noisy_policy (str, optional): The policy of noisy function. Now we have 'Jitter' and 'Gaussian'.
          'Jitter' can be found in `Switch Transformer paper`_.
          'Gaussian' can be found in `ViT-MoE paper`_.
      drop_tks (bool, optional): Whether drops tokens in evaluation
      use_residual (bool, optional): Makes this MoE layer a Residual MoE.
          More information can be found in `Microsoft paper`_.
      residual_instance (nn.Module, optional): The instance of residual module in Residual MoE
      expert_instance (MoeExperts, optional): The instance of experts module in MoeLayer
      expert_cls (Type[nn.Module], optional): The class of each expert when no instance is given
      expert_args (optional): The args of expert when no instance is given
  
  .. _Switch Transformer paper:
      https://arxiv.org/abs/2101.03961
  .. _ViT-MoE paper:
      https://arxiv.org/abs/2106.05974
  .. _Microsoft paper:
      https://arxiv.org/abs/2201.05596
docstr: without overlap:
             |    C    |
  |     A    |         |    R    |
  
  with overlap:
        |    C1   ||    C2   ||    C3   ||    C4   |
  | A1 || A2 |     | R1 | A3 || R2 | A4 || R3 |     | R4 |
  
  where C is computation, A is all gather, R is reduce scatter.
  
  Args:
      dispatch_data (torch.Tensor): (num_experts, capacity, hidden_size)
  
  Returns:
      torch.Tensor: (num_experts, capacity, hidden_size)
docstr_tok: ["without","overlap",":","|","C","|","|","A","|","|","R","|","with","overlap",":","|","C1","||","C2","||","C3","||","C4","|","|","A1","||","A2","|","|","R1","|","A3","||","R2","|","A4","||","R3","|","|","R4","|","where","C","is","computation",",","A","is","all","gather",",","R","is","reduce","scatter",".","Args",":","dispatch_data","(","torch.Tensor",")",":","(","num_experts",",","capacity",",","hidden_size",")","Returns",":","torch.Tensor",":","(","num_experts",",","capacity",",","hidden_size",")"]
code: def _tp_process(
      self,
      dispatch_data: torch.Tensor,
      used_capacity: torch.Tensor,
      overlap: bool = False,
  ) -> torch.Tensor:
      """
      without overlap:
                 |    C    |
      |     A    |         |    R    |
  
      with overlap:
            |    C1   ||    C2   ||    C3   ||    C4   |
      | A1 || A2 |     | R1 | A3 || R2 | A4 || R3 |     | R4 |
  
      where C is computation, A is all gather, R is reduce scatter.
  
      Args:
          dispatch_data (torch.Tensor): (num_experts, capacity, hidden_size)
  
      Returns:
          torch.Tensor: (num_experts, capacity, hidden_size)
      """
      if not overlap or dist.get_world_size(self.ep_group) == 1:
          expert_in = AllGather.apply(
              dispatch_data, self.ep_group, False
          )[0]
          expert_out = self.experts(expert_in)
          expert_out = ReduceScatter.apply(
              expert_out, self.ep_group, False
          )[0]
          return expert_out
      else:
  
          @dataclasses.dataclass
          class Capsule:
              data: torch.Tensor
              handle: Any
              indices: Tuple
  
          NUM_CHUNK = 4
          NUM_STAGES = 4
  
          assert (
              dispatch_data.shape[0] % NUM_CHUNK == 0
          ), "arbitrary chunk num is not supported yet, please use chunk num that can divide num_experts"
          chunk_size = dispatch_data.shape[0] // NUM_CHUNK
          chunk_data = torch.split(dispatch_data, chunk_size, dim=0)
          output = torch.empty_like(dispatch_data)
  
          def get_chunk_slice(
              idx: int, chunk_size: int
          ) -> Tuple[slice]:
              return (slice(idx * chunk_size, (idx + 1) * chunk_size),)
  
          _expert_in, expert_in, _expert_out, expert_out = (
              None,
              None,
              None,
              None,
          )
  
          for i in range(NUM_CHUNK + NUM_STAGES - 1):
              if expert_out is not None:
                  expert_out.handle.wait()
                  output[expert_out.indices] = expert_out.data
                  expert_out = None
  
              # reduce scatter last output
              if _expert_out is not None:
                  expert_out = Capsule(
                      *ReduceScatter.apply(
                          _expert_out.data, self.ep_group, True
                      ),
                      indices=_expert_out.indices,
                  )
                  _expert_out = None
  
              # all gather next input
              if 0 <= i < NUM_CHUNK:
                  _expert_in = Capsule(
                      *AllGather.apply(
                          chunk_data[i].contiguous(),
                          self.ep_group,
                          True,
                      ),
                      indices=get_chunk_slice(i, chunk_size),
                  )
  
              # compute
              if expert_in is not None:
                  expert_in.handle.wait()
                  _expert_out = Capsule(
                      self.experts(expert_in.data, expert_in.indices),
                      handle=None,
                      indices=expert_in.indices,
                  )
                  expert_in = None
  
              if _expert_in is not None:
                  expert_in = _expert_in
                  _expert_in = None
  
          return output
  
imports: import dataclasses
  import math
  from typing import Any, Optional, Tuple
  import torch
  import torch.distributed
  import torch.nn
  import torch.nn.functional
  from colossalai.moe._operation import AllGather, AllToAll, HierarchicalAllToAll, MoeCombine, MoeDispatch, ReduceScatter
  from colossalai.moe.experts import MLPExperts
  from colossalai.moe.load_balance import LoadBalancer
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.moe.routers import MoeRouter, get_router_cls
  from colossalai.moe.utils import create_ep_hierarchical_group, get_noise_generator
  from colossalai.tensor.moe_tensor.api import get_dp_group, get_ep_group, get_ep_group_ranks, get_ep_size
label_desc: Use image node_id 6 for calling the SparseMLP obj's underlying member method code with example usage: obj._tp_process(dispatch_data, used_capacity, overlap) and returns: expert_out, output, slice