class_docstr: ConvStrategyGenerator is a generic class to generate strategies.
  The operation data is defined as `output = input x other + bias`.
class_docstr_tok: ["ConvStrategyGenerator","is","a","generic","class","to","generate","strategies",".","The","operation","data","is","defined","as","`","output","=","input","x","other","+","bias","`","."]
docstr: Compute the computation cost per device with this specific strategy.
  
  Note: compute_cost need to be divided by TFLOPS, now it just shows the computation size.
docstr_tok: ["Compute","the","computation","cost","per","device","with","this","specific","strategy",".","Note",":","compute_cost","need","to","be","divided","by","TFLOPS",",","now","it","just","shows","the","computation","size","."]
code: def update_compute_cost(self, strategy: ShardingStrategy):
      """
      Compute the computation cost per device with this specific strategy.
  
      Note: compute_cost need to be divided by TFLOPS, now it just shows the computation size.
      """
      # TODO: compute_cost need to be divided by TFLOPS, now it just shows the computation size.
      # 1D: (L) * N * Cout * Cin * kernel
      # 2D: (H * W) * N * Cout * Cin * kernel
      # 3D: (H * W  * D) * N * Cout * Cin * kernel
      sharded_input_shape = strategy.sharding_specs[
          self.op_data["input"]
      ].get_sharded_shape_per_device()
      sharded_other_shape = strategy.sharding_specs[
          self.op_data["other"]
      ].get_sharded_shape_per_device()
      sharded_output_shape = strategy.sharding_specs[
          self.op_data["output"]
      ].get_sharded_shape_per_device()
      if self.has_bias:
          # bias add is an element wise operation, so the cost is equal to product of output shape.
          bias_compute_cost = reduce(operator.mul, sharded_output_shape)
  
      output_size = sharded_output_shape[2:]
      output_size_product = reduce(operator.mul, output_size)
      input_size = sharded_input_shape[2:]
      input_size_product = reduce(operator.mul, input_size, 1)
      kernel_size = sharded_other_shape[2:]
      kernel_size_product = reduce(operator.mul, kernel_size, 1)
      batch_size = sharded_input_shape[0]
      channel_in = sharded_input_shape[1]
      channel_out = sharded_other_shape[1]
  
      forward_compute_cost = (
          output_size_product
          * batch_size
          * channel_in
          * channel_out
          * kernel_size_product
      )
  
      backward_activation_cost = (
          input_size_product
          * batch_size
          * channel_in
          * channel_out
          * kernel_size_product
      )
      backward_weight_cost = (
          output_size_product
          * batch_size
          * channel_in
          * channel_out
          * kernel_size_product
      )
      backward_compute_cost = (
          backward_weight_cost + backward_activation_cost
      )
      if self.has_bias:
          forward_compute_cost += bias_compute_cost
          backward_compute_cost += bias_compute_cost
      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
  
imports: import copy
  import operator
  from functools import reduce
  from typing import List
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import CommType, MemoryCost, ShardingStrategy, TrainCycleItem
  from colossalai.auto_parallel.tensor_shard.utils import ignore_sharding_exception
  from colossalai.tensor.shape_consistency import CollectiveCommPattern
  from .strategy_generator import StrategyGenerator
label_desc: Use image node_id 2 for calling the ConvStrategyGenerator obj's underlying member method code with example usage: obj.update_compute_cost(strategy) without return types