code: def aot_export_joint_simple(
      func: Callable,
      args,
      *,
      trace_joint: bool,
      # It looks like the main consequence of this API is that for dynamic shapes,
      # it will assume that parms/buffers are static.
      # With the new inferred dynamic shapes API, maybe this doesn't matter?
      num_params_buffers: int = 0,
      decompositions: Optional[Dict] = None,
  ) -> torch.fx.GraphModule:
      """
      A simplified version of export. Used by higher order operators.
  
      This function makes a high-level "no calling convention changes" guarantee:
      - If no inputs require grad (so we export an inference graph),
        there are *no* calling convention change between the exported graph, and "func".
      - If at least one input requires grad (so we trace out and export a joint fw-bw graph),
        Then if you were partition the graph into a separate forward and backward graph,
        The forward graph will have no calling convention changes compared to "func".
  
      The above also relies on some strong restrictions around which functions this API accepts:
      (1) `args` cannot contain any pytrees (they must have been pytree_flattened already)
      (2) `func` cannot mutate any inputs
      (3) The outputs of `func` cannot alias any inputs.
  
      Note: this function is only lightly tested today. It will probably be tested more heavily by higher order ops.
      """
      if trace_joint:
          ctx = nullcontext
      else:
          # Run under no_grad, so our tracing machinery only traces an inference graph.
          ctx = torch.no_grad
  
      with ctx():
          fx_g, metadata, in_spec, out_spec = _aot_export_function(
              func,
              args,
              decompositions=decompositions,
          )
      # At this point, we can just directly return the (joint or inference graph) that we traced.
      # First though: a bunch of assertions to make sure that our graph doesn't require
      # any calling convention changes compared to the original function.
      # These restrictions are *in addition to* the general restrictions on export.
  
      # No input mutations
      if (
          len(
              [
                  x
                  for x in metadata.input_info
                  if x.mutates_data or x.mutates_metadata
              ]
          )
          != 0
      ):
          raise RuntimeError(
              f"aot_export_joint_simple does not support input mutations. {str(metadata)}"
          )
      # No output aliasing
      if (
          len(
              [
                  x
                  for x in metadata.output_info
                  if x.output_type != OutputType.non_alias
              ]
          )
          != 0
      ):
          raise RuntimeError(
              f"aot_export_joint_simple does not support outputs that alias inputs. {str(metadata)}"
          )
      # No pytrees
      if type(in_spec) == pytree.LeafSpec:
          raise RuntimeError(
              f"aot_export_joint_simple requires inputs to be a single list/tuple. in_spec={str(in_spec)}"
          )
      if (
          len(
              [
                  x
                  for x in in_spec.children_specs
                  if type(x) != pytree.LeafSpec
              ]
          )
          != 0
      ):
          raise RuntimeError(
              f"aot_export_joint_simple requires individual inputs not to be pytrees. in_spec={str(in_spec)}"
          )
      if type(out_spec) == pytree.LeafSpec:
          raise RuntimeError(
              f"aot_export_joint_simple requires outputs to be a single list/tuple. out_spec={str(out_spec)}"
          )
      if (
          len(
              [
                  x
                  for x in out_spec.children_specs
                  if type(x) != pytree.LeafSpec
              ]
          )
          != 0
      ):
          raise RuntimeError(
              f"aot_export_joint_simple requires individual outputs not to be pytrees. out_spec={str(out_spec)}"
          )
      # TODO: we might have to temporarily patch config.functionalize_rng
      # so that it doesn't run when we're exporting a higher order op.
  
      if config.debug_assert:
          # Smoke test that after partitioning, we can run the forward without any calling convention changes.
          (
              fw_module,
              bw_module,
          ) = aot_config.default_partition(  # noqa: F821
              fx_g,
              args,
              num_fwd_outputs=len(
                  fw_metadata.output_infos
              ),  # noqa: F821
          )
          # Attempt to run the fw_module with the original user inputs
          fake_mode = detect_fake_mode(args)
          if fake_mode is None:
              fake_mode = FakeTensorMode()
          with fake_mode:
              fw_module(*args)
      return fx_g
  
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 6 for calling a global function with example usage: aot_export_joint_simple(func, args) and returns: fx_g