code: def create_functionalized_fn(
      fn,
      args,
      *,
      meta: ViewAndMutationMeta,
      aot_config: AOTConfig,
      trace_joint: bool,
  ) -> Any:
      @wraps(fn)
      def _functionalized_f_helper(*args):
          # See Note [Disabling Functionalize TLS Above Python Functionalization]
          disable_above = torch._C._ExcludeDispatchKeyGuard(
              torch._C.DispatchKeySet(
                  torch._C.DispatchKey.Functionalize
              )
          )
          with disable_above, FunctionalTensorMode(
              aot_config.pre_dispatch
          ):
              # Wrap inputs into functional wrappers
              f_args = pytree.tree_map(to_fun, args)
              # Run the joint
              f_outs = fn(*f_args)
  
          if trace_joint:
              # We support a limited amount of mutation of graph inputs during the backward pass.
              # (This is used e.g. by Float8, which needs to update buffers during the backward pass)
              # Here, we perform extra checks for primals that were mutated in the **backward**
              # We're doing the checks here instead of doing them with the rest of the input mutation handling because:
              # - We need to detect inputs that were mutated in the backward **separately** from mutations that happened
              #   during the forward, because the handling is different: some input mutations from the the forward
              #   can be only handled in a fw-only runtime epilogue, and in theory if we wanted to handle those same
              #   types of mutations in the backward we would need a bw-only runtime epilogue.
              # - We could in theory have our analysis pass differentiate mutations in the fw from mutations in
              #   the bw by running our analysis first on the fw-only graph, and then on the joint graph. This would
              #   require an extra round of tracing though, so it's more efficient to do in-line here.
              assert (
                  isinstance(args, tuple)
                  and len(args) == 2
                  and isinstance(args[0], (list, tuple))
              )
              # Only look at mutations that happened to forward inputs (e.g. fw buffers that were saved for bw)
              primals_before = args[0]
              primals_after = pytree.tree_map(from_fun, f_args[0])
              for f_inpt, before, after, inpt_info in zip(
                  f_args[0],
                  primals_before,
                  primals_after,
                  meta.input_info,
              ):
                  # Ban metadata mutations on fw inputs during the bw
                  if not inpt_info.mutates_metadata:
                      assert not has_metadata_mutation(
                          f_inpt,
                          before,
                          check_only_storage_mutation=False,
                      ), "Found a graph input that had its metadata mutated in the backward. This is not supported"
                  # Allow data mutations on fw inputs during the bw, but only if they do not require grad
                  # So we can guarantee that we can keep the mutations in the graph
                  if (
                      has_data_mutation(f_inpt)
                      and not inpt_info.mutates_data
                  ):
                      assert (
                          not inpt_info.requires_grad
                      ), "Found a graph input that requires_grad and was mutated in the backward. This is not supported"
                      # Otherwise, put the mutation in the graph
                      before.copy_(after)
              # Now that we covered mutations to *forward* inputs during the backward,
              # we also need to cover mutations to *backward-only* inputs during the backward (e.g. mutation to a grad_out).
              # Today, we will just error in all cases of this happening unless someone needs us to support it.
              tangents_before = args[1]
              tangents_after = pytree.tree_map(from_fun, f_args[1])
              for f_inpt, before, after in zip(
                  f_args[1], tangents_before, tangents_after
              ):
                  assert not has_metadata_mutation(
                      f_inpt, before, check_only_storage_mutation=False
                  ) and not has_data_mutation(
                      f_inpt
                  ), "Found an input to the backward that was mutated during the backward pass. This is not supported"
  
          if aot_config.keep_inference_input_mutations:
              # Note: This is a bit annoying. There's a layering issue here, where:
              # (1) functionalization needs to operate on **synthetic base** inputs, before unpacking them into the "real" inputs.
              # (2) For keep_input_mutations, we support tracing a call to copy_() directly on mutated inputs.
              #     However, we **only** want to support this for inputs that have data-only (and no metadata) mutations,
              #     because inductor (and backends in generally) would prefer not to see these (e.g. as_strided_(), resize_()).
              #     This makes it pretty difficult for this logic to operate on synthetic bases.
              # (3) In addition, there are cases where it's significantly cheaper to perform the copy on the individual
              #     (unpacked) input aliases, instead of the synthetic base.
              # Example case where (3) could be important:
              #
              #     def f(x, y):
              #         x.mul_(2)
              #         y.mul_(3)
              #         return x, y
              #    a = torch.ones(1'000'000)
              #    x, y = out(a[0:9], a[1:10])
              #
              # It would be much better to add copy_() calls into the graph for the two tiny slices, instead of materializing
              # a giant "updated synthetic base" and copying into a's entire storage.
              #
              # For now, we are pessimistically not performing the optimization from (3);
              # we will materialize an "updated" synthetic base, and copy it back to the synthetic input base.
              # This allows us to factor aot autograd much more nicely, since only one area of the code needs to worry
              # about synthetic bases.
              for i, (inpt_old, inpt_f) in enumerate(
                  zip(args, f_args)
                  if not trace_joint
                  else zip(args[0], f_args[0])
              ):
                  if not isinstance(inpt_f, torch.Tensor):
                      continue
                  assert is_fun(inpt_f)
                  inpt_new = from_fun(inpt_f)
                  if (
                      meta.input_info[i].mutation_type
                      == MutationType.MUTATED_IN_GRAPH
                  ):
                      # We found an input that had a (data-only) mutation.
                      # Since keep_input_mutations is set, we need to faithfully apply a copy_()
                      # so the compiler will see the input mutation in the graph.
                      if meta.input_info[
                          i
                      ].mutations_hidden_from_autograd:
                          # Hidden from autograd = run under no_grad, **and** don't bump VC
                          with torch.no_grad(), torch.autograd._unsafe_preserve_version_counter(
                              inpt_old
                          ):
                              inpt_old.copy_(inpt_new)
                      elif meta.input_info[
                          i
                      ].mutations_under_no_grad_or_inference_mode:
                          # Under no_grad = run under no_grad (we still bump the VC though)
                          # (inference_mode will also bump the VC, as long as the tensor in question
                          # was created outside of inference_mode)
                          with torch.no_grad():
                              inpt_old.copy_(inpt_new)
                      else:
                          inpt_old.copy_(inpt_new)
  
          return pytree.tree_map(from_fun, f_outs)
  
      # Kinda annoying, but needed to make sure that the fx graph we trace out has "primals"
      # and "tangents" as its input names (which are special-cased by the partitioner)
      # TODO (tmanlaibaatar) revisit this if we ever need to turn on non-strict joint graph export
      def joint_helper(primals, tangents):
          return _functionalized_f_helper(primals, tangents)
  
      helper = joint_helper if trace_joint else _functionalized_f_helper
      if config.functionalize_rng_ops:
          # Setup the wrapper for functionalization of rng ops
          helper, args = create_functionalized_rng_ops_wrapper(
              helper, args, trace_joint
          )
  
      return helper, args
  
