class_docstr: SumGenerator deals with the sharding strategies of torch.sum op.
code: def update_compute_cost(self, strategy: ShardingStrategy):
      sharded_input_shape = strategy.sharding_specs[
          self.op_data["input"]
      ].get_sharded_shape_per_device()
      sharded_output_shape = strategy.sharding_specs[
          self.op_data["output"]
      ].get_sharded_shape_per_device()
      input_size_product = reduce(operator.mul, sharded_input_shape)
      output_size_product = reduce(operator.mul, sharded_output_shape)
  
      compute_cost = TrainCycleItem(
          fwd=input_size_product,
          bwd=output_size_product,
          total=input_size_product + output_size_product,
      )
  
      strategy.compute_cost = compute_cost
  
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 2 for calling the SumGenerator obj's underlying member method code with example usage: obj.update_compute_cost(strategy) without return types