class_docstr: SumGenerator deals with the sharding strategies of torch.sum op.
class_docstr_tok: ["SumGenerator","deals","with","the","sharding","strategies","of","torch.sum","op","."]
code: def collate_strategies(self) -> List[ShardingStrategy]:
      strategy_list = []
      for index, strategy in enumerate(
          self.predecessor_node.strategies_vector
      ):
          dim_partition_dict_mapping = {}
          communication_action_mapping = {}
          input_sharding_spec = strategy.output_sharding_specs[
              self.op_data["input"]
          ]
          dim_partition_dict_for_input = copy.deepcopy(
              input_sharding_spec.dim_partition_dict
          )
          sum_dims, sum_mapping_dict = self.op_data["sum_info"].data
  
          # TODO: a better way to handle the distributed sum is sum all the data on chip and then do all reduce
          # among all the shard groups
          recover_dims = []
          dim_partition_dict_for_output = {}
          for dim in dim_partition_dict_for_input:
              if dim in sum_dims:
                  recover_dims.append(dim)
              elif dim in sum_mapping_dict:
                  dim_partition_dict_for_output[
                      sum_mapping_dict[dim]
                  ] = dim_partition_dict_for_input[dim]
              else:
                  raise RuntimeError(
                      f"dim {dim} is not in sum_mapping_dict or sum_dims"
                  )
  
          for dim in recover_dims:
              dim_partition_dict_for_input.pop(dim)
  
          dim_partition_dict_mapping = {
              "input": dim_partition_dict_for_input,
              "output": dim_partition_dict_for_output,
          }
          sharding_spec_mapping = self.to_sharding_spec_mapping(
              dim_partition_dict_mapping
          )
          # add index into name to pass the duplicated check
          # we keep same strategies with different name for node merging, and it will not increase the searching space,
          # because in solver, this node will be merged into other nodes, and solver will not create a new variable for this node.
          name = f'{sharding_spec_mapping["input"].sharding_sequence} -> {sharding_spec_mapping["output"].sharding_sequence}_{index}'
  
          strategy = self.get_sharding_strategy(
              name=name,
              sharding_spec_mapping=sharding_spec_mapping,
              communication_action_mapping=communication_action_mapping,
          )
          strategy_list.append(strategy)
  
      return strategy_list
  
file_path: ColossalAI/colossalai/auto_parallel/tensor_shard/node_handler/strategy/sum_generator.py
imports: import copy
  import operator
  from functools import reduce
  from typing import List
  from colossalai.auto_parallel.tensor_shard.node_handler.strategy.strategy_generator import FollowingStrategyGenerator
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, ShardingStrategy, TrainCycleItem
label_desc: Use image node_id 4 for calling the SumGenerator obj's underlying member method code with example usage: obj.collate_strategies() and returns: strategy_list