class_docstr: Generate sharding strategies for the batched matrix multiplication.
  
  A batched matrix multiplication can be viewed as
  [b, i, k] x [b, k, j] -> [b, i, j]
  
  The bias term is considered to have a 2D logical shape.
  
  Note: This class will be used to generate strategies for torch.bmm
  and torch.addbmm. However, the result of torch.addbmm is not correct,
  some extra runtime apply actions are required to keep numerical correctness.
class_docstr_tok: ["Generate","sharding","strategies","for","the","batched","matrix","multiplication",".","A","batched","matrix","multiplication","can","be","viewed","as","[","b",",","i",",","k","]","x","[","b",",","k",",","j","]","-",">","[","b",",","i",",","j","]","The","bias","term","is","considered","to","have","a","2D","logical","shape",".","Note",":","This","class","will","be","used","to","generate","strategies","for","torch.bmm","and","torch.addbmm",".","However",",","the","result","of","torch.addbmm","is","not","correct",",","some","extra","runtime","apply","actions","are","required","to","keep","numerical","correctness","."]
code: def update_compute_cost(
      self, strategy: ShardingStrategy
  ) -> ShardingStrategy:
      fwd_compute_cost = self.op_data["input"].data.shape[-1] * reduce(
          operator.mul, self.op_data["output"].data.shape
      )
      bwd_compute_cost = fwd_compute_cost * 2
      compute_cost = TrainCycleItem(
          fwd=fwd_compute_cost,
          bwd=bwd_compute_cost,
          total=fwd_compute_cost + bwd_compute_cost,
      )
      strategy.compute_cost = compute_cost
  
imports: import operator
  from functools import reduce
  from typing import List
  from colossalai.auto_parallel.tensor_shard.options import SolverPerference
  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 4 for calling the BatchedMatMulStrategyGenerator obj's underlying member method code with example usage: obj.update_compute_cost(strategy) without return types