class_docstr: SoftmaxGenerator is used to generate strategies for torch.nn.Softmax or F.softmax.
class_docstr_tok: ["SoftmaxGenerator","is","used","to","generate","strategies","for","torch.nn.Softmax","or","F.softmax","."]
docstr: Compute the computation cost per device with this specific strategy.
docstr_tok: ["Compute","the","computation","cost","per","device","with","this","specific","strategy","."]
code: def update_compute_cost(self, strategy: ShardingStrategy):
      """
      Compute the computation cost per device with this specific strategy.
      """
      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)
  
      forward_compute_cost = output_size_product * 2
      backward_compute_cost = input_size_product
      total_compute_cost = forward_compute_cost + backward_compute_cost
      compute_cost = TrainCycleItem(
          fwd=forward_compute_cost,
          bwd=backward_compute_cost,
          total=total_compute_cost,
      )
      strategy.compute_cost = compute_cost
  
file_path: ColossalAI/colossalai/auto_parallel/tensor_shard/node_handler/strategy/softmax_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 SoftmaxGenerator obj's underlying member method code with example usage: obj.update_compute_cost(strategy) without return types