code: def create_aot_dispatcher_function(
      flat_fn, flat_args: List[Any], aot_config: AOTConfig
  ):
      """
      Traces the forward and backward graphs of the attr:`flat_fn` to generate a
      joint graph. The joint graph is an Fx graph with Aten ops. Please refer to
      the tracing mechanism to understand the graph capturing details.
  
      The joint graph is then passed through attr:`partition_fn` to isolate the
      forward and backward portions, which are then respectively compiled via the
      provided attr:`fw_compiler` and attr:`bw_compiler`.
  
      The resulting compiled forward and backward graphs are then wrapped up in a
      ``torch.autograd.Function`` object.
  
      The calling convention here is that the first aot_config.num_params_buffers
      inputs in flat_args are parameters and buffers, and the rest are inputs.
  
      We use this to assume that parameters/buffer's shapes don't change.
  
      Note: this function is used both by aot_function and aot_export (controlled by aot_config.is_export)
          When aot_config.is_export is True, we return an FX graph + metadata
          When aot_config.is_export is False, we return an ordinary runtime function
      """
  
      # This is the main entry point.
      # TODO: Chillee argues that dynamo itself should pass in fake tensors to
      # the list of arguments when compiling; at the moment we do not do this
  
      if aot_config.decompositions is None:
          aot_config.decompositions = {}
  
      aot_config.decompositions = {
          **aot_autograd_decompositions,
          **aot_config.decompositions,
      }
  
      if config.functionalize_rng_ops:
          # Update the decompositions with functionalized random decompositions
          aot_config.decompositions = {
              **rng_decompositions,
              **aot_config.decompositions,
          }
  
      # Check flat_args to see if they're already fake.  If so, use that fake
      # mode instead.
  
      fake_mode = detect_fake_mode(flat_args)
      if fake_mode is None:
          shape_env = ShapeEnv() if aot_config.dynamic_shapes else None
          fake_mode = FakeTensorMode(shape_env=shape_env)
      else:
          shape_env = fake_mode.shape_env
  
      python_dispatcher_mode = (
          enable_python_dispatcher()
          if shape_env is not None
          else nullcontext()
      )
  
      with torch.autograd.set_multithreading_enabled(
          False
      ), preserve_rng_state(), (
          fake_mode
      ), python_dispatcher_mode, PhiloxStateTracker():
  
          def process_inputs(flat_args):
              def convert(idx, x):
                  if shape_env is not None:
                      from torch._dynamo.source import ConstantSource
  
                      if isinstance(x, int):
                          # We always specialize on scalar values in export.
                          if aot_config.is_export:
                              return x
                          source = ConstantSource(f"sym_{idx}")
                          return shape_env.create_symintnode(
                              shape_env.create_symbol(x, source),
                              hint=x,
                              source=source,
                          )
                  if not isinstance(x, torch.Tensor):
                      return x
                  if isinstance(x, FakeTensor):
                      assert x.fake_mode is fake_mode
                      return x
                  if is_traceable_wrapper_subclass(x):
                      attrs, _ = x.__tensor_flatten__()
                      if all(
                          isinstance(getattr(x, attr), FakeTensor)
                          for attr in attrs
                      ):
                          assert all(
                              getattr(x, attr).fake_mode is fake_mode
                              for attr in attrs
                          )
                          return x
  
                  # see note [Tensor Fakification and Symbol Caching]
                  symbolic_context = None
                  source = None
                  if (
                      tracing_context := torch._guards.TracingContext.try_get()
                  ):
                      if x in tracing_context.tensor_to_context:
                          symbolic_context = (
                              tracing_context.tensor_to_context[x]
                          )
                          source = symbolic_context.tensor_source
                  if (
                      idx < aot_config.num_params_buffers
                      and config.static_weight_shapes
                      and not symbolic_context
                  ):
                      # TODO: Ensure that this codepath is never exercised from
                      # Dynamo
                      return fake_mode.from_tensor(
                          x, static_shapes=True
                      )
  
                  return fake_mode.from_tensor(
                      x,
                      static_shapes=False,
                      symbolic_context=symbolic_context,
                      source=source,
                  )
  
              return [
                  convert(idx, x) for idx, x in enumerate(flat_args)
              ]
  
          fake_flat_args = process_inputs(flat_args)
  
          needs_autograd = (
              any(
                  x.requires_grad
                  for x in fake_flat_args
                  if isinstance(x, Tensor)
              )
              and torch.is_grad_enabled()
          )
  
          with enable_python_dispatcher():
              # Patch set_rng_state as set_rng_state with fake tensors is
              # nonsensical. This does not affect the collection of metadata.
              with patch(
                  "torch.cuda.set_rng_state", lambda *args: None
              ):
                  fw_metadata = run_functionalized_fw_and_collect_metadata(
                      flat_fn,
                      keep_input_mutations=aot_config.keep_inference_input_mutations,
                      is_train=needs_autograd,
                      pre_dispatch=aot_config.pre_dispatch,
                  )(
                      *fake_flat_args
                  )
  
                  req_subclass_dispatch = requires_subclass_dispatch(
                      fake_flat_args, fw_metadata
                  )
  
                  if needs_autograd and not any(
                      x.requires_grad for x in fw_metadata.output_info
                  ):
                      # We realized that none of the outputs require grad,
                      # so we actually have an inference graph.
                      needs_autograd = False
                      # A bit silly: right now in the subclass codepath, our ViewAndMutationMeta
                      # changes depending on whether we pass in is_train / keep_input_mutations,
                      # so we're forced to recompute the metadata.
                      # TODO: refactor the subclass path of run_functionalized_fw_and_collect_metadata
                      # so that this is unnecessary.
                      if req_subclass_dispatch:
                          fw_metadata = run_functionalized_fw_and_collect_metadata(
                              flat_fn,
                              keep_input_mutations=aot_config.keep_inference_input_mutations
                              and not needs_autograd,
                              is_train=needs_autograd,
                              pre_dispatch=aot_config.pre_dispatch,
                          )(
                              *fake_flat_args
                          )
                      else:
                          fw_metadata = ViewAndMutationMeta(
                              input_info=fw_metadata.input_info,
                              output_info=fw_metadata.output_info,
                              num_intermediate_bases=fw_metadata.num_intermediate_bases,
                              keep_input_mutations=aot_config.keep_inference_input_mutations
                              and not needs_autograd,
                              traced_tangents=fw_metadata.traced_tangents,
                              subclass_inp_meta=fw_metadata.subclass_inp_meta,
                              subclass_fw_graph_out_meta=fw_metadata.subclass_fw_graph_out_meta,
                              subclass_tangent_meta=fw_metadata.subclass_tangent_meta,
                              is_train=needs_autograd,
                          )
  
          if fw_metadata.num_intermediate_bases > 0:
              assert not req_subclass_dispatch, f"""\
  torch.compile is currently being used with tensor subclass inputs:
  {','.join([str(type(x)) for x in fake_flat_args])}. We are attempting to a compile a graph with two graph outputs
  that alias one another, which is currently unsupported in the subclass use case. If you run into this,
  please file a github issue"""
  
          if aot_config.is_export:
              # aot_export: ban input metadata mutations for now to keep shared code paths simpler.
              # Keeping .resize_() in the graph will require some work
              # Allowing it but keeping the graph functional will require some calling convention changes.
              if (
                  len(
                      [
                          x
                          for x in fw_metadata.input_info
                          if x.mutates_metadata
                      ]
                  )
                  != 0
              ):
                  raise RuntimeError(
                      f"""\
  Found an input that received a metadata mutation, through e.g. a call to `.resize_()` or `.transpose_()`.
  This is currently banned in the aot_export workflow. If you need this functionality, please file a github issue.
  
  fw_metadata={str(fw_metadata)}"""
                  )
              # In export, banning data mutations on inputs that require grad for now.
              # This should be rare, and is tricky to get right. When we trace the backward,
              # we currently trace with autograd.grad instead of .backward(), which makes it difficult
              # to ensure that we run autograd all the way through the input **before** it saw the mutation.
              if (
                  len(
                      [
                          x
                          for x in fw_metadata.input_info
                          if x.requires_grad and x.mutates_data
                      ]
                  )
                  != 0
              ):
                  raise RuntimeError(
                      f"""\
  Found a graph input that requires gradients, and received a mutation.
  This is currently banned in the aot_export workflow. If you need this functionality, please file a github issue.
  
  fw_metadata={str(fw_metadata)}"""
                  )
              if req_subclass_dispatch:
                  raise RuntimeError(
                      """\
  aot_export is not currently supported with traceable tensor subclass.
  If you need this feature, please comment on <CREATE_ISSUE_LINK>"""
                  )
  
              # Need to decide on a strategy for functionalized RNG: toggling via global config seems bad,
              # and turning it on will require a non-trivial calling convention change for any export runtime.
              if config.functionalize_rng_ops:
                  raise RuntimeError(
                      """\
  Functionalized RNG is not currently supported in the aot_export workflow. Please file a github issue,
  or otherwise set torch._functorch.config.functionalize_rng_ops = False."""
                  )
  
          # crappy version of dispatcher
          # TODO: Do this properly
          if needs_autograd:
              # For now, aot_dispatch_autograd knows to explicitly return a graph
              # when run with export, and an opaque callable otherwise.
              # In theory we could factor these out, but I wanted to let the dust
              # settle on how functionalized rng fits into export first.
              compiler_fn = (
                  aot_dispatch_autograd_graph
                  if aot_config.is_export
                  else aot_dispatch_autograd
              )
          else:
              # aot_dispatch_base_graph contains only the "graph bits", while aot_dispatch_base
              # includes some extra work around handling a runtime epilogue.
              compiler_fn = (
                  aot_dispatch_base_graph
                  if aot_config.is_export
                  else aot_dispatch_base
              )
  
          compiler_fn = partial(
              aot_wrapper_synthetic_base,
              compiler_fn=compiler_fn,
              needs_autograd=needs_autograd,
          )
          compiler_fn = partial(
              aot_wrapper_dedupe, compiler_fn=compiler_fn
          )
          # You can put more passes here
  
          compiled_fn = compiler_fn(
              flat_fn,
              fake_flat_args,
              aot_config,
              fw_metadata=fw_metadata,
          )
          if aot_config.is_export:
              mutated_user_inp_locs = [
                  idx - aot_config.num_params_buffers
                  for idx in fw_metadata.mutated_inp_runtime_indices
                  if idx >= aot_config.num_params_buffers
              ]
              if len(mutated_user_inp_locs) > 0:
                  raise RuntimeError(
                      f"""
  Found following user inputs located at {mutated_user_inp_locs} are mutated. This is currently banned in the aot_export workflow.
  If you need this functionality, please file a github issue.
  
  fw_metadata={str(fw_metadata)}"""
                  )
  
              # During export, we don't get back a callable - we get back the raw fx graph
              # (either a joint or an inference-only graph)
              assert isinstance(compiled_fn, torch.fx.GraphModule)
              return compiled_fn, fw_metadata
  
          if not hasattr(compiled_fn, "_boxed_call"):
              compiled_fn = make_boxed_func(compiled_fn)
  
          return compiled_fn
  
