code: def _convert_output(
      gm: fx.GraphModule,
      node: fx.Node,
      node_to_obj: Dict[fx.Node, Any],
  ) -> fx.Node:
      new_args = []
      has_partial = False
      for argument in node.args[0]:  # type: ignore[union-attr]
          if not isinstance(argument, fx.Node):
              new_args.append(argument)
              continue
  
          obj = node_to_obj[argument]
  
          if not _is_partial_dtensor(obj):
              new_args.append(argument)
              continue
  
          has_partial = True
  
          # we know it's a dtensor from is partial DT check...
          dt = cast(DTensor, obj)
  
          traced_dispatch, result_obj = _build_dummy_add_graph(
              dt, node_to_obj
          )
  
          wait = [
              n
              for n in traced_dispatch.graph.nodes
              if n.name == "wait_comm" or n.name == "wait_tensor"
          ]
          add = [
              n for n in traced_dispatch.graph.nodes if n.name == "add"
          ]
          assert len(wait) == 1 and len(add) == 1
  
          # remove add node and replace it with wait node
          add[0].replace_all_uses_with(wait[0])
          traced_dispatch.graph.eliminate_dead_code()
          # also update the actual DTensor corresponding to the node
          # TODO(anj): We require mapping of the final DTensor output to the wait
          # comm node.
          node_to_obj[wait[0]] = result_obj
  
          value_remap: Dict[fx.Node, fx.Node] = {}
          for dtn in traced_dispatch.graph.nodes:
              if dtn.op == OP.PLACEHOLDER:
                  # do nothing, ignore placeholders, as it has
                  # already been prepared in value_remap
                  value_remap[dtn] = argument
              elif dtn.op == OP.OUTPUT:
                  assert (
                      len(dtn.args) == 1 and len(dtn.args[0]) == 1
                  ), f"Expecting single output, but got {dtn.args} {len(dtn.args)}"
                  new_args.append(value_remap[dtn.args[0][0]])
                  # the concrete DTensor value of output was added when creating the
                  # inner graph (in _build_dummy_add_graph). Just add it to the final
                  # output node so that we can report the final output specs correctly.
                  # TODO(anj): We are depending on the concrete DTensor output of the dummy add.
                  node_to_obj[
                      value_remap[dtn.args[0][0]]
                  ] = node_to_obj[dtn.args[0][0]]
  
              else:
                  if dtn.op == OP.GET_ATTR:
                      setattr(
                          gm,
                          dtn.target,
                          getattr(traced_dispatch, dtn.target),
                      )
                  with gm.graph.inserting_before(node):
                      value_remap[dtn] = gm.graph.node_copy(
                          dtn, lambda n: value_remap[n]
                      )
      if has_partial:
          gm.graph.erase_node(node)
          return gm.graph.output(new_args)
      else:
          return node
  
imports: import logging
  import operator
  from dataclasses import dataclass
  from enum import auto, Enum
  from functools import partial
  from typing import Any, Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
  import torch
  import torch.distributed._spmd.experimental_ops
  import torch.fx
  from torch.distributed._spmd.comm_tensor import _get_tracer
  from torch.distributed._spmd.graph_utils import OP
  from torch.distributed._spmd.log_utils import get_logger
  from torch.distributed._tensor import DeviceMesh, DTensor
  from torch.distributed._tensor.op_schema import OpSchema
  from torch.distributed._tensor.placement_types import _Partial, DTensorSpec, Placement, Replicate, Shard, TensorMeta
  from torch.distributed._tensor.redistribute import redistribute_local_tensor
  from torch.fx.experimental.proxy_tensor import make_fx, proxy_slot
  from torch.utils import _pytree
  from torch.utils._pytree import tree_flatten, tree_map, tree_map_only, tree_unflatten
label_desc: Use image node_id 14 for calling a global function with example usage: _convert_output(gm, node, node_to_obj) and returns: gm, node