code: def build_data_parallel_strategies(
      train_step_graph: GraphModule,
      num_params: int,
      num_states: int,
      mesh: DeviceMesh,
      batch_dim: int = 0,
  ) -> Dict[fx.Node, StrategyType]:
      """Loop through the train step graph and build the data parallel strategy for each fx Node."""
      activation_idx = num_params + num_states
      non_compute_ops = [
          aten.clone.default,
          aten.detach.default,
          aten.ones_like.default,
          aten.reshape.default,
          aten.t.default,
          aten.view.default,
          torch.ops._spmd.tag_grad.default,
          operator.getitem,
      ]
  
      tuple_strategy_ops = [aten._fused_adam.default]
  
      dp_strategy_map: Dict[fx.Node, StrategyType] = {}
      batch_dim_analyzer = BatchDimAnalyzer(batch_dim)
      placeholder_idx = 0
      num_param_grad = 0
  
      # first we backward propagate to mark the param gradients sharding
      # with tag_grad node helps and then delete the tag_grad nodes
      for node in reversed(list(train_step_graph.graph.nodes)):
          # find a param_grad node via the tagging
          if node.target == torch.ops._spmd.tag_grad.default:
              cur_node = node
              while cur_node.target in non_compute_ops:
                  cur_node = cur_node.args[0]
                  partial_strategy = _gen_partial_strategy(mesh)
                  dp_strategy_map[cur_node] = DataParallelStrategy(
                      NodeType.GRAD, [partial_strategy]
                  )
              num_param_grad += 1
              # remove the tag_grad node from graph
              node.replace_all_uses_with(node.args[0])
              train_step_graph.graph.erase_node(node)
  
              if num_param_grad == num_params:
                  # early break if we have already processed all param_grads
                  break
  
      # next we forward propagate to mark all the sharding
      for node in train_step_graph.graph.nodes:
          if node.op == "placeholder":
              if "val" not in node.meta:
                  # NOTE: There're certain cases where the placeholder nodes do
                  # not have real tensor values:
                  # 1. optimizer states can be None sometimes, i.e. SGD with
                  #    no momentum, optimizer states populate `momentum` state
                  #    as None, the full graph we get from `compile` would have
                  #    None as the placeholder value
                  # 2. function args might not only contain params or activations,
                  #    but also contain other non-tensor inputs, i.e. the model
                  #    and optimizer instances baked in as a placeholder, there might
                  #    also be some scalar argument which is not a tensor
                  #
                  # For the above cases, we create a NON_TENSOR stratgy so that we
                  # know it's not a tensor and we don't need to shard it
                  dp_strategy_map[node] = DataParallelStrategy(
                      NodeType.NON_TENSOR, []
                  )
  
              elif placeholder_idx < num_params:
                  # during compilation there's an assumption that the first num_params
                  # placeholders should be parameters
                  shard_strategy = _gen_shard_strategy(mesh, 0)
                  replica_strategy = _gen_replicate_strategy(mesh)
                  dp_strategy_map[node] = DataParallelStrategy(
                      NodeType.PARAM, [replica_strategy, shard_strategy]
                  )
  
              elif placeholder_idx < activation_idx:
                  # optimizer states follow the same strategy as
                  # the corresponding parameters
                  replica_strategy = _gen_replicate_strategy(mesh)
                  shard_strategy = _gen_shard_strategy(mesh, 0)
  
                  dp_strategy_map[node] = DataParallelStrategy(
                      NodeType.STATE, [replica_strategy, shard_strategy]
                  )
              else:
                  activation_batch_dim_size = node.meta["val"].shape[
                      batch_dim
                  ]
                  # find the first activation node and use its batch dim size
                  if batch_dim_analyzer.batch_dim_size == -1:
                      batch_dim_analyzer.init_batch_dim_size(
                          activation_batch_dim_size
                      )
  
                  batch_dim_analyzer.set_batch_dim(node, batch_dim)
                  shard_strategy = _gen_shard_strategy(mesh, batch_dim)
                  dp_strategy_map[node] = DataParallelStrategy(
                      NodeType.ACT, [shard_strategy]
                  )
              placeholder_idx += 1
          elif node.op == "call_function":
              # Annotate node types for the computation graph
              # Data Parallel node propagation logic:
              # param (non-compute) -> out: param
              # grad (non-compute before/after) -> out: grad
              # state -> output: state
              #
              # param + activation (param must be replicate, act be sharded) -> out: activation
              # param/state + grad (param/state/grad be the same spec) -> out: param/state
              # param + state -> out: param
  
              if node.target in non_compute_ops:
                  # At this point, we should have removed all the `tag_grad` nodes in the graph
                  assert node.target != torch.ops._spmd.tag_grad.default
  
                  input_nodes = node.all_input_nodes
                  assert (
                      len(input_nodes) == 1
                  ), f"non-compute op only support one input now, found node: {node} with length of inputs: {len(node.args)}"
                  arg_strategy = dp_strategy_map[input_nodes[0]]
  
                  if node.target == operator.getitem:
                      # for getitem call, just forward the strategy from the input
                      getitem_idx = node.args[1]
                      if isinstance(arg_strategy, TupleStrategy):
                          # for tuple strategy, we need to get the child strategy from the tuple
                          dp_strategy_map[node] = arg_strategy.childs[
                              getitem_idx
                          ]
                      else:
                          # if it's not a tuple strategy, we just forward the arg strategy
                          dp_strategy_map[node] = arg_strategy
                  else:
                      assert isinstance(
                          arg_strategy, DataParallelStrategy
                      )
                      arg_node_type = arg_strategy.node_type
                      if arg_node_type == NodeType.PARAM:
                          replica_strategy = _gen_replicate_strategy(
                              mesh
                          )
                          dp_strategy_map[node] = DataParallelStrategy(
                              NodeType.PARAM, [replica_strategy]
                          )
                      elif arg_node_type == NodeType.GRAD:
                          partial_sig = _gen_partial_strategy(mesh)
                          dp_strategy_map[node] = DataParallelStrategy(
                              NodeType.GRAD, [partial_sig]
                          )
                      elif arg_node_type == NodeType.ACT:
                          arg_node_spec = (
                              batch_dim_analyzer.compute_act_spec(
                                  input_nodes[0], mesh
                              )
                          )
  
                          output_spec = (
                              batch_dim_analyzer.compute_act_spec(
                                  node, mesh
                              )
                          )
  
                          shard_strategy = PlacementStrategy(
                              output_spec=output_spec,
                              input_specs=[arg_node_spec],
                          )
                          dp_strategy_map[node] = DataParallelStrategy(
                              NodeType.ACT, [shard_strategy]
                          )
                      else:
                          raise RuntimeError(
                              f"non compute op not supporting {arg_node_type}! "
                          )
  
                  # finished processing this non-compute node
                  continue
  
              # for computatation nodes, we need to check all the inputs
              input_args = node.all_input_nodes
              input_specs = []
              if node in dp_strategy_map:
                  # found a param_grad node that already have output pre-filled spec
                  # fill in the expected input specs for the pre-filled strategy
                  node_strategy = dp_strategy_map[node]
                  assert isinstance(node_strategy, DataParallelStrategy)
                  node_type = node_strategy.node_type
                  assert node_type == NodeType.GRAD
                  produce_param_grad_strat = node_strategy.strategies
                  has_activation = False
                  for arg in input_args:
                      arg_strategy = dp_strategy_map[arg]
                      assert isinstance(
                          arg_strategy, DataParallelStrategy
                      )
                      arg_node_type = arg_strategy.node_type
                      if arg_node_type == NodeType.ACT:
                          # activation sharded
                          has_activation = True
                          act_spec = (
                              batch_dim_analyzer.compute_act_spec(
                                  arg, mesh
                              )
                          )
  
                          input_specs.append(act_spec)
  
                  if has_activation:
                      assert len(produce_param_grad_strat) == 1
                      produce_param_grad_strat[
                          0
                      ].input_specs = input_specs
              elif node.target in tuple_strategy_ops:
                  # ops that need to build tuple strategy instead of normal strategy
                  # This should happen rarely and only needed when we need to generate
                  # different node strategy for multiple outputs (i.e. fused_adam op)
                  # TODO: Currently this specializes to fused optimizer ops, but we need
                  # to see how to generalize this strategy building logic
                  output_strategy_len = len(node.args) - 1
                  tuple_strategies = []
                  for i in range(output_strategy_len):
                      if not isinstance(node.args[i], list):
                          raise RuntimeError(
                              f"Expecting list as arg to build Tuple Strategy, but found type {type(node.args[i])}!"
                          )
                      # for list/tuple arg, use the first one to find out the node type
                      if len(node.args[i]) > 0:
                          arg_strategy = dp_strategy_map[
                              node.args[i][0]
                          ]
                          assert isinstance(
                              arg_strategy, DataParallelStrategy
                          )
                          assert arg_strategy.node_type in [
                              NodeType.PARAM,
                              NodeType.GRAD,
                              NodeType.STATE,
                          ], "Expecting param/grad/state as arg to build Tuple Strategy!"
                          replica_strategy = _gen_replicate_strategy(
                              mesh
                          )
                          shard_strategy = _gen_shard_strategy(
                              mesh, shard_dim=0
                          )
                          out_node_strategy: StrategyType = (
                              DataParallelStrategy(
                                  arg_strategy.node_type,
                                  [replica_strategy, shard_strategy],
                              )
                          )
  
                          tuple_strategies.append(out_node_strategy)
  
                  output_tuple_strategy = TupleStrategy(
                      tuple(tuple_strategies)
                  )
                  dp_strategy_map[node] = output_tuple_strategy
              else:
                  # NOTE: This is the common region for all regular computation ops
  
                  input_node_types = [
                      cast(
                          DataParallelStrategy, dp_strategy_map[arg]
                      ).node_type
                      for arg in input_args
                      if isinstance(
                          dp_strategy_map[arg], DataParallelStrategy
                      )
                  ]
                  if NodeType.GRAD in input_node_types:
                      # param/state + grad, build up acceptable strategy
                      # the strategy should be the same for all the inputs/outputs
                      # TODO: optimizer parts should follow the dtensor prop logic
                      # to support more general cases that allows optimizer states
                      # to have different shardings compare to the params
                      replica_strategy = _gen_replicate_strategy(mesh)
                      shard_strategy = _gen_shard_strategy(
                          mesh, shard_dim=0
                      )
                      output_node_type = NodeType.PARAM
  
                      non_grad_types = [
                          t
                          for t in input_node_types
                          if t != NodeType.GRAD
                      ]
  
                      output_node_type = non_grad_types[0]
                      for non_grad_type in non_grad_types:
                          assert (
                              non_grad_type == output_node_type
                          ), f"Found more than one non grad types! Expect {output_node_type} but found {non_grad_type}!"
                      assert output_node_type in [
                          NodeType.PARAM,
                          NodeType.STATE,
                      ], f"Expecting output node type to be either state or param, but found {output_node_type}!"
  
                      dp_strategy_map[node] = DataParallelStrategy(
                          output_node_type,
                          [replica_strategy, shard_strategy],
                      )
                  elif NodeType.STATE in input_node_types:
                      # either param + state or state + state
                      replica_strategy = _gen_replicate_strategy(mesh)
                      shard_strategy = _gen_shard_strategy(
                          mesh, shard_dim=0
                      )
                      output_node_type = (
                          NodeType.PARAM
                          if NodeType.PARAM in input_node_types
                          else NodeType.STATE
                      )
  
                      dp_strategy_map[node] = DataParallelStrategy(
                          output_node_type,
                          [replica_strategy, shard_strategy],
                      )
                  elif NodeType.PARAM in input_node_types:
                      if NodeType.ACT in input_node_types:
                          # param + activation, build up acceptable strategy
                          # param must be replicated, activation must be sharded
                          for arg in input_args:
                              arg_strategy = dp_strategy_map[arg]
                              assert isinstance(
                                  arg_strategy, DataParallelStrategy
                              )
                              node_type = arg_strategy.node_type
                              if node_type == NodeType.ACT:
                                  # compute activation spec
                                  act_spec = batch_dim_analyzer.compute_act_spec(
                                      arg, mesh
                                  )
  
                                  input_specs.append(act_spec)
                              elif node_type == NodeType.PARAM:
                                  # param must be replicated
                                  input_specs.append(
                                      DTensorSpec(
                                          mesh=mesh,
                                          placements=(Replicate(),),
                                      )
                                  )
                              else:
                                  raise RuntimeError(
                                      f"Expecting node with parameter and activation, but found {input_node_types}! "
                                  )
                          # produce activation type sharding for output
                          output_spec = (
                              batch_dim_analyzer.compute_act_spec(
                                  node, mesh
                              )
                          )
  
                          act_strategy = PlacementStrategy(
                              output_spec=output_spec,
                              input_specs=input_specs,
                          )
  
                          dp_strategy_map[node] = DataParallelStrategy(
                              NodeType.ACT, [act_strategy]
                          )
                      else:
                          # If inputs only have parameters, the
                          # strategy of this node should follow input
                          dp_strategy_map[node] = dp_strategy_map[
                              input_args[0]
                          ]
                  else:
                      # If input nodes does not have PARAM/GRAD/STATE, then
                      # it should be a pure activation computation, it should
                      # produce activation output.
                      # Activations are usually sharded unless model creates
                      # new tensors during computation, which depend on whether
                      # the new tensor associate with a batch dim or not, it could
                      # be shard/replicate/partial, batch dim analyzer should tell
                      # us the correct sharding.
                      for arg in input_args:
                          arg_strategy = dp_strategy_map[arg]
                          assert isinstance(
                              arg_strategy, DataParallelStrategy
                          )
                          input_spec = (
                              batch_dim_analyzer.compute_act_spec(
                                  arg, mesh
                              )
                          )
  
                          input_specs.append(input_spec)
  
                      act_spec = batch_dim_analyzer.compute_act_spec(
                          node, mesh
                      )
                      op_strategy = PlacementStrategy(
                          output_spec=act_spec, input_specs=input_specs
                      )
                      dp_strategy_map[node] = DataParallelStrategy(
                          NodeType.ACT, [op_strategy]
                      )
  
          elif node.op == "output":
              dp_strategy_map[node] = DataParallelStrategy(
                  NodeType.NON_TENSOR, []
              )
          else:
              raise RuntimeError(f"op code {node.op} not supported")
  
      return dp_strategy_map
  
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 5 for calling a global function with example usage: build_data_parallel_strategies(train_step_graph, num_params, num_states, mesh, batch_dim) and returns: dp_strategy_map