code: def do_auto_functionalize(
      op: torch._ops.OpOverload,
      args: Tuple[Any, ...],
      kwargs: Dict[str, Any],
  ) -> Any:
      """Functionalizes a call to op(*args, **kwargs) by emitting a call to
      `outs = auto_functionalized(op, mutated_args_names, normalized_kwargs)`
      and replacing the mutated (args, kwargs) with the corresponding outputs.
  
      The normalized_kwargs are just the (args, kwargs), but all in kwarg form.
      This makes handling easier for the auto_functionalized HOP.
      """
      from torch._subclasses.functional_tensor import (
          PythonFunctionalizeAPI,
      )
  
      ctx = PythonFunctionalizeAPI()
  
      # List of the name of args that get mutated (according to the schema)
      mutable_args_names = []
      # All of the (args, kwargs), but all as kwargs. The names for the
      # args come from the schema. This makes it easier for us to work with them.
      normalized_kwargs = {}
      schema = op._schema
      for idx, arg in enumerate(schema.arguments):
          if arg.alias_info is not None and arg.alias_info.is_write:
              mutable_args_names.append(arg.name)
          # NB: torch_dispatch kwargs are the args defined as kwarg-only in the schema
          if arg.name in kwargs:
              normalized_kwargs[arg.name] = kwargs[arg.name]
          else:
              normalized_kwargs[arg.name] = args[idx]
  
      unwrapped_kwargs = ctx.unwrap_tensors(normalized_kwargs)  # type: ignore[arg-type]
      with ctx.redispatch_to_next():
          unwrapped_actual_out, unwrapped_outs = auto_functionalized(
              op, mutable_args_names, unwrapped_kwargs  # type: ignore[arg-type]
          )
      assert len(unwrapped_outs) == len(mutable_args_names)
      for name, unwrapped_out in zip(
          mutable_args_names, unwrapped_outs
      ):
          # Can be None if input was `Tensor(a!)?`
          if unwrapped_out is None:
              continue
          assert isinstance(unwrapped_out, torch.Tensor)
          orig_arg = normalized_kwargs[name]
          ctx.replace(orig_arg, unwrapped_out)
          ctx.commit_update(orig_arg)
          ctx.sync(orig_arg)
  
      return ctx.wrap_tensors(unwrapped_actual_out)
  
file_path: pytorch/torch/_higher_order_ops/auto_functionalize.py
imports: from typing import Any, Dict, List, Tuple
  import torch
  import torch.utils._pytree
  from torch import Tensor
  from torch._C import DispatchKey
  from torch._ops import HigherOrderOperator
  from torch._prims_common import clone_preserve_strides
  from torch._subclasses.fake_tensor import FakeTensorMode
  from torch.fx.experimental.proxy_tensor import disable_proxy_modes_tracing, ProxyTorchDispatchMode, track_tensor_tree
label_desc: Use image node_id 5 for calling a global function with example usage: do_auto_functionalize(op, args, kwargs) and returns: ctx