class_docstr: MatMulStrategyGenerator is a generic class to cover all matrix multiplication cases.
  The operation data is defined as `output = input x other + bias`.
class_docstr_tok: ["MatMulStrategyGenerator","is","a","generic","class","to","cover","all","matrix","multiplication","cases",".","The","operation","data","is","defined","as","`","output","=","input","x","other","+","bias","`","."]
code: def update_memory_cost(
      self, strategy: ShardingStrategy
  ) -> ShardingStrategy:
      size_mapping = {
          "input": self._compute_size_in_bytes(strategy, "input"),
          "other": self._compute_size_in_bytes(strategy, "other"),
          "output": self._compute_size_in_bytes(strategy, "output"),
      }
  
      if self.has_bias:
          bias_size = self._compute_size_in_bytes(strategy, "bias")
          size_mapping["bias"] = bias_size
  
      # compute fwd cost incurred
      # fwd_cost = input + other + bias + output
      fwd_activation_cost = sum(
          [v for k, v in size_mapping.items() if not self.is_param(k)]
      )
      fwd_parameter_cost = sum(
          [v for k, v in size_mapping.items() if self.is_param(k)]
      )
      fwd_mem_cost = MemoryCost(
          activation=fwd_activation_cost, parameter=fwd_parameter_cost
      )
  
      # compute bwd cost incurred
      # bwd_cost = input_grad + bias_grad
      bwd_activation_cost = sum(
          [
              v
              for k, v in size_mapping.items()
              if k in ["input", "other", "bias"]
          ]
      )
      bwd_mem_cost = MemoryCost(
          activation=bwd_activation_cost, parameter=0
      )
  
      # compute total cost
      total_mem_cost = MemoryCost(
          activation=fwd_activation_cost + bwd_activation_cost,
          parameter=fwd_parameter_cost + 0,
      )
      memory_cost = TrainCycleItem(
          fwd=fwd_mem_cost, bwd=bwd_mem_cost, total=total_mem_cost
      )
      strategy.memory_cost = memory_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 1 for calling the MatMulStrategyGenerator obj's underlying member method code with example usage: obj.update_memory_cost(strategy) without return types