content: #!/usr/bin/env python3
  
  import dataclasses as dc
  import logging
  import operator
  import typing as t
  
  import torch.fx as fx
  
  _LOGGER = logging.getLogger(__name__)
  
  RemoveDuplicateOutputArgsFunc = t.Callable[
      [
          fx.GraphModule,
          t.Collection[str],
      ],
      t.Mapping[str, "RemoveDuplicateResult"],
  ]
  
  
  def remove_duplicate_output_args(
      top_level: fx.GraphModule, target_subnets: t.Collection[str]
  ) -> t.Mapping[str, "RemoveDuplicateResult"]:
      """Removes duplicate output args.
  
      This pass removes duplicate output args from the target subnets and fixes
      their uses in the top level module where the subnets are called. This pass
      must be called after acc split on the top-level net and subsequent calls to
      the acc trace on the subnets.
  
      This pass will change both the subnets and top level module.
  
      Returns:
          a mapping of the target subnet name to its dedupcate result
      """
  
      processed_subnets = {}
      for node in top_level.graph.nodes:  # type: fx.Node
          if node.op == "call_module" and node.name in target_subnets:
              assert isinstance(node.target, str)
              sub_gm = top_level.get_submodule(node.target)
              assert isinstance(sub_gm, fx.GraphModule)
  
              replace_res = _remove_duplicate_output_args(sub_gm)
              processed_subnets[node.name] = replace_res
              if replace_res.replacement_map is None:
                  continue
              sub_gm.recompile()
  
              needs_recompile = False
              # iterate on the copy since we will be changing elements of node.users
              for user in list(node.users):
                  idx = _ensure_proper_output_use(user, node)
                  idx_new = replace_res.replacement_map[idx]
                  if idx_new != idx:
                      user.args = (user.args[0], idx_new)
                      needs_recompile = True
  
              if needs_recompile:
                  top_level.recompile()
      return processed_subnets
  
  
  @dc.dataclass(frozen=True)
  class RemoveDuplicateResult:
      replacement_map: t.Optional[t.List[int]]
      module: fx.GraphModule
  
  
  def _ensure_proper_output_use(user: fx.Node, target_node: fx.Node) -> int:
      """
      Ensures the node looks in proper form of calling the output of an fx2trt
      splitter sub-net. Specifically:
  
      1. op is call function, target: operator.getitem
      2. args is a 2-element tuple
      3. args[0] is the name of the subnet's output
      4. args[1] is the index into the subnet output tuple
  
      E.g.:
  
          %getitem_4 : [#users=1] = call_function[target=operator.getitem](args = (%_run_on_acc_1, 4), kwargs = {})
  
      returns the index into the subnet output tuple
      """
      assert (
          user.op == "call_function"
          and user.target == operator.getitem
          and len(user.args) == 2
          and isinstance(user.args[0], fx.Node)
          and user.args[0].name == target_node.name
          and isinstance(user.args[1], int)
      ), f"Node is not a proper user of splitter output: {user.format_node()}"
  
      return user.args[1]
  
  
  def _remove_duplicate_output_args(gm: fx.GraphModule) -> RemoveDuplicateResult:
      output_nodes = [n for n in gm.graph.nodes if n.op == "output"]
      assert (
          len(output_nodes) == 1
      ), f"Expecting exactly one `output` node, but got {len(output_nodes)}"
  
      changed = False
      # arg node name to its index in the new output args tuple
      name_to_idx: t.Dict[str, int] = {}
      output_node = output_nodes[0]
  
      # Output op only uses its `args[0]`, and it does not have `kwargs`.
      # https://pytorch.org/docs/stable/fx.html#torch.fx.Node
      args: t.Sequence[t.Any] = output_node.args[0]
  
      # Only concern outselves to the case where the args is an iterable of fx.Node.
      # Other return cases (e.g., a single value) is possible and we don't handle
      # that in this pass.
      if not (isinstance(args, t.Iterable) and all(isinstance(a, fx.Node) for a in args)):
          return RemoveDuplicateResult(replacement_map=None, module=gm)
  
      # Map old index of the arg node to the remaining node's idx,
      # initialized to `i => i`
      replacement_map: t.List[int] = list(range(len(args)))
      args_new = []
      for idx, a in enumerate(args):
          assert isinstance(a, fx.Node), f"Expecting fx.Node instance, but got: {type(a)}"
  
          if a.name not in name_to_idx:
              args_new.append(a)
              name_to_idx[a.name] = len(args_new) - 1
          else:
              changed = True
              _LOGGER.warning(
                  f"Replaced duplicate output arg '{a.name}': "
                  f"{idx} -> {name_to_idx[a.name]}"
              )
          replacement_map[idx] = name_to_idx[a.name]
  
      output_node.args = (tuple(args_new),)
      if changed:
          gm.recompile()
      return RemoveDuplicateResult(replacement_map, module=gm)
  