imports: import itertools
  from contextlib import nullcontext
  from functools import partial, wraps
  from typing import Any, Callable, Dict, List, Optional, Tuple
  from unittest.mock import patch
  import torch
  import torch.nn
  import torch.utils._pytree
  import torch.utils.dlpack
  from torch import Tensor
  from torch._dispatch.python import enable_python_dispatcher
  from torch._dynamo import compiled_autograd
  from torch._dynamo.utils import dynamo_timed, preserve_rng_state
  from torch._guards import detect_fake_mode
  from torch._subclasses import FakeTensor, FakeTensorMode
  from torch.fx.experimental.proxy_tensor import make_fx
  from torch.fx.experimental.symbolic_shapes import ShapeEnv
  from torch.utils._python_dispatch import is_traceable_wrapper_subclass
  from torch._decomp.decompositions_for_rng import PhiloxStateTracker, rng_decompositions
  from .None import config
  from .partitioners import default_partition
  from ._aot_autograd.utils import strict_zip, _get_symint_hints, KNOWN_TYPES, partial_flatten_asdict, normalize_as_list, _get_autocast_states, make_boxed_func, make_boxed_compiler, call_func_at_runtime_with_args, create_tree_flattened_fn, maybe_to_fresh_input
  from ._aot_autograd.logging_utils import graph_being_compiled, nth_graph, model_name, set_model_name, get_aot_compilation_context, get_aot_graph_name, get_graph_being_compiled, track_graph_compiling, callback_set, setup_stacktrace_preservation_hooks, describe_input, format_guard_bug_msg
  from ._aot_autograd.functional_utils import is_fun, to_fun, from_fun, sync_functional_tensor, has_metadata_mutation, has_data_mutation, are_all_mutations_hidden_from_autograd, are_all_mutations_under_no_grad_or_inference_mode, gen_alias_from_base, assert_functional_graph, _get_mutation_type, _check_if_mutation_can_be_in_graph
  from ._aot_autograd.schemas import OutputType, OutputAliasInfo, MutationType, InputAliasInfo, SubclassCreationMeta, ViewAndMutationMeta, SubclassMeta, TensorAlias, BackwardSignature, GraphOutputName, GraphInputName, FQN, GraphSignature, AOTConfig
  from ._aot_autograd.subclass_utils import requires_subclass_dispatch, unwrap_tensor_subclasses, wrap_tensor_subclasses, wrap_tensor_subclasses_maybe_joint, create_metadata_for_subclass
  from ._aot_autograd.collect_metadata_analysis import run_functionalized_fw_and_collect_metadata
  from ._aot_autograd.input_output_analysis import remove_dupe_metadata, create_synthetic_base_metadata, _tensors_definitely_do_not_overlap, compute_overlapping_inputs, create_graph_signature
  from ._aot_autograd.traced_function_transforms import fn_input_mutations_to_outputs, fn_prepped_for_autograd, create_functionalized_fn, create_functionalized_rng_ops_wrapper, aot_dispatch_subclass, create_functional_call, create_joint
  from ._aot_autograd.runtime_wrappers import create_runtime_wrapper, functionalized_rng_runtime_epilogue, aot_dispatch_subclass_wrapper, aot_wrapper_dedupe, aot_wrapper_synthetic_base, merge_view_inputs
  from ._aot_autograd.dispatch_and_compile_graph import aot_dispatch_base_graph, aot_dispatch_autograd_graph
  from ._aot_autograd.jit_compile_runtime_wrappers import aot_dispatch_base, aot_dispatch_autograd
label_desc: Use image node_id 1 for calling a global function with example usage: create_aot_dispatcher_function(flat_fn, flat_args, aot_config) and returns: compiled_fn, unknown, compiled_fn,  fw_metadata, fake_mode, x, x, fake_mode, shape_env, x, x