class_docstr: A StrategyGenerator which deals with the sharding strategies of batch normalization.
  
  To keep the math consistency, there are two way to do BatchNorm if the input
  shards on batch dimension:
  1. We gather the input partitions through batch dimension, then do the normal BatchNorm.
  2. We do the SyncBatchNorm on the each input partition separately, the SyncBN op will help
     us to keep the computing correctness.
  In this generator, both methods will be considered.
class_docstr_tok: ["A","StrategyGenerator","which","deals","with","the","sharding","strategies","of","batch","normalization",".","To","keep","the","math","consistency",",","there","are","two","way","to","do","BatchNorm","if","the","input","shards","on","batch","dimension",":","1",".","We","gather","the","input","partitions","through","batch","dimension",",","then","do","the","normal","BatchNorm",".","2",".","We","do","the","SyncBatchNorm","on","the","each","input","partition","separately",",","the","SyncBN","op","will","help","us","to","keep","the","computing","correctness",".","In","this","generator",",","both","methods","will","be","considered","."]
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: a constant coefficient need to be added.
      # 1D: (L) * N * Cin
      # 2D: (H * W) * N  * Cin
      # 3D: (H * W  * D) * N  * Cin
      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()
      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)
      input_product = reduce(operator.mul, sharded_input_shape, 1)
      forward_compute_cost = input_product
      backward_activation_compute_cost = input_product
      backward_weight_compute_cost = input_product
      backward_compute_cost = (
          backward_weight_compute_cost
          + backward_activation_compute_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
  
file_path: colossalai/colossalai/auto_parallel/tensor_shard/node_handler/strategy/batch_norm_generator.py
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 BatchNormStrategyGenerator obj's underlying member method code with example usage: obj.update_compute_cost(strategy) without return types