code: def fn_prepped_for_autograd(
      fn: Callable,
      meta: ViewAndMutationMeta,
  ) -> Any:
      @wraps(fn)
      def inner_fn(*args):
          args_maybe_cloned = [
              maybe_to_fresh_input(i, t, meta)
              for i, t in enumerate(args)
          ]
  
          outs = fn(*args_maybe_cloned)
          assert isinstance(outs, (tuple, list))
          outs = list(outs)
          assert len(meta.output_info) == len(outs)
  
          mutated_inputs_to_return = [
              x
              for (i, x) in enumerate(args_maybe_cloned)
              if i in meta.mutated_inp_runtime_indices
          ]
  
          intermediate_bases = []
          for i, (o, info) in enumerate(zip(outs, meta.output_info)):
              if (
                  info.output_type
                  == OutputType.alias_of_intermediate_save_as_output
              ):
                  intermediate_bases.append(o._base)
  
          assert meta.num_intermediate_bases == len(intermediate_bases)
  
          # the compiled forward should return (mutated_inputs, user_outs, intermediate_bases)
          fw_outs_to_return = (
              *mutated_inputs_to_return,
              *outs,
              *intermediate_bases,
          )
  
          # Also return a boolean mask specifying which outputs to this function will be used as tangents
          mutated_inputs_grad_mask = [
              meta.input_info[
                  meta.mutated_inp_runtime_indices[i]
              ].mutates_data
              and meta.input_info[
                  meta.mutated_inp_runtime_indices[i]
              ].requires_grad
              for (i, x) in enumerate(mutated_inputs_to_return)
          ]
  
          # Pass any (non-aliased) outputs in as tangents, since they'll be returned as outputs in the fw
          # For outputs that are aliases of intermediates, we will have returned the output's _base as an output in the graph instead,
          # which we *should* send to grad()
          output_grad_mask = [
              meta.output_info[i].output_type
              in [
                  OutputType.non_alias,
                  OutputType.unsafe_view_alias,
                  OutputType.custom_function_view,
              ]
              # Also, only tensor outputs should participate in the backward
              # (in particular, Symint outputs in the forward graph shouldn't get tangents)
              and issubclass(meta.output_info[i].raw_type, Tensor)
              and meta.output_info[i].requires_grad
              for (i, x) in enumerate(outs)
          ]
  
          intermediate_base_grad_mask = [
              True for _ in range(len(intermediate_bases))
          ]
  
          out_grad_mask = (
              mutated_inputs_grad_mask
              + output_grad_mask
              + intermediate_base_grad_mask
          )
          assert len(out_grad_mask) == len(fw_outs_to_return)
  
          # Take care to grab and sync the updated inputs from primals_after_cloning (the inputs we actually mutate!)
          # and not primals (the preserved inputs, pre-mutation, that we pass to grad())
          # This is annoying: our joint function needs to be aware of functionalization
          # (syncing mutated inputs before calling autograd.grad())
          # In theory, we could make the autograd engine do this automatically, although that probably isn't any cleaner.
          for arg in args_maybe_cloned:
              if not isinstance(arg, Tensor):
                  continue
              sync_functional_tensor(arg)
  
          return fw_outs_to_return, out_grad_mask
  
      return inner_fn
  
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 2 for calling a global function with example usage: fn_prepped_for_autograd(fn, meta) and returns: inner_fn, fw_outs_to_return,  out_grad_mask