code: def aot_dispatch_autograd_graph(
      flat_fn,
      flat_args: List[Any],
      aot_config: AOTConfig,
      *,
      fw_metadata: ViewAndMutationMeta,
  ) -> Union[
      Callable, Tuple[Callable, List[Any], Optional[SubclassMeta]]
  ]:
      # traced_tangents corresponds to the set of outputs in the traced forward that should get grad_outputs in the traced backward.
      # It includes outputs of the original forward, *and* any updated inputs due to input mutations.
      # However, it does *not* include any outputs that are aliases of inputs or intermediates, or any metadata-only input mutations.
      traced_tangents = pytree.tree_map(
          lambda x: x.detach().contiguous()
          if isinstance(x, Tensor)
          else x,
          fw_metadata.traced_tangents,
      )
  
      joint_inputs = (flat_args, traced_tangents)
  
      fn_prepared_for_autograd = fn_prepped_for_autograd(
          flat_fn,
          fw_metadata,
      )
      joint_fn_to_trace = create_joint(
          fn_prepared_for_autograd, aot_config=aot_config
      )
  
      (
          joint_fn_to_trace,
          updated_joint_inputs,
      ) = create_functionalized_fn(
          joint_fn_to_trace,
          joint_inputs,
          meta=fw_metadata,
          aot_config=aot_config,
          trace_joint=True,
      )
  
      subclass_tracing_info = aot_dispatch_subclass(
          joint_fn_to_trace,
          updated_joint_inputs,
          is_joint_structure=True,
          meta=fw_metadata,
          fw_only=flat_fn,
      )
  
      joint_fn_to_trace = subclass_tracing_info.plain_tensor_trace_fn
      updated_joint_inputs = subclass_tracing_info.plain_tensor_args
      maybe_subclass_meta = subclass_tracing_info.maybe_subclass_meta
  
      fx_g = _create_graph(
          joint_fn_to_trace, updated_joint_inputs, aot_config=aot_config
      )
  
      # There should be *NO* mutating ops in the graph at this point.
      assert_functional_graph(fx_g.graph)
  
      # Redundant with the check above, but worth having in case tracing introduced
      # a fake tensor. Unlikely.
      # See Note: [Fake Modules and AOTAutograd]
      torch._dynamo.utils.assert_no_fake_params_or_buffers(fx_g)
      fx_g.graph.eliminate_dead_code()
      fx_g.recompile()
      # TODO: in AOTAutograd, we create metadata like _indices_of_inps_to_detach to detect
      # when we need to manually detach() some inputs in the forward.
      # Higher order ops might eventually need to do the same.
      if aot_config.is_export:
          assert (
              maybe_subclass_meta is None
          ), "aot_export_module does not support tensor subclass inputs for now."
          return fx_g
      return fx_g, updated_joint_inputs, maybe_subclass_meta
  
file_path: pytorch/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py
imports: from typing import Any, Callable, List, Optional, Tuple, Union
  import torch
  import torch.utils._pytree
  import torch.utils.dlpack
  from torch import Tensor
  from torch._dispatch.python import enable_python_dispatcher
  from torch._dynamo.utils import lazy_format_graph_code
  from torch._logging import getArtifactLogger
  from torch._subclasses.functional_tensor import FunctionalTensorMode
  from torch.fx.experimental.proxy_tensor import make_fx
  from .functional_utils import assert_functional_graph
  from .schemas import AOTConfig, SubclassMeta, ViewAndMutationMeta
  from .traced_function_transforms import aot_dispatch_subclass, create_functionalized_fn, create_joint, fn_input_mutations_to_outputs, fn_prepped_for_autograd
label_desc: Use image node_id 3 for calling a global function with example usage: aot_dispatch_autograd_graph(flat_fn, flat_args, aot_config) and returns: fx_g,  updated_joint_inputs,  maybe_subclass_meta, fx_g