code: def partitioner(graph: GraphModule) -> GraphModule:
      """Graph partitioner that partitions the single device graph to distributed graph."""
      shape_adjustment_ops = {
          aten._unsafe_view.default: 1,
          aten.expand.default: 1,
          aten.new_zeros.default: 1,
          aten.ones.default: 0,
          aten.reshape.default: 1,
          aten.view.default: 1,
          aten.zeros.default: 0,
      }
      # partition the graph to distributed
      for node in graph.graph.nodes:
          node_sharding = node.meta["sharding"]
          # None sharding means this node don't need sharding
          if node_sharding is None:
              continue
  
          if node.op == "placeholder":
              out_spec = node_sharding.output_spec
              if not hasattr(out_spec, "from_local"):
                  local_val = _partition_val(node.meta["val"], out_spec)
                  # update node value
                  node.meta["val"] = local_val
          elif node.op == "call_function":
              out_spec = node_sharding.output_spec
  
              # check if there's misaligned sharding, insert reshard if there is
              expected_input_specs = node_sharding.input_specs
              for idx, input_arg in enumerate(node.all_input_nodes):
                  input_arg_sharding = input_arg.meta["sharding"]
  
                  input_arg_spec = input_arg_sharding.output_spec
                  desired_spec = (
                      out_spec
                      if expected_input_specs is None
                      else expected_input_specs[idx]
                  )
                  if input_arg_spec != desired_spec:
                      input_arg_spec.tensor_meta = input_arg.meta[
                          "tensor_meta"
                      ]
                      desired_spec.tensor_meta = input_arg.meta[
                          "tensor_meta"
                      ]
                      input_arg_tensor = input_arg.meta["val"]
  
                      # insert reshard operation
                      def reshard_fn(
                          local_tensor: torch.Tensor,
                      ) -> torch.Tensor:
                          return redistribute_local_tensor(
                              local_tensor,
                              input_arg_spec,
                              desired_spec,
                          )
  
                      reshard_gm = make_fx(reshard_fn)(input_arg_tensor)
                      reshard_gm_nodes = list(reshard_gm.graph.nodes)
                      input_node = reshard_gm_nodes[0]
                      with graph.graph.inserting_before(node):
                          output_node = graph.graph.graph_copy(
                              reshard_gm.graph,
                              val_map={
                                  input_node: input_arg,
                              },
                          )
                      node.replace_input_with(input_arg, output_node)
  
              output_val = node.meta["val"]
  
              if node.target == torch.ops.aten.repeat.default:
                  # for repeat op, we need to infer the repeat sizes
                  assert isinstance(output_val, torch.Tensor)
                  local_shape = compute_local_shape(
                      output_val.shape,
                      out_spec.mesh,
                      out_spec.placements,
                  )
                  input_shape = node.args[0].meta["val"].shape
  
                  def infer_repeat_sizes(repeated_shape, input_shape):
                      repeated_size = [1] * len(repeated_shape)
                      padded_length = len(repeated_shape) - len(
                          input_shape
                      )
                      for i in range(len(repeated_shape)):
                          if i < padded_length:
                              repeated_size[i] = repeated_shape[i]
                          else:
                              repeated_size[i] = (
                                  repeated_shape[i]
                                  // input_shape[i - padded_length]
                              )
  
                      return repeated_size
  
                  node.update_arg(
                      1, infer_repeat_sizes(local_shape, input_shape)
                  )
  
              elif node.target in shape_adjustment_ops:
                  # for view related op that needs shape, adjust shape to local shape if needed
                  assert isinstance(output_val, torch.Tensor)
                  local_shape = compute_local_shape(
                      output_val.shape,
                      out_spec.mesh,
                      out_spec.placements,
                  )
                  shape_arg_num = shape_adjustment_ops[node.target]
                  node.update_arg(shape_arg_num, local_shape)
  
              # convert output val to its local component
              node.meta["val"] = _partition_val(output_val, out_spec)
  
          elif node.op == "output":
              break
          else:
              raise RuntimeError(f"op code {node} not supported")
  
      # clean up the graph by removing sharding and partitioning related metadata
      for node in graph.graph.nodes:
          if "sharding" in node.meta:
              del node.meta["sharding"]
          if "val" in node.meta and isinstance(
              node.meta["val"], torch.Tensor
          ):
              local_tensor_meta = _extract_tensor_metadata(
                  node.meta["val"]
              )
              node.meta["tensor_meta"] = local_tensor_meta
  
      graph.graph.lint()
      graph.recompile()
      return graph
  
imports: import operator
  from contextlib import contextmanager
  from enum import Enum
  from typing import Any, cast, Dict, List, Optional, Tuple
  import torch
  import torch.distributed.distributed_c10d
  import torch.fx
  import torch.library
  import torch.nn
  import torch.utils._pytree
  from torch.distributed._spmd.batch_dim_utils import BatchDimAnalyzer
  from torch.distributed._tensor import DeviceMesh, distribute_tensor, Replicate, Shard
  from torch.distributed._tensor._utils import compute_local_shape
  from torch.distributed._tensor.op_schema import OpStrategy, PlacementStrategy, StrategyType, TupleStrategy
  from torch.distributed._tensor.placement_types import _Partial, DTensorSpec, Placement
  from torch.distributed._tensor.redistribute import redistribute_local_tensor
  from torch.fx import GraphModule
  from torch.fx.experimental.proxy_tensor import make_fx
  from torch.fx.passes.shape_prop import _extract_tensor_metadata
  from torch.nn.utils._named_member_accessor import NamedMemberAccessor
label_desc: Use image node_id 8 for calling a global function with example usage: partitioner(graph) and returns: graph, repeated_size, redistribute_local_tensor