code: def size_value_converting_pass(
      gm: torch.fx.GraphModule, device_mesh: DeviceMesh
  ):
      """
      In the auto parallel system, tensors may get shard on different devices, so the size of tensors
      need to be converted to the size of original tensor and managed by the users, such as torch.view,
      torch.reshape, etc. These nodes have enough information like input sharding_spec and
      output sharding_spec to decide how to convert the size value.
      """
      mod_graph = gm.graph
      nodes = tuple(mod_graph.nodes)
      node_pairs = {}
  
      # DeviceMesh information instructs the scaling of the size value
      device_mesh_info = {}
      for dim, dim_size in enumerate(device_mesh.shape):
          device_mesh_info[dim] = dim_size
  
      def _extract_target_dim(node):
          """
          A helper function to extract the target dimension from size node.
          There are two usages of torch.Tensor.size:
          1. tensor.size()
          2. tensor.size(dim)
  
          If a target_dim is assigned, then the output will be in type of int, instead of torch.Size.
          Otherwise, the output will be in type of torch.Size and this function will return None.
          """
          target_dim = None
          if len(node.args) > 1:
              target_dim = node.args[1]
              if target_dim < 0:
                  target_dim += node.args[0]._meta_data.dim()
          return target_dim
  
      def _post_processing(node, size_processing_node):
          """
          This function is used to process the dependency between the size node and its users after
          inserting the size_process_node.
          """
          # store original node and processing node pair in node_pairs dictionary
          # It will be used to replace the original node with processing node in slice object
          node_pairs[node] = size_processing_node
          size_processing_node._meta_data = node._meta_data
  
          if hasattr(node.meta["info"], "activation_checkpoint"):
              MetaInfo(
                  size_processing_node,
                  mod_dir=node.meta["info"].mod_dir,
                  activation_checkpoint=tuple(
                      node.meta["info"].activation_checkpoint
                  ),
              )
  
          user_list = list(node.users.keys())
          for user in user_list:
              if user == size_processing_node:
                  continue
              new_args = list(user.args)
              new_kwargs = dict(user.kwargs)
              # the origin node may be a positional argument or key word argument of user node
              if node in new_args:
                  # substitute the origin node with size_processing_node
                  new_args[new_args.index(node)] = size_processing_node
                  user.args = tuple(new_args)
              elif str(node) in new_kwargs:
                  # substitute the origin node with size_processing_node
                  new_kwargs[str(node)] = size_processing_node
                  user.kwargs = new_kwargs
  
      def _update_slice_object_args(slice_object):
          """
          This function is used to update the slice object argument list.
          If the slice object contains the Node argument, then the size node will be replaced with
          """
          if isinstance(slice_object, slice):
              start = slice_object.start
              stop = slice_object.stop
              step = slice_object.step
              if start in node_pairs:
                  start = node_pairs[start]
              if stop in node_pairs:
                  stop = node_pairs[stop]
              if step in node_pairs:
                  step = node_pairs[step]
              return slice(start, stop, step)
          elif isinstance(slice_object, int):
              if slice_object in node_pairs:
                  return node_pairs[slice_object]
              else:
                  return slice_object
          else:
              raise RuntimeError(
                  f"Unsupported slice object type: {type(slice_object)}"
              )
  
      for node in nodes:
          if node.op == "call_method" and node.target == "size":
              # extract useful information from size node
              # dim_partition_dict will instruct the size value on which
              # dimension should be enlarged.
              sharding_spec = node.args[0].sharding_spec
              dim_partition_dict = sharding_spec.dim_partition_dict
  
              target_dim = _extract_target_dim(node)
  
              # insert size_processing node
              with mod_graph.inserting_after(node):
                  size_processing_node = mod_graph.create_node(
                      "call_function",
                      size_processing,
                      args=(
                          node,
                          dim_partition_dict,
                          device_mesh_info,
                          target_dim,
                          node.name,
                      ),
                  )
              _post_processing(node, size_processing_node)
  
          if (
              node.op == "call_function"
              and node.target == operator.getitem
          ):
              getitem_index = node.args[1]
              # slice object is quite special in torch.fx graph,
              # On one side, we treat slice object same as type of int,
              # so we do not create a node for slice object. On the other side,
              # slice object could take fx.Node as its argument. And the user
              # relationship cannot be tracked in fx graph.
              # Therefore, I record the node_pairs in this pass, and use the it
              # to replace the original node argument inside the slice object if
              # it has been processed in above pass.
  
              # There are three main usages of operator.getitem:
              #   getitem(input, int)
              #   getitem(input, slice)
              #   getitem(input, Tuple[slice])
              # In this pass, we need process the last two cases because
              # node arguments may potentially appear in these cases.
              if isinstance(getitem_index, slice):
                  new_slice_item = _update_slice_object_args(
                      getitem_index
                  )
                  new_args = (node.args[0], new_slice_item)
                  node.args = new_args
  
              elif isinstance(getitem_index, (tuple, list)):
                  if not isinstance(getitem_index[0], slice):
                      continue
                  new_slice_items = []
  
                  for slice_item in getitem_index:
                      if slice_item is None:
                          new_slice_items.append(None)
                          continue
                      new_slice_item = _update_slice_object_args(
                          slice_item
                      )
                      new_slice_items.append(new_slice_item)
  
                  new_args = (node.args[0], tuple(new_slice_items))
                  node.args = new_args
  
      return gm
  
file_path: colossalai/colossalai/auto_parallel/passes/runtime_preparation_pass.py
imports: import operator
  from typing import Dict, List, Union
  import torch
  from torch.fx.node import Node
  from colossalai._analyzer.fx.node_util import MetaInfo
  from colossalai.auto_parallel.tensor_shard.constants import RESHAPE_FUNC_OP
  from colossalai.auto_parallel.tensor_shard.sharding_strategy import CommType, OperationDataType
  from colossalai.auto_parallel.tensor_shard.solver.strategies_constructor import StrategiesConstructor
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.tensor.comm_spec import _all_reduce
  from colossalai.tensor.shape_consistency import ShapeConsistencyManager
  from colossalai.tensor.sharding_spec import ShardingSpec
  from .constants import SHAPE_ARGUMENT_OPS
label_desc: Use image node_id 3 for calling a global function with example usage: size_value_converting_pass(gm, device_mesh) and returns: gm, target_dim, slice, node_pairs, slice_object