imports: import warnings
  from contextlib import nullcontext
  from functools import wraps
  from typing import Any, Callable, List, Tuple, Union
  from unittest.mock import patch
  import torch
  import torch.fx.traceback
  import torch.utils._pytree
  from torch import Tensor
  from torch._decomp.decompositions_for_rng import PhiloxStateTracker
  from torch._guards import detect_fake_mode
  from torch._prims_common import CUDARngStateHelper
  from torch._subclasses.functional_tensor import FunctionalTensorMode
  from torch.fx import Interpreter
  from torch.fx.experimental.symbolic_shapes import definitely_false, sym_eq
  from torch.nn.utils import stateless
  from ..None import config
  from .collect_metadata_analysis import run_functionalized_fw_and_collect_metadata
  from .functional_utils import from_fun, has_data_mutation, has_metadata_mutation, is_fun, sync_functional_tensor, to_fun
  from .logging_utils import setup_stacktrace_preservation_hooks
  from .schemas import AOTConfig, MutationType, OutputType, SubclassMeta, SubclassTracingInfo, ViewAndMutationMeta
  from .subclass_utils import create_subclass_meta, requires_subclass_dispatch, unwrap_tensor_subclasses, wrap_tensor_subclasses_maybe_joint
  from .utils import maybe_to_fresh_input
label_desc: Use image node_id 5 for calling a global function with example usage: create_functionalized_fn(fn, args) and returns: helper,  args, pytree, _functionalized_f_helper