code: def _split_fused_adam(
      gm: IterGraphModule,
      orig_optim_block: FusedOptimizerBlock,
      split_gradients: Set[fx.Node],
  ) -> Tuple[FusedOptimizerBlock, FusedOptimizerBlock]:
      """Split the `orig_optim_block` into two FusedOptimizerBlock.
  
      The first one will be the optimizer that optimize `split_gradients`. The second one is
      used to optimize the remaining gradients.
      An assert will be raised if one of the optimizer optimize zero gradients.
      """
      orig_optim_args = AdamArgs(
          *orig_optim_block.optim.optim_node.args
      )
      optim_args = (
          AdamArgs([], [], [], [], [], []),
          AdamArgs([], [], [], [], [], []),
      )
      # The only hint we can use to split the optimizer is the order/indices.
      orig_optim_indices: Tuple[List[int], List[int]] = ([], [])
      orig_step_indices: Tuple[List[int], List[int]] = ([], [])
  
      for idx, gradient in enumerate(orig_optim_args.grads):
          group_idx = 0 if gradient in split_gradients else 1
          orig_optim_indices[group_idx].append(idx)
          # Get the argument for idx-th gradient from orig_optim_args
          for orig_arg, optim_arg in zip(
              orig_optim_args, optim_args[group_idx]
          ):
              # Only add the argument to the list if the original argument list
              # is not empty. If the original argument list is empty, the new
              # one must be an empty list as well.
              if orig_arg:
                  optim_arg.append(orig_arg[idx])
  
          # If argument order of step is the same as optimizer, nothing has to be
          # done. However, it is risky to rely on this assumption so we populate
          # the orig_step_indices.
          orig_step_output = optim_args[group_idx].state_steps[-1]
          assert str(orig_step_output.target).startswith(
              "aten.copy_"
          ), f"The copy output is {orig_step_output.target}, expect aten.copy_"
          orig_step_getitem = orig_step_output.args[1]
          assert "getitem" in str(
              orig_step_getitem.target
          ), f"The copy getitem is {orig_step_getitem.target}, expect operator.getitem"
          orig_step_idx = orig_step_getitem.args[1]
          orig_step_indices[group_idx].append(orig_step_idx)
  
      if not all(l for l in (orig_step_indices + orig_optim_indices)):
          raise ValueError(
              "At least one split optimizer does not have input."
          )
  
      output = get_output(gm.graph)
      results: List[FusedOptimizerBlock] = []
      flatten_output_args, spec = tree_flatten(
          (output.args, output.kwargs)
      )
      flatten_output_args_indices: DefaultDict[
          fx.Node, Set[int]
      ] = collections.defaultdict(set)
      for idx, output_arg in enumerate(flatten_output_args):
          if isinstance(output_arg, fx.Node):
              flatten_output_args_indices[output_arg].add(idx)
  
      def replace_flatten_output_args(
          orig_node: fx.Node, new_node: fx.Node
      ):
          for idx in flatten_output_args_indices[orig_node]:
              flatten_output_args[idx] = new_node
  
      # Create the new step and optim nodes and blocks.
      for group_idx in range(2):
          step_args: List[fx.Node] = []
          orig_step_outputs: List[fx.Node] = []
          # We have to create the new step node and block first because it is used
          # for the new optim node as the input.
          with gm.graph.inserting_after(
              orig_optim_block.optim.optim_node
          ):
              for idx in orig_step_indices[group_idx]:
                  step_args.append(
                      cast(
                          Tuple[fx.Node, ...],
                          orig_optim_block.step.add_node.args[0],
                      )[idx]
                  )
                  orig_step_outputs.append(
                      orig_optim_block.step.outputs[idx]
                  )
              step = gm.graph.call_function(
                  aten._foreach_add.Scalar,
                  (step_args, 1),
              )
          step_block = ForeachAddBlock(step, generate_output=True)
          for i, step_output in enumerate(step_block.outputs):
              # Replace the original step output in the graph output node with
              # the new one.
              orig_step_output = orig_step_outputs[i]
              replace_flatten_output_args(orig_step_output, step_output)
              # Also need to replace the step output used for the new optimizer.
              assert (
                  optim_args[group_idx].state_steps[i]
                  == orig_step_output
              ), (
                  f"The expected step output node mismatched, {orig_step_output} "
                  f"{optim_args[group_idx].state_steps[i]}"
              )
              optim_args[group_idx].state_steps[i] = step_output
  
          # Insert the optimizer node after the first step output because its
          # topo sort order is the last.
          with gm.graph.inserting_after(step_block.outputs[0]):
              optim = gm.graph.call_function(
                  aten._fused_adam.default,
                  optim_args[group_idx],
                  orig_optim_block.optim.optim_node.kwargs,
              )
          optim_block = FusedAdamBlock(optim, generate_output=True)
          for curr_idx, orig_idx in enumerate(
              orig_optim_indices[group_idx]
          ):
              list_names = (
                  "param_outputs",
                  "exp_avgs_outputs",
                  "exp_avg_sqs_outputs",
              )
              for name in list_names:
                  orig_list = getattr(orig_optim_block.optim, name)
                  curr_list = getattr(optim_block, name)
                  replace_flatten_output_args(
                      orig_list[orig_idx], curr_list[curr_idx]
                  )
  
          results.append(FusedOptimizerBlock(step_block, optim_block))
  
      # Optimizer is used as the output of the train_step. Therefore, we have to
      # update the output node of the graph.
      output_args, output_kwargs = tree_unflatten(
          flatten_output_args, spec
      )
      gm.graph.node_set_args(output, output_args)
      gm.graph.node_set_kwargs(output, output_kwargs)
      # Remove the original copy_ nodes as they won't be DCE.
      for copy_output in itertools.chain(
          orig_optim_block.optim.param_outputs,
          orig_optim_block.optim.exp_avgs_outputs,
          orig_optim_block.optim.exp_avg_sqs_outputs,
      ):
          gm.graph.erase_node(copy_output)
      # Call DCE once to get rid of the old optimizer. By doing so, we will be
      # able to erase the copy_ nodes of step later.
      gm.graph.eliminate_dead_code()
      for copy_output in orig_optim_block.step.outputs:
          gm.graph.erase_node(copy_output)
      # This is not required but calling this for consistency.
      gm.graph.eliminate_dead_code()
  
      return results[0], results[1]
  
imports: import collections
  import itertools
  import logging
  import operator
  import tempfile
  import time
  from dataclasses import dataclass, field
  from functools import wraps
  from typing import Any, Callable, cast, DefaultDict, Dict, Iterable, List, Optional, Set, Tuple, Union
  import torch
  import torch.fx
  from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
  from torch.distributed._spmd.graph_utils import CommType, dump_graphs_to_files, find_node, get_output, OP
  from torch.distributed._spmd.iter_graph_module import IterGraphModule
  from torch.fx.passes.shape_prop import TensorMetadata
  from torch.utils import _pytree
  from torch.utils._pytree import tree_flatten, tree_unflatten
label_desc: Use image node_id 16 for calling a global function with example usage: _split_fused_adam(gm, orig_optim_block, split_gradients) and returns: results,  results