class_docstr: The NodeHandler is an abstract class used to generate every possible strategies for an operator node.
  
  Args:
      node (Node): the input node in node argument list.
      device_mesh (DeviceMesh): A logical view of a physical mesh.
      strategies_vector (StrategiesVector): all the strategies generated in this handler will be recorded into the strategies_vector.
class_docstr_tok: ["The","NodeHandler","is","an","abstract","class","used","to","generate","every","possible","strategies","for","an","operator","node",".","Args",":","node","(","Node",")",":","the","input","node","in","node","argument","list",".","device_mesh","(","DeviceMesh",")",":","A","logical","view","of","a","physical","mesh",".","strategies_vector","(","StrategiesVector",")",":","all","the","strategies","generated","in","this","handler","will","be","recorded","into","the","strategies_vector","."]
docstr: Compute the resharding costs and save the costs in the ShardingStrategy object.
docstr_tok: ["Compute","the","resharding","costs","and","save","the","costs","in","the","ShardingStrategy","object","."]
code: def update_resharding_cost(self, strategy: ShardingStrategy) -> None:
      """
      Compute the resharding costs and save the costs in the ShardingStrategy object.
      """
      # TODO: test this function when other handlers are ready
      resharding_costs = {}
      shape_consistency_manager = ShapeConsistencyManager()
  
      for node in self.predecessor_node:
          node_name = str(node)
          # get the current sharding spec generated by this node handler
  
          # we will not compute the resharding costs for the node not counted in the strategy.
          # And the node with tuple or list output need to be handled below.
          node_in_strategy = [
              op_data.name for op_data in strategy.sharding_specs.keys()
          ]
          if str(node) not in node_in_strategy:
              continue
  
          op_data = strategy.get_op_data_by_name(node_name)
          current_sharding_spec = strategy.sharding_specs[op_data]
          # get the sharding specs for this node generated
          # in its own node handler
          assert hasattr(
              node, "strategies_vector"
          ), f"The predecessor node {node_name} has no strategy vector to compute the resharding cost."
          prev_strategy_vector = node.strategies_vector
          prev_sharding_specs = [
              prev_strategy.get_sharding_spec_by_name(node_name)
              for prev_strategy in prev_strategy_vector
          ]
  
          # create data structure to store costs
          if node not in resharding_costs:
              resharding_costs[node] = []
  
          def _compute_resharding_cost(
              prev_sharding_spec: Union[
                  ShardingSpec, List[ShardingSpec]
              ],
              current_sharding_spec: Union[
                  ShardingSpec, List[ShardingSpec]
              ],
              data: Union[
                  torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor]
              ],
          ) -> TrainCycleItem:
              """
              This is a helper function to compute the resharding cost for a specific strategy of a node.
              """
              if prev_sharding_spec is None:
                  return TrainCycleItem(fwd=0, bwd=0, total=0)
              elif isinstance(prev_sharding_spec, ShardingSpec):
                  if isinstance(data, torch.Tensor):
                      dtype = data.dtype
                      size_per_elem_bytes = torch.tensor(
                          [], dtype=dtype
                      ).element_size()
                      (
                          _,
                          _,
                          consistency_cost,
                      ) = shape_consistency_manager.shape_consistency(
                          prev_sharding_spec, current_sharding_spec
                      )
  
                      resharding_cost = TrainCycleItem(
                          fwd=consistency_cost["forward"]
                          * size_per_elem_bytes,
                          bwd=consistency_cost["backward"]
                          * size_per_elem_bytes,
                          total=consistency_cost["total"]
                          * size_per_elem_bytes,
                      )
                      return resharding_cost
                  else:
                      # This raise is used to check if we have missed any type of data.
                      # It could be merged into Parameter branch, which means we won't handle
                      # non-tensor arguments.
                      raise ValueError(
                          f"Unsupported data type {type(data)}"
                      )
              else:
                  assert isinstance(
                      prev_sharding_spec, (tuple, list)
                  ), f"prev_sharding_spec should be in type of ShardingSpec, List[ShardingSpec], \
                              or Tuple[ShardingSpec], but got {type(prev_sharding_spec)}"
  
                  fwd_cost = 0
                  bwd_cost = 0
                  total_cost = 0
                  for index, (
                      prev_sharding_spec_item,
                      current_sharding_spec_item,
                  ) in enumerate(
                      zip(prev_sharding_spec, current_sharding_spec)
                  ):
                      item_cost = _compute_resharding_cost(
                          prev_sharding_spec_item,
                          current_sharding_spec_item,
                          data[index],
                      )
                      fwd_cost += item_cost.fwd
                      bwd_cost += item_cost.bwd
                      total_cost += item_cost.total
                  resharding_cost = TrainCycleItem(
                      fwd=fwd_cost, bwd=bwd_cost, total=total_cost
                  )
                  return resharding_cost
  
          # for each sharding spec generated by the predecessor's node handler
          # compute the resharding cost to switch to the sharding spec generated
          # by the current node handler
          for prev_sharding_spec in prev_sharding_specs:
              resharding_cost = _compute_resharding_cost(
                  prev_sharding_spec,
                  current_sharding_spec,
                  op_data.data,
              )
              resharding_costs[node].append(resharding_cost)
      strategy.resharding_costs = resharding_costs
      return strategy
  
imports: from abc import ABC, abstractmethod
  from typing import Dict, List, Tuple, Union
  import torch
  from torch.fx.node import Node
  from colossalai.auto_parallel.meta_profiler.shard_metainfo import ShardMetaInfo, meta_register
  from colossalai.auto_parallel.tensor_shard.options import ShardOption, SolverPerference
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import OperationData, ShardingSpec, ShardingStrategy, StrategiesVector, TrainCycleItem
  from colossalai.auto_parallel.tensor_shard.utils import check_sharding_spec_validity
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.logging import get_dist_logger
  from colossalai.tensor.shape_consistency import ShapeConsistencyManager
  from .strategy import StrategyGenerator
label_desc: Use image node_id 2 for calling the NodeHandler obj's underlying member method code with example usage: obj.update_resharding_cost(strategy) and returns: strategy, TrainCycleItem, resharding_cost, resharding_cost