docstr: Beam search for the best swap combination.
  Specifically, we swap two elements from two groups and calculate the score.
  The score is the difference between the origin group sum and the new group sum.
  The larger the score, the better the swap combination.
  
  Args:
      inputs (Tuple): (data, origin_score, swap_list)
      beam_width (int): beam width for beam search
      avg (float): average value of the data
      group_swap_factor (float): group loss for group swap loss
  
  Returns:
      List: results list
docstr_tok: ["Beam","search","for","the","best","swap","combination",".","Specifically",",","we","swap","two","elements","from","two","groups","and","calculate","the","score",".","The","score","is","the","difference","between","the","origin","group","sum","and","the","new","group","sum",".","The","larger","the","score",",","the","better","the","swap","combination",".","Args",":","inputs","(","Tuple",")",":","(","data",",","origin_score",",","swap_list",")","beam_width","(","int",")",":","beam","width","for","beam","search","avg","(","float",")",":","average","value","of","the","data","group_swap_factor","(","float",")",":","group","loss","for","group","swap","loss","Returns",":","List",":","results","list"]
code: def _beam_search(
      self,
      inputs: Tuple[List, float, List],
      beam_width: int,
      avg: float,
      group_swap_factor: float,
  ) -> List:
      """
      Beam search for the best swap combination.
      Specifically, we swap two elements from two groups and calculate the score.
      The score is the difference between the origin group sum and the new group sum.
      The larger the score, the better the swap combination.
  
      Args:
          inputs (Tuple): (data, origin_score, swap_list)
          beam_width (int): beam width for beam search
          avg (float): average value of the data
          group_swap_factor (float): group loss for group swap loss
  
      Returns:
          List: results list
      """
      data, origin_score, swap_list = inputs
      results = []
      group_num = len(data)
      group_size = len(data[0])
      origin_diff_list = [
          self._get_diff_from_avg(data, i, avg)
          for i in range(group_num)
      ]
  
      for group_num_i in range(group_num):
          for group_size_i in range(group_size):
              for group_num_j in range(group_num_i + 1, group_num):
                  for group_size_j in range(group_size):
                      new_data = deepcopy(data)
                      # calculate origin group sum
                      origin_diff = (
                          origin_diff_list[group_num_i]
                          + origin_diff_list[group_num_j]
                      )
                      # swap data
                      self._swap_data(
                          new_data,
                          group_num_i,
                          group_size_i,
                          group_num_j,
                          group_size_j,
                      )
                      # calculate new group sum
                      new_diff = self._get_diff_from_avg(
                          new_data, group_num_i, avg
                      ) + self._get_diff_from_avg(
                          new_data, group_num_j, avg
                      )
                      # caculate score
                      new_score = origin_diff - new_diff
                      if new_score > 0:
                          new_score = origin_score + new_score
                          # get swap loss
                          swap_loss = self._get_swap_loss(
                              group_swap_factor,
                              swap_list,
                              group_num_i,
                              group_size_i,
                              group_num_j,
                              group_size_j,
                          )
                          new_score = new_score - swap_loss
                          # update swap list
                          new_swap_list = swap_list + [
                              (
                                  group_num_i,
                                  group_size_i,
                                  group_num_j,
                                  group_size_j,
                              )
                          ]
                          results.append(
                              (new_data, new_score, new_swap_list)
                          )
      # sort results
      results.sort(key=lambda x: x[1], reverse=True)
      # select top k results
      results = results[:beam_width]
      return results
  
imports: from copy import deepcopy
  from typing import List, Optional, Tuple
  import torch
  import torch.distributed
  from torch import Tensor, nn
  from torch.distributed import ProcessGroup
  from colossalai.cluster import ProcessGroupMesh
  from colossalai.moe.experts import MLPExperts
  from colossalai.moe.manager import MOE_MANAGER
  from colossalai.zero.low_level import LowLevelZeroOptimizer
label_desc: Use image node_id 9 for calling the LoadBalancer obj's underlying member method code with example usage: obj._beam_search(inputs, beam_width, avg, group_swap_factor) and returns: results