class_docstr: UnaryElementwiseGenerator which deals with the sharding strategies of UnaryElementwiseOp.
docstr: Compute the memory cost per device with this specific strategy.
code: def update_memory_cost(self, strategy: ShardingStrategy):
      """
      Compute the memory cost per device with this specific strategy.
      """
      forward_size_mapping = {
          "input": self._compute_size_in_bytes(strategy, "input"),
          "output": self._compute_size_in_bytes(strategy, "output"),
      }
  
      backward_size_mapping = copy.deepcopy(forward_size_mapping)
      backward_size_mapping.pop("output")
      # compute fwd cost incurred
      # fwd_cost = input + output
      fwd_activation_cost = sum(
          [
              v
              for k, v in forward_size_mapping.items()
              if not self.is_param(k)
          ]
      )
      fwd_parameter_cost = sum(
          [
              v
              for k, v in forward_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
      bwd_activation_cost = sum(
          [
              v
              for k, v in backward_size_mapping.items()
              if not self.is_param(k)
          ]
      )
      bwd_parameter_cost = sum(
          [
              v
              for k, v in backward_size_mapping.items()
              if self.is_param(k)
          ]
      )
      bwd_mem_cost = MemoryCost(
          activation=bwd_activation_cost, parameter=bwd_parameter_cost
      )
  
      # compute total cost
      total_mem_cost = MemoryCost(
          activation=fwd_activation_cost + bwd_activation_cost,
          parameter=fwd_parameter_cost + bwd_parameter_cost,
      )
      memory_cost = TrainCycleItem(
          fwd=fwd_mem_cost, bwd=bwd_mem_cost, total=total_mem_cost
      )
      strategy.memory_cost = memory_cost
  
file_path: ColossalAI/colossalai/auto_parallel/tensor_shard/node_handler/strategy/unary_elementwise_generator.py
imports: import copy
  from typing import List
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import MemoryCost, ShardingStrategy, TrainCycleItem
  from .strategy_generator import FollowingStrategyGenerator
label_desc: Use image node_id 3 for calling the UnaryElementwiseGenerator obj's underlying member method code with example usage: obj.update_memory_cost(strategy) without return types