{"kind":"task","effective_mode":"full","benchmark":{"kind":"benchmark","effective_mode":"full","slug":"longbench-v2","formal_name":"LongBench v2","introduction":"LongBench v2 evaluates deep understanding and reasoning over long contexts through multiple-choice questions. Its official description lists 503 questions spanning tasks such as single-document and multi-document QA and code-repository understanding.","introduction_ja":"","introduction_en":"","category":"Category not supplied","task_count":null,"acquisition_status":"Acquisition status not supplied","official_url":"https://huggingface.co/datasets/zai-org/LongBench-v2","indexing_mode":"noindex","profile":{"resources":[],"task_format":"","scoring":"","metric":"","size":"","answer_access":"","license":"","citation":"","maintainer":"","released":"","why_hard":"","related":[]}},"task_id":"e56f2733-f03e-5f9c-87e6-4004a7f8c104","task_key":"train--66fa542bbb02136c067c686d","task_revision_id":"3","upstream_id":"66fa542bbb02136c067c686d","short_description":"In the FileManager class, which of the following wrongly describes the purpose…","config":"","split":"train","body":"{\"choice_A\":\"Template Substitution Handling: The write_with_template method calls substitute_with_template to replace patterns in the template file using only the dictionary. The result is then written to the file only if dry_run is False.\",\"choice_B\":\"Duplicate File Write Prevention: The write_with_template method maintains a set called filenames that tracks written files. It raises an assertion error if the same file is attempted to be written again within the same execution.\",\"choice_C\":\"File Writing if Contents Changed: The _write_if_changed method checks if the contents of a file have changed by comparing them with the new contents. If they differ, it overwrites the file, ensuring only modified files are rewritten.\",\"choice_D\":\"Generated Comment Insertion: When the env_callable returns a dictionary, the substitute_with_template method checks if the key generated_comment exists. If it does not, it creates one using the generator path and the template filename. This is added to the dictionary before substitution.\",\"context\":\"from __future__ import annotations\\n\\nimport argparse\\nimport os\\nimport re\\nfrom collections import Counter, defaultdict, namedtuple\\nfrom pathlib import Path\\nfrom typing import Sequence\\n\\nimport yaml\\n\\nimport torchgen.api.dispatcher as dispatcher\\nimport torchgen.dest as dest\\nfrom torchgen.api.types import DispatcherSignature\\nfrom torchgen.code_template import CodeTemplate\\nfrom torchgen.context import native_function_manager\\nfrom torchgen.gen import get_grouped_native_functions, parse_native_yaml\\nfrom torchgen.model import (\\n BackendIndex,\\n BackendMetadata,\\n DispatchKey,\\n NativeFunction,\\n NativeFunctionsGroup,\\n OperatorName,\\n)\\nfrom torchgen.selective_build.selector import SelectiveBuilder\\nfrom torchgen.utils import concatMap, context, FileManager, NamespaceHelper, Target\\nfrom torchgen.yaml_utils import YamlLoader\\n\\n\\n# Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key.\\n# Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping)\\nParsedExternalYaml = namedtuple(\\n \\\"ParsedExternalYaml\\\",\\n [\\\"backend_key\\\", \\\"autograd_key\\\", \\\"class_name\\\", \\\"cpp_namespace\\\", \\\"backend_indices\\\"],\\n)\\n\\n\\ndef parse_backend_yaml(\\n backend_yaml_path: str,\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n backend_indices: dict[DispatchKey, BackendIndex],\\n) -> ParsedExternalYaml:\\n native_functions_map: dict[OperatorName, NativeFunction] = {\\n f.func.name: f\\n for f in concatMap(\\n lambda f: [f] if isinstance(f, NativeFunction) else list(f.functions()),\\n grouped_native_functions,\\n )\\n }\\n\\n with open(backend_yaml_path) as f:\\n yaml_values = yaml.load(f, Loader=YamlLoader)\\n assert isinstance(yaml_values, dict)\\n\\n valid_keys = [\\n \\\"backend\\\",\\n \\\"class_name\\\",\\n \\\"cpp_namespace\\\",\\n \\\"extra_headers\\\",\\n \\\"supported\\\",\\n \\\"autograd\\\",\\n \\\"full_codegen\\\",\\n \\\"non_native\\\",\\n \\\"ir_gen\\\",\\n \\\"symint\\\",\\n ]\\n\\n backend = yaml_values.pop(\\\"backend\\\", None)\\n assert backend is not None, 'You must provide a value for \\\"backend\\\"'\\n\\n class_name = yaml_values.pop(\\\"class_name\\\", None)\\n\\n cpp_namespace = yaml_values.pop(\\\"cpp_namespace\\\", None)\\n assert cpp_namespace is not None, 'You must provide a value for \\\"cpp_namespace\\\"'\\n\\n # Mostly just defaulting to false to stick with LazyTensor convention.\\n use_out_as_primary = yaml_values.pop(\\\"use_out_as_primary\\\", False)\\n assert isinstance(\\n use_out_as_primary, bool\\n ), f\\\"You must provide either True or False for use_out_as_primary. Provided: {use_out_as_primary}\\\"\\n\\n use_device_guard = yaml_values.pop(\\\"device_guard\\\", False)\\n assert isinstance(\\n use_device_guard, bool\\n ), f\\\"You must provide either True or False for device_guard. Provided: {use_device_guard}\\\"\\n\\n supported = yaml_values.pop(\\\"supported\\\", [])\\n if supported is None:\\n supported = [] # Allow an empty list of supported ops\\n assert isinstance(\\n supported, list\\n ), f'expected \\\"supported\\\" to be a list, but got: {supported} (of type {type(supported)})'\\n\\n symint = yaml_values.pop(\\\"symint\\\", [])\\n if symint is None:\\n symint = [] # Allow an empty list of symint ops\\n assert isinstance(\\n symint, list\\n ), f'expected \\\"symint\\\" to be a list, but got: {supported} (of type {type(supported)})'\\n symint_set = set(symint)\\n\\n supported_autograd = yaml_values.pop(\\\"autograd\\\", [])\\n assert isinstance(\\n supported_autograd, list\\n ), f'expected \\\"autograd\\\" to be a list, but got: {supported_autograd}'\\n\\n # full_codegen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py\\n full_codegen = yaml_values.pop(\\\"full_codegen\\\", [])\\n supported.extend(full_codegen)\\n\\n # non_native is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py\\n yaml_values.pop(\\\"non_native\\\", {})\\n\\n # ir_gen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py\\n yaml_values.pop(\\\"ir_gen\\\", {})\\n\\n assert (\\n len(yaml_values.keys()) == 0\\n ), f'{backend_yaml_path} contains unexpected keys: {\\\", \\\".join(yaml_values.keys())}. \\\\\\nOnly the following keys are supported: {\\\", \\\".join(valid_keys)}'\\n\\n def create_backend_index(\\n backend_ops: list[str],\\n symint_ops: set[str],\\n dispatch_key: DispatchKey,\\n *,\\n use_out_as_primary: bool,\\n use_device_guard: bool,\\n ) -> BackendIndex:\\n metadata: dict[OperatorName, BackendMetadata] = {}\\n for op in backend_ops:\\n op_name = OperatorName.parse(op)\\n assert (\\n op_name in native_functions_map\\n ), f\\\"Found an invalid operator name: {op_name}\\\"\\n # See Note [External Backends Follow Dispatcher API]\\n kernel_name = dispatcher.name(native_functions_map[op_name].func)\\n if op in symint_ops:\\n kernel_name += \\\"_symint\\\"\\n # TODO: allow structured external backends later.\\n m = BackendMetadata(\\n kernel=kernel_name, structured=False, cpp_namespace=cpp_namespace\\n )\\n metadata[op_name] = m\\n return BackendIndex(\\n dispatch_key=dispatch_key,\\n use_out_as_primary=use_out_as_primary,\\n external=True,\\n device_guard=use_device_guard,\\n index=metadata,\\n )\\n\\n backend_key: DispatchKey | None = None\\n if len(supported) > 0:\\n with context(\\n lambda: f'The provided value for \\\"backend\\\" must be a valid DispatchKey, but got {backend}.'\\n ):\\n backend_key = DispatchKey.parse(backend)\\n\\n backend_idx = create_backend_index(\\n supported,\\n symint_set,\\n backend_key,\\n use_out_as_primary=use_out_as_primary,\\n use_device_guard=use_device_guard,\\n )\\n assert backend_key not in backend_indices\\n backend_indices[backend_key] = backend_idx\\n\\n autograd_key: DispatchKey | None = None\\n if len(supported_autograd) > 0:\\n with context(\\n lambda: f'The \\\"autograd\\\" key was specified, which indicates that you would like to override \\\\\\nthe behavior of autograd for some operators on your backend. However \\\"Autograd{backend}\\\" is not a valid DispatchKey.'\\n ):\\n autograd_key = DispatchKey.parse(f\\\"Autograd{backend}\\\")\\n\\n autograd_idx = create_backend_index(\\n supported_autograd,\\n symint_set,\\n autograd_key,\\n use_out_as_primary=use_out_as_primary,\\n use_device_guard=use_device_guard,\\n )\\n assert autograd_key not in backend_indices\\n backend_indices[autograd_key] = autograd_idx\\n\\n for g in grouped_native_functions:\\n if isinstance(g, NativeFunction):\\n forward_kernels = (\\n []\\n if backend_key is None\\n else [\\n m\\n for m in [backend_indices[backend_key].get_kernel(g)]\\n if m is not None\\n ]\\n )\\n backward_kernels = (\\n []\\n if autograd_key is None\\n else [\\n m\\n for m in [backend_indices[autograd_key].get_kernel(g)]\\n if m is not None\\n ]\\n )\\n else:\\n forward_kernels = (\\n []\\n if backend_key is None\\n else [\\n m\\n for m in [\\n backend_indices[backend_key].get_kernel(f)\\n for f in g.functions()\\n ]\\n if m is not None\\n ]\\n )\\n backward_kernels = (\\n []\\n if autograd_key is None\\n else [\\n m\\n for m in [\\n backend_indices[autograd_key].get_kernel(f)\\n for f in g.functions()\\n ]\\n if m is not None\\n ]\\n )\\n\\n forward_kernels = [f for f in forward_kernels if f is not None]\\n backward_kernels = [f for f in backward_kernels if f is not None]\\n assert (\\n len(forward_kernels) == 0 or len(backward_kernels) == 0\\n ), f'Currently, all variants of an op must either be registered to a backend key, or to a backend\\\\'s \\\\\\nautograd key. They cannot be mix and matched. If this is something you need, feel free to create an issue! \\\\\\n{forward_kernels[0].kernel} is listed under \\\"supported\\\", but {backward_kernels[0].kernel} is listed under \\\"autograd\\\".'\\n\\n return ParsedExternalYaml(\\n backend_key, autograd_key, class_name, cpp_namespace, backend_indices\\n )\\n\\n\\ndef error_on_missing_kernels(\\n native_functions: Sequence[NativeFunction],\\n backend_indices: dict[DispatchKey, BackendIndex],\\n backend_key: DispatchKey,\\n autograd_key: DispatchKey | None,\\n class_name: str,\\n kernel_defn_file_path: str,\\n full_codegen: list[OperatorName] | None = None,\\n) -> None:\\n try:\\n with open(kernel_defn_file_path) as f:\\n backend_defns = f.read()\\n except OSError as e:\\n raise AssertionError(\\n f\\\"Unable to read from the specified impl_path file: {kernel_defn_file_path}\\\"\\n ) from e\\n\\n if full_codegen is None:\\n full_codegen = []\\n\\n indices = [backend_indices[backend_key].index] + (\\n [] if autograd_key is None else [backend_indices[autograd_key].index]\\n )\\n # Quick mapping from each OperatorName used by the external backend\\n # to its backend kernel name\\n expected_backend_op_names: dict[OperatorName, str] = dict(\\n list(\\n concatMap(\\n lambda index: [\\n (op_name, metadata.kernel) for op_name, metadata in index.items()\\n ],\\n indices,\\n )\\n )\\n )\\n expected_backend_native_funcs: list[NativeFunction] = [\\n f\\n for f in native_functions\\n if f.func.name in expected_backend_op_names.keys()\\n and f.func.name not in full_codegen\\n ]\\n expected_backend_kernel_name_counts: dict[str, list[NativeFunction]] = defaultdict(\\n list\\n )\\n for native_f in expected_backend_native_funcs:\\n expected_backend_kernel_name_counts[\\n expected_backend_op_names[native_f.func.name]\\n ].append(native_f)\\n\\n # This just looks for lines containing \\\"foo(\\\", and assumes that the kernel foo has been implemented.\\n # It might cause false negatives (we won't catch all cases), but that's ok - if we catch a missing kernel\\n # here, then we get a nicer error message. If we miss it, you get a linker error.\\n kernel_defn_regex = rf\\\"(.*){class_name}::\\\\s*([\\\\w\\\\d]*)\\\\(\\\"\\n actual_backend_kernel_name_counts = Counter(\\n # A bit unwieldy (this could probably be moved into regex),\\n # but we don't want to include kernel names that come from function calls,\\n # like \\\"return torch_xla::XLANativeFunctions::empty_strided_symint(...)\\\".\\n # Easy check is to ignore any lines with colons before the class name.\\n [\\n y\\n for (x, y) in re.findall(kernel_defn_regex, backend_defns)\\n if not x.endswith(\\\":\\\")\\n ]\\n )\\n\\n missing_kernels_err_msg = \\\"\\\"\\n for expected_name, funcs in expected_backend_kernel_name_counts.items():\\n expected_overload_count = len(funcs)\\n actual_overload_count = actual_backend_kernel_name_counts[expected_name]\\n if expected_overload_count != actual_overload_count:\\n\\n def create_decl(f: NativeFunction) -> str:\\n with native_function_manager(f):\\n return DispatcherSignature.from_schema(f.func).decl()\\n\\n expected_schemas_str = \\\"\\\\n\\\".join([create_decl(f) for f in funcs])\\n missing_kernels_err_msg += f\\\"\\\"\\\"\\n{class_name} is missing a kernel definition for {expected_name}. We found {actual_overload_count} kernel(s) with that name,\\nbut expected {expected_overload_count} kernel(s). The expected function schemas for the missing operator are:\\n{expected_schemas_str}\\n\\n\\\"\\\"\\\"\\n assert missing_kernels_err_msg == \\\"\\\", missing_kernels_err_msg\\n\\n\\ndef main() -> None:\\n parser = argparse.ArgumentParser(description=\\\"Generate backend stub files\\\")\\n parser.add_argument(\\n \\\"-s\\\",\\n \\\"--source-yaml\\\",\\n \\\"--source_yaml\\\",\\n help=\\\"path to source yaml file containing operator external definitions\\\",\\n )\\n parser.add_argument(\\\"-o\\\", \\\"--output-dir\\\", \\\"--output_dir\\\", help=\\\"output directory\\\")\\n parser.add_argument(\\n \\\"--dry-run\\\", \\\"--dry_run\\\", type=bool, default=False, help=\\\"output directory\\\"\\n )\\n parser.add_argument(\\n \\\"--impl-path\\\",\\n \\\"--impl_path\\\",\\n type=str,\\n default=None,\\n help=\\\"path to the source C++ file containing kernel definitions\\\",\\n )\\n options = parser.parse_args()\\n\\n run(options.source_yaml, options.output_dir, options.dry_run, options.impl_path)\\n\\n\\ndef gen_dispatchkey_nativefunc_headers(\\n fm: FileManager,\\n class_name: str,\\n cpp_namespace: str,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n backend_dispatch_key: DispatchKey,\\n autograd_dispatch_key: DispatchKey | None,\\n backend_name: str = \\\"\\\",\\n) -> None:\\n assert class_name is not None\\n generated_comment = (\\n \\\"Autogenerated file by gen_backend_stubs.py. Do not edit directly!\\\"\\n )\\n\\n # Convert to a set first to remove duplicate kernel names.\\n # Backends are allowed to repeat kernel names; only generate the declaration once!\\n # Sort for deterministic output.\\n backend_declarations = sorted(\\n set(\\n concatMap(\\n lambda f: dest.compute_native_function_declaration(\\n f, backend_indices[backend_dispatch_key]\\n ),\\n grouped_native_functions,\\n )\\n )\\n )\\n autograd_declarations = sorted(\\n set(\\n concatMap(\\n lambda f: []\\n if autograd_dispatch_key is None\\n else dest.compute_native_function_declaration(\\n f, backend_indices[autograd_dispatch_key]\\n ),\\n grouped_native_functions,\\n )\\n )\\n )\\n\\n ns_helper = NamespaceHelper(cpp_namespace)\\n fm.write_with_template(\\n f\\\"{backend_dispatch_key}NativeFunctions.h\\\",\\n \\\"DispatchKeyNativeFunctions.h\\\",\\n lambda: {\\n \\\"generated_comment\\\": generated_comment,\\n \\\"namespace_prologue\\\": ns_helper.prologue,\\n \\\"class_name\\\": class_name,\\n \\\"namespace_epilogue\\\": ns_helper.epilogue,\\n \\\"dispatch_declarations\\\": backend_declarations + autograd_declarations,\\n \\\"BackendName\\\": backend_name,\\n \\\"DispatchKey\\\": backend_dispatch_key,\\n },\\n )\\n\\n\\ndef gen_dispatcher_registrations(\\n fm: FileManager,\\n output_dir: str,\\n class_name: str,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n backend_dispatch_key: DispatchKey,\\n dispatch_key: DispatchKey,\\n selector: SelectiveBuilder,\\n # build_in_tree is true for lazy TS backend and affects include paths, not used for external backends\\n build_in_tree: bool = False,\\n per_operator_headers: bool = False,\\n backend_name: str = \\\"\\\",\\n eager_registration: bool = True,\\n) -> None:\\n headers = [\\n f\\\"{output_dir}/{backend_dispatch_key}NativeFunctions.h\\\",\\n ]\\n if build_in_tree:\\n external_backend_headers_str = \\\"\\\\n\\\".join(f\\\"#include <{h}>\\\" for h in headers)\\n else:\\n external_backend_headers_str = \\\"\\\\n\\\".join(f'#include \\\"{h}\\\"' for h in headers)\\n\\n assert class_name is not None\\n backend_index = backend_indices[dispatch_key]\\n\\n dispatch_registrations_body = list(\\n concatMap(\\n dest.RegisterDispatchKey(\\n backend_index,\\n Target.REGISTRATION,\\n selector,\\n rocm=False,\\n symint=True,\\n class_method_name=f\\\"{class_name}\\\",\\n skip_dispatcher_op_registration=False,\\n ),\\n grouped_native_functions,\\n )\\n )\\n newline = \\\"\\\\n\\\"\\n ns_helper = NamespaceHelper(namespace_str=\\\"at\\\")\\n deferred_dispatch_registrations = \\\"\\\"\\n static_init_dispatch_registrations = \\\"\\\"\\n if eager_registration:\\n static_template = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nTORCH_LIBRARY_IMPL(aten, $dispatch_key, m) {\\n $dispatch_registrations_body\\n};\\\"\\\"\\\"\\n )\\n static_init_dispatch_registrations = static_template.substitute(\\n dispatch_key=dispatch_key,\\n dispatch_registrations_body=dispatch_registrations_body,\\n )\\n else:\\n deferred_template = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nTORCH_API void Register${backend_name}${dispatch_key}NativeFunctions();\\nTORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() {\\n static auto m = MAKE_TORCH_LIBRARY_IMPL(aten, $dispatch_key);\\n $dispatch_registrations_body\\n}\\\"\\\"\\\"\\n )\\n deferred_dispatch_registrations = deferred_template.substitute(\\n backend_name=backend_name,\\n dispatch_key=dispatch_key,\\n dispatch_registrations_body=dispatch_registrations_body,\\n )\\n\\n fm.write_with_template(\\n f\\\"Register{dispatch_key}.cpp\\\",\\n \\\"RegisterDispatchKey.cpp\\\",\\n lambda: {\\n \\\"extra_cuda_headers\\\": \\\"\\\",\\n \\\"external_backend_headers\\\": external_backend_headers_str,\\n \\\"ops_headers\\\": \\\"#include \\\"\\n if not per_operator_headers\\n else \\\"\\\",\\n \\\"DispatchKey\\\": dispatch_key,\\n \\\"dispatch_namespace\\\": dispatch_key.lower(),\\n \\\"dispatch_headers\\\": dest.gen_registration_headers(\\n backend_index, per_operator_headers=per_operator_headers, rocm=False\\n ),\\n \\\"dispatch_definitions\\\": fm.substitute_with_template(\\n \\\"RegisterDispatchDefinitions.ini\\\",\\n lambda: {\\n \\\"ns_prologue\\\": ns_helper.prologue,\\n \\\"ns_epilogue\\\": ns_helper.epilogue,\\n \\\"static_init_dispatch_registrations\\\": static_init_dispatch_registrations,\\n \\\"deferred_dispatch_registrations\\\": deferred_dispatch_registrations,\\n \\\"dispatch_helpers\\\": dest.gen_registration_helpers(backend_index),\\n \\\"dispatch_namespace\\\": dispatch_key.lower(),\\n \\\"dispatch_namespaced_definitions\\\": \\\"\\\",\\n \\\"dispatch_anonymous_definitions\\\": list(\\n concatMap(\\n dest.RegisterDispatchKey(\\n backend_index,\\n Target.ANONYMOUS_DEFINITION,\\n selector,\\n rocm=False,\\n symint=True,\\n class_method_name=f\\\"{class_name}\\\",\\n skip_dispatcher_op_registration=False,\\n ),\\n grouped_native_functions,\\n )\\n ),\\n },\\n ).split(newline),\\n },\\n )\\n\\n\\ndef run(\\n source_yaml: str, output_dir: str, dry_run: bool, impl_path: str | None = None\\n) -> None:\\n # Assumes that this file lives at PYTORCH_ROOT/torchgen/gen_backend_stubs.py\\n pytorch_root = Path(__file__).parent.parent.absolute()\\n template_dir = os.path.join(pytorch_root, \\\"aten/src/ATen/templates\\\")\\n\\n def make_file_manager(install_dir: str) -> FileManager:\\n return FileManager(\\n install_dir=install_dir, template_dir=template_dir, dry_run=dry_run\\n )\\n\\n fm = make_file_manager(output_dir)\\n\\n native_yaml_path = os.path.join(\\n pytorch_root, \\\"aten/src/ATen/native/native_functions.yaml\\\"\\n )\\n tags_yaml_path = os.path.join(pytorch_root, \\\"aten/src/ATen/native/tags.yaml\\\")\\n parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)\\n native_functions, backend_indices = (\\n parsed_yaml.native_functions,\\n parsed_yaml.backend_indices,\\n )\\n grouped_native_functions = get_grouped_native_functions(native_functions)\\n parsed_backend_yaml = parse_backend_yaml(\\n source_yaml, grouped_native_functions, backend_indices\\n )\\n backend_key = parsed_backend_yaml.backend_key\\n autograd_key = parsed_backend_yaml.autograd_key\\n cpp_namespace = parsed_backend_yaml.cpp_namespace\\n class_name = parsed_backend_yaml.class_name\\n backend_indices = parsed_backend_yaml.backend_indices\\n\\n selector = SelectiveBuilder.get_nop_selector()\\n\\n if backend_key is None:\\n # This could be useful if a backend wants to quickly set up a noop yaml file but doesn't have any kernels ready yet.\\n return\\n\\n if class_name is None:\\n # class_name is an optional argument to backend yaml file.\\n # if specified it allows an external backend to override\\n # the name of the class that all generated kernel definitions live under.\\n # if not specified, its value is given as native_function_class_name.\\n class_name = backend_indices[backend_key].native_function_class_name()\\n assert class_name is not None\\n\\n if impl_path is not None:\\n error_on_missing_kernels(\\n native_functions,\\n backend_indices,\\n backend_key,\\n autograd_key,\\n class_name,\\n impl_path,\\n )\\n\\n gen_dispatchkey_nativefunc_headers(\\n fm,\\n class_name,\\n cpp_namespace,\\n backend_indices,\\n grouped_native_functions,\\n backend_key,\\n autograd_key,\\n )\\n\\n for dispatch_key in (\\n [backend_key] if autograd_key is None else [backend_key, autograd_key]\\n ):\\n gen_dispatcher_registrations(\\n fm,\\n output_dir,\\n class_name,\\n backend_indices,\\n grouped_native_functions,\\n backend_key,\\n dispatch_key,\\n selector,\\n )\\n\\n\\nif __name__ == \\\"__main__\\\":\\n main()\\n\\n\\nfrom __future__ import annotations\\n\\nimport textwrap\\nfrom dataclasses import dataclass\\nfrom typing import Sequence\\n\\nfrom torchgen.api.types import DispatcherSignature\\nfrom torchgen.api.types.signatures import CppSignature, CppSignatureGroup\\nfrom torchgen.context import method_with_native_function\\nfrom torchgen.model import (\\n Argument,\\n BackendIndex,\\n BaseTy,\\n BaseType,\\n DispatchKey,\\n FunctionSchema,\\n ListType,\\n NativeFunction,\\n NativeFunctionsGroup,\\n OperatorName,\\n OptionalType,\\n Type,\\n)\\nfrom torchgen.utils import mapMaybe\\n\\n\\nbase_type_to_c_type = {\\n BaseTy.Tensor: \\\"AtenTensorHandle\\\",\\n BaseTy.bool: \\\"int32_t\\\", # Use int to pass bool\\n BaseTy.int: \\\"int64_t\\\",\\n BaseTy.SymInt: \\\"int64_t\\\", # Inductor-generated code won't see a SymInt\\n BaseTy.Scalar: \\\"double\\\", # Use double to pass both integer and floating point\\n BaseTy.float: \\\"double\\\", # TODO: how about other floating point types?\\n BaseTy.str: \\\"const char*\\\",\\n BaseTy.DeviceIndex: \\\"int32_t\\\",\\n BaseTy.Layout: \\\"int32_t\\\", # Represent enum as int\\n BaseTy.MemoryFormat: \\\"int32_t\\\", # Represent enum as int\\n BaseTy.ScalarType: \\\"int32_t\\\", # Represent enum as int\\n BaseTy.Generator: \\\"AtenGeneratorHandle\\\",\\n}\\n\\nbase_type_to_aten_type = {\\n BaseTy.Tensor: \\\"at::Tensor\\\",\\n BaseTy.bool: \\\"bool\\\",\\n BaseTy.int: \\\"int64_t\\\",\\n BaseTy.SymInt: \\\"c10::SymInt\\\",\\n BaseTy.Scalar: \\\"c10::Scalar\\\",\\n BaseTy.float: \\\"double\\\",\\n BaseTy.str: \\\"c10::string_view\\\",\\n BaseTy.DeviceIndex: \\\"c10::DeviceIndex\\\",\\n BaseTy.Layout: \\\"c10::Layout\\\",\\n BaseTy.MemoryFormat: \\\"c10::MemoryFormat\\\",\\n BaseTy.ScalarType: \\\"c10::ScalarType\\\",\\n BaseTy.Generator: \\\"at::Generator\\\",\\n}\\n\\nbase_type_to_callsite_expr = {\\n BaseTy.Tensor: \\\"*tensor_handle_to_tensor_pointer\\\",\\n BaseTy.bool: \\\"\\\",\\n BaseTy.int: \\\"\\\",\\n BaseTy.SymInt: \\\"\\\",\\n BaseTy.Scalar: \\\"\\\",\\n BaseTy.float: \\\"\\\",\\n BaseTy.str: \\\"\\\",\\n BaseTy.DeviceIndex: \\\"static_cast\\\",\\n BaseTy.Layout: \\\"static_cast\\\",\\n BaseTy.MemoryFormat: \\\"static_cast\\\",\\n BaseTy.ScalarType: \\\"static_cast\\\",\\n BaseTy.Generator: \\\"*generator_handle_to_generator_pointer\\\",\\n}\\n\\n\\n# convert args to C types, names in declarations, and expressions in function bodies\\ndef convert_arg_type_and_name(typ: Type, name: str) -> tuple[list[str], list[str], list[str], list[str]]: # type: ignore[return]\\n if isinstance(typ, BaseType):\\n if typ.name in base_type_to_c_type:\\n return (\\n [base_type_to_c_type[typ.name]],\\n [name],\\n [base_type_to_aten_type[typ.name]],\\n [\\n f\\\"{base_type_to_callsite_expr[typ.name]}({name})\\\"\\n if base_type_to_callsite_expr[typ.name]\\n else name\\n ],\\n )\\n elif typ.name == BaseTy.Device:\\n return (\\n [\\\"int32_t\\\", \\\"int32_t\\\"],\\n [name, name + \\\"_index_\\\"],\\n [\\\"c10::Device\\\"],\\n [\\n f\\\"c10::Device(static_cast({name}), static_cast({name}_index_))\\\"\\n ],\\n )\\n else:\\n # TODO: BaseTy.Dimname, etc.\\n raise NotImplementedError(f\\\"TODO: add support for arg type {repr(typ)}\\\")\\n elif isinstance(typ, OptionalType):\\n c_types, names, aten_types, callsite_exprs = convert_arg_type_and_name(\\n typ.elem, name\\n )\\n j = 0 # index for names\\n new_aten_types = []\\n new_callsite_exprs = []\\n for aten_type in aten_types:\\n # Use pointer to denote optional type\\n c_types[j] = c_types[j] + \\\"*\\\"\\n if aten_type.startswith(\\\"c10::ArrayRef<\\\"):\\n # ArrayRef is passed as pointer + size, but no need to add \\\"*\\\" to the size argument\\n new_aten_types.append(f\\\"::std::optional<{aten_type}>\\\")\\n base_type = aten_type[len(\\\"c10::ArrayRef<\\\") : -1]\\n new_callsite_exprs.append(\\n f\\\"pointer_to_optional_list<{base_type}>({names[j]}, {names[j+1]})\\\"\\n )\\n j += 2\\n elif aten_type == \\\"c10::Device\\\":\\n # Device is passed as device_type + device_index\\n new_aten_types.append(\\\"::std::optional\\\")\\n new_callsite_exprs.append(\\n f\\\"pointer_to_optional_device({names[j]}, {names[j+1]})\\\"\\n )\\n j += 2\\n else:\\n new_aten_types.append(f\\\"::std::optional<{aten_type}>\\\")\\n new_callsite_exprs.append(\\n f\\\"pointer_to_optional<{aten_type}>({names[j]})\\\"\\n )\\n j += 1\\n\\n return (\\n c_types,\\n names,\\n new_aten_types,\\n new_callsite_exprs,\\n )\\n elif isinstance(typ, ListType):\\n # Need to explictly pass the list as pointer + length\\n c_types, names, aten_types, _ = convert_arg_type_and_name(typ.elem, name)\\n assert len(c_types) == 1, \\\"ListType with unsupported element type \\\" + repr(typ)\\n\\n # The list content should never be modified\\n c_types[0] = f\\\"const {c_types[0]}*\\\"\\n c_types.append(\\\"int64_t\\\")\\n name = names[0]\\n names.append(name + \\\"_len_\\\")\\n\\n atype = aten_types[0]\\n callsite_exprs = []\\n if atype == \\\"bool\\\":\\n # no converter from std::vector to c10::ArrayRef\\n # construct std::array instead\\n assert typ.size is not None\\n callsite_exprs.append(f\\\"pointer_to_list<{typ.size}>({name})\\\")\\n elif atype == \\\"::std::optional\\\":\\n # convert from std::vector<::std::optional> to c10::List<::std::optional>\\n callsite_exprs.append(\\n f\\\"c10::List<{atype}>(c10::ArrayRef<{atype}>(pointer_to_list<{atype}>({name}, {name}_len_)))\\\"\\n )\\n else:\\n callsite_exprs.append(f\\\"pointer_to_list<{atype}>({name}, {name}_len_)\\\")\\n\\n aten_types = [f\\\"c10::ArrayRef<{t}>\\\" for t in aten_types]\\n return (\\n c_types,\\n names,\\n aten_types,\\n callsite_exprs,\\n )\\n\\n\\ndef zip_type_and_name(types: list[str], names: list[str]) -> list[str]:\\n return [typ + \\\" \\\" + name for typ, name in zip(types, names)]\\n\\n\\n# Generate argument declarations and callsite expressions\\ndef gen_arguments(flat_arguments: Sequence[Argument]) -> tuple[list[str], list[str]]:\\n types = []\\n new_names = []\\n callsite_exprs = []\\n for arg in flat_arguments:\\n new_types, names, _, new_callsite_exprs = convert_arg_type_and_name(\\n arg.type, arg.name\\n )\\n types.extend(new_types)\\n new_names.extend(names)\\n callsite_exprs.extend(new_callsite_exprs)\\n return zip_type_and_name(types, new_names), callsite_exprs\\n\\n\\n# Return values are passed out as pointer arguments because all the C shim functions\\n# are expected to return AOTITorchError.\\n# Generate returns as declarations and callsite expressions\\ndef gen_returns(schema: FunctionSchema) -> tuple[list[str], list[str]]:\\n types = []\\n names = []\\n for idx, ret in enumerate(schema.returns):\\n names.append(f\\\"ret{idx}\\\")\\n if isinstance(ret.type, BaseType) and ret.type.name in base_type_to_c_type:\\n types.append(base_type_to_c_type[ret.type.name] + \\\"*\\\")\\n else:\\n raise NotImplementedError(\\n f\\\"TODO: add support for return type {repr(ret.type)}\\\"\\n )\\n\\n def convert_return(typ: BaseType, val: str) -> str:\\n if typ.name == BaseTy.Tensor:\\n return f\\\"new_tensor_handle(std::move({val}));\\\"\\n elif typ.name == BaseTy.SymInt:\\n return f\\\"{val}.expect_int()\\\"\\n elif typ.name == BaseTy.Scalar:\\n return f\\\"{val}.toDouble()\\\"\\n else:\\n return val\\n\\n ret_pointer_can_be_null = False\\n unambiguous_name = schema.name.unambiguous_name()\\n for name in [\\n \\\"_scaled_dot_product_flash_attention\\\",\\n \\\"_scaled_dot_product_efficient_attention\\\",\\n \\\"_scaled_dot_product_cudnn_attention\\\",\\n \\\"convolution_backward\\\",\\n ]:\\n if name in unambiguous_name:\\n ret_pointer_can_be_null = True\\n break\\n\\n callsite_exprs: list[str] = []\\n for idx, ret in enumerate(schema.returns):\\n tmp = \\\"tmp_result\\\" if len(names) == 1 else f\\\"std::get<{idx}>(tmp_result)\\\"\\n assert isinstance(ret.type, BaseType)\\n rval = convert_return(ret.type, tmp)\\n if ret_pointer_can_be_null:\\n callsite_exprs.append(f\\\"if ({names[idx]}) {{ *{names[idx]} = {rval}; }}\\\")\\n else:\\n callsite_exprs.append(f\\\"*{names[idx]} = {rval};\\\")\\n\\n return zip_type_and_name(types, names), callsite_exprs\\n\\n\\n# gen.py generates header first and then src, so caching the result here to avoid duplicate work\\ndeclaration_definition_cache: dict[tuple[str, str, str], tuple[str, str]] = {}\\n\\n\\ndef gen_declaration_and_definition(\\n schema: FunctionSchema, device: str, backend_call: str\\n) -> tuple[str, str]:\\n func_name = schema.name.unambiguous_name()\\n\\n global declaration_definition_cache\\n if (func_name, device, backend_call) in declaration_definition_cache:\\n return declaration_definition_cache[(func_name, device, backend_call)]\\n\\n if schema.is_out_fn():\\n # out_variant has out arguments in the front, and it's ok to ignore return values\\n # because C shim functions only return AOTITorchError\\n args, callsite_exprs = gen_arguments(\\n [*schema.arguments.out, *schema.arguments.flat_non_out]\\n )\\n ret_assignments: list[str] = []\\n else:\\n args, callsite_exprs = gen_arguments(schema.arguments.flat_all)\\n # ignore return values for inplace ops\\n ret_declarations, ret_assignments = (\\n ([], []) if schema.name.name.inplace else gen_returns(schema)\\n )\\n args.extend(ret_declarations)\\n\\n declaration = f\\\"AOTITorchError aoti_torch_{device}_{func_name}({', '.join(args)})\\\"\\n\\n tmp_result = \\\"auto tmp_result = \\\" if ret_assignments else \\\"\\\"\\n ret_assignments_str = \\\"\\\\n\\\" + \\\"\\\\n\\\".join(ret_assignments) if ret_assignments else \\\"\\\"\\n definition = f\\\"\\\"\\\"\\n{declaration} {{\\n AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({{\\n {tmp_result}{backend_call}(\\n{textwrap.indent(', '.join(callsite_exprs), \\\" \\\")}\\n );{textwrap.indent(ret_assignments_str, \\\" \\\")}\\n }});\\n}}\\n\\\"\\\"\\\"\\n declaration_definition_cache[(func_name, device, backend_call)] = (\\n declaration,\\n definition,\\n )\\n return declaration, definition\\n\\n\\ndef gen_static_dispatch_backend_call_signature(\\n sig: CppSignature | DispatcherSignature,\\n f: NativeFunction,\\n) -> CppSignature:\\n sig = DispatcherSignature.from_schema(f.func)\\n cpp_sigs = CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=False\\n )\\n if sig.symint and f.func.has_symint():\\n cpp_sig = cpp_sigs.symint_signature\\n else:\\n cpp_sig = cpp_sigs.signature\\n assert cpp_sig is not None\\n return cpp_sig\\n\\n\\ndef gen_static_dispatch_backend_call(\\n f: NativeFunction,\\n backend_index: BackendIndex,\\n) -> str:\\n sig = DispatcherSignature.from_schema(f.func)\\n cpp_sig = gen_static_dispatch_backend_call_signature(sig, f)\\n return f\\\"at::{backend_index.dispatch_key.lower()}::{cpp_sig.name()}\\\"\\n\\n\\ndef get_backend_index_for_aoti(\\n func: NativeFunction,\\n func_group_mapping: dict[OperatorName, NativeFunctionsGroup],\\n dispatch_key: DispatchKey,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n) -> BackendIndex | None:\\n backend_index = None\\n if backend_indices[dispatch_key].has_kernel(func) or (\\n func.structured_delegate is not None\\n and func.structured_delegate in func_group_mapping\\n and backend_indices[dispatch_key].has_kernel(\\n func_group_mapping[func.structured_delegate]\\n )\\n ):\\n backend_index = backend_indices[dispatch_key]\\n elif backend_indices[DispatchKey.CompositeExplicitAutograd].has_kernel(func):\\n # We need to create C shim wrappers for CompositeExplicitAutograd kernels\\n backend_index = backend_indices[DispatchKey.CompositeExplicitAutograd]\\n elif backend_indices[DispatchKey.CompositeExplicitAutogradNonFunctional].has_kernel(\\n func\\n ):\\n # We need to create C shim wrappers for CompositeExplicitAutogradNonFunctional kernels\\n backend_index = backend_indices[\\n DispatchKey.CompositeExplicitAutogradNonFunctional\\n ]\\n elif backend_indices[DispatchKey.CompositeImplicitAutograd].has_kernel(func):\\n backend_index = backend_indices[DispatchKey.CompositeImplicitAutograd]\\n\\n return backend_index\\n\\n\\ndef get_header_for_aoti(\\n func: NativeFunction,\\n func_group_mapping: dict[OperatorName, NativeFunctionsGroup],\\n dispatch_key: DispatchKey,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n) -> str | None:\\n backend_index = get_backend_index_for_aoti(\\n func, func_group_mapping, dispatch_key, backend_indices\\n )\\n return (\\n None\\n if backend_index is None\\n else f\\\"#include \\\"\\n )\\n\\n\\ndef get_fallback_op_name(func: NativeFunction) -> str:\\n return (\\n f\\\"{func.namespace}.{func.func.name.name}.{func.func.name.overload_name}\\\"\\n if func.func.name.overload_name\\n else f\\\"{func.namespace}.{func.func.name.name}.default\\\"\\n )\\n\\n\\ndef gen_c_shim(\\n func: NativeFunction,\\n func_group_mapping: dict[OperatorName, NativeFunctionsGroup],\\n dispatch_key: DispatchKey,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n header: bool,\\n) -> str | None:\\n backend_index = get_backend_index_for_aoti(\\n func, func_group_mapping, dispatch_key, backend_indices\\n )\\n if backend_index is None:\\n return None\\n\\n schema = func.func\\n device = dispatch_key.lower()\\n backend_call = gen_static_dispatch_backend_call(\\n func,\\n backend_index,\\n )\\n\\n try:\\n if header:\\n declaration, _ = gen_declaration_and_definition(\\n schema, device, backend_call\\n )\\n return f\\\"AOTI_TORCH_EXPORT {declaration};\\\"\\n else:\\n _, definition = gen_declaration_and_definition(schema, device, backend_call)\\n return definition\\n\\n except NotImplementedError:\\n return None\\n\\n\\n@dataclass(frozen=True)\\nclass ShimGenerator:\\n func_group_mapping: dict[OperatorName, NativeFunctionsGroup]\\n dispatch_key: DispatchKey\\n backend_indices: dict[DispatchKey, BackendIndex]\\n header: bool # True to generate .h and False to generate .cpp\\n\\n @method_with_native_function\\n def __call__(\\n self,\\n func: NativeFunction,\\n ) -> str | None:\\n result = gen_c_shim(\\n func,\\n self.func_group_mapping,\\n self.dispatch_key,\\n self.backend_indices,\\n self.header,\\n )\\n return result\\n\\n\\ndef gen_aoti_c_shim(\\n native_functions: Sequence[NativeFunction],\\n func_group_mapping: dict[OperatorName, NativeFunctionsGroup],\\n dispatch_key: DispatchKey,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n header: bool,\\n includes: str = \\\"\\\",\\n) -> str:\\n body = \\\"\\\\n\\\".join(\\n list(\\n mapMaybe(\\n ShimGenerator(\\n func_group_mapping, dispatch_key, backend_indices, header\\n ),\\n native_functions,\\n )\\n )\\n )\\n device = dispatch_key.lower()\\n\\n warning = \\\"\\\"\\\"\\n// WARNING: THIS FILE IS AUTOGENERATED BY torchgen. DO NOT MODIFY BY HAND.\\n// See https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 for details\\\"\\\"\\\"\\n\\n if header:\\n return f\\\"\\\"\\\"\\n{warning}\\n\\n#pragma once\\n\\n#include \\n\\n#ifdef __cplusplus\\nextern \\\"C\\\" {{\\n#endif\\n\\n{body}\\n\\n#ifdef __cplusplus\\n}} // extern \\\"C\\\"\\n#endif\\n\\\"\\\"\\\"\\n\\n else:\\n return f\\\"\\\"\\\"\\n{warning}\\n\\n#include \\n#include \\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#include \\n#include \\n#include \\n#else\\n{includes}\\n#endif\\n\\nusing namespace torch::aot_inductor;\\n\\n{body}\\\"\\\"\\\"\\n\\n\\nfrom __future__ import annotations\\n\\nimport argparse\\nimport os\\nfrom collections import namedtuple\\nfrom pathlib import Path\\nfrom typing import Any, Callable, Iterable, Iterator, Sequence\\n\\nimport yaml\\n\\nimport torchgen.dest as dest\\nfrom torchgen.api.lazy import setValueT\\nfrom torchgen.api.types import BaseCppType\\nfrom torchgen.dest.lazy_ir import GenLazyIR, GenLazyNativeFuncDefinition, GenTSLazyIR\\nfrom torchgen.gen import get_grouped_native_functions, parse_native_yaml\\nfrom torchgen.gen_backend_stubs import (\\n error_on_missing_kernels,\\n gen_dispatcher_registrations,\\n gen_dispatchkey_nativefunc_headers,\\n parse_backend_yaml,\\n)\\nfrom torchgen.model import NativeFunction, NativeFunctionsGroup, OperatorName\\nfrom torchgen.selective_build.selector import SelectiveBuilder\\nfrom torchgen.utils import FileManager, NamespaceHelper\\nfrom torchgen.yaml_utils import YamlLoader\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Lazy Tensor Codegen\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n# Overview\\n# ~~~~~~~~\\n#\\n# This codegen script builds on existing data models and helpers used\\n# by all ATen backends, and adds new functionality specific to lazy\\n# tensor backends.\\n#\\n# Inputs:\\n# - _native_functions.yaml: controls which operators are\\n# supported by the backend.\\n#\\n# Outputs:\\n# (for all backends)\\n# Ir.h defines Lazy IR classes to be constructed during tracing\\n# - opt-in: also generate 'lowering' methods for the TorchScript backend only\\n# NativeFunctions.cpp defines implementations of native functions which perform lazy tracing\\n# - opt-in: 'full_codegen' section of backend yaml; 'supported' section omits these implementations\\n# NativeFunctions.h declares implementations of native functions for both 'supported' and 'full_codegen'\\n# ops\\n#\\n# Register.cpp registers all op implementations with the dispatcher\\n# RegisterAutograd.cpp registers all autograd implementations with the dispatcher\\n#\\n# Validation Helpers:\\n# - Shape Inference: errs if any ops in backend yaml require shape inference not provided by meta kernels or\\n# implementations in torch/csrc/lazy/core/shape_inference.*\\n# - native function impls: errs if any 'supported' ops do not have an implementation defined in the backend\\n# (non-codegen) implementation file\\n#\\n#\\n# About the Data Model\\n# ~~~~~~~~~~~~~~~~~~~~\\n#\\n# Modeled after ATen codegen, the first step is to parse yaml and build a data model for the operators\\n# we care about. In this case, the _native_functions yaml defines a subset of the core operators\\n# (defined in more detail in the main native_functions.yaml), which will be supported by your backend.\\n# Backends can list ops in two categories:\\n# - `supported` ops require hand-implementations but still get codegenned declarations and registrations\\n# - `full_codegen` ops get implementations (and IR classes) generated too\\n#\\n# Each native function is modeled as an object with a schema, and each schema has objects representing their\\n# arguments. Much of the codegen is manipulation of the arguments and their types. For example, lazy tensor\\n# backends need to transform 'at::Tensor' arguments into 'lazy::Value' objects, as well as replacing reference\\n# types (stringref) with actual string objects, and this is done by manipulating the data model objects.\\n# - see api/lazy.py for the lazy data model\\n#\\n# Once the data model is set up, the rest of this script processes a number of templates for output CPP file\\n# and fills in the template values using helpers in `dest/lazy_ir.py` and `dest/lazy_ts_lowering.py`. These\\n# helpers mostly iterate over functions and their arguments, outputting different c++ snippets.\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\n# Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key.\\n# Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping, full_codegen)\\nParsedExternalYaml = namedtuple(\\n \\\"ParsedExternalYaml\\\",\\n [\\\"backend_key\\\", \\\"autograd_key\\\", \\\"cpp_namespace\\\", \\\"backend_indices\\\", \\\"full_codegen\\\"],\\n)\\n\\n\\ndef parse_native_functions_keys(\\n backend_yaml_path: str,\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n) -> tuple[list[OperatorName], list[Any], list[OperatorName]]:\\n with open(backend_yaml_path) as f:\\n yaml_values = yaml.load(f, Loader=YamlLoader)\\n assert isinstance(yaml_values, dict)\\n\\n full_codegen = yaml_values.pop(\\\"full_codegen\\\", [])\\n non_native = yaml_values.pop(\\\"non_native\\\", [])\\n ir_gen = yaml_values.pop(\\\"ir_gen\\\", [])\\n assert isinstance(full_codegen, list)\\n assert isinstance(non_native, list)\\n assert isinstance(ir_gen, list)\\n full_codegen_opnames = [OperatorName.parse(name) for name in full_codegen]\\n ir_gen_opnames = [OperatorName.parse(name) for name in ir_gen]\\n return full_codegen_opnames, non_native, ir_gen_opnames\\n\\n\\ndef validate_shape_inference_header(\\n shape_inference_hdr: str, expected_shape_infr_decls: list[str]\\n) -> None:\\n try:\\n with open(shape_inference_hdr) as f:\\n shape_infr_decls = f.read()\\n shape_infr_decl_lines = set(shape_infr_decls.split(\\\"\\\\n\\\"))\\n except OSError as e:\\n raise AssertionError(\\n f\\\"Unable to read from the specified shape_inference_hdr file: {shape_inference_hdr}\\\"\\n ) from e\\n\\n # TODO(whc) add a check for shape inference functions that have meta kernels implement and should be retired.\\n\\n missing_decls = [\\n decl for decl in expected_shape_infr_decls if decl not in shape_infr_decl_lines\\n ]\\n if missing_decls:\\n raise Exception( # noqa: TRY002\\n f\\\"\\\"\\\"Missing shape inference function.\\\\n\\nPlease add declare this function in {shape_inference_hdr}:\\\\n\\nand implement it in the corresponding shape_inference.cpp file.\\\\n\\n{os.linesep.join(missing_decls)}\\\"\\\"\\\"\\n )\\n\\n\\n# Some helper functions for the codegen.\\ndef get_ltc_helper_fns() -> str:\\n return \\\"\\\"\\\"\\\\\\nat::Tensor to_meta(const at::Tensor& tensor) {\\n // undefined tensors can't be converted to the meta device, since they don't have sizes/strides\\n if (!tensor.defined()) return tensor;\\n auto out = at::native::empty_strided_meta_symint(tensor.sym_sizes(), tensor.sym_strides(), \\\\\\n/*dtype=*/std::make_optional(tensor.scalar_type()), /*layout=*/std::make_optional(tensor.layout()), \\\\\\n/*device=*/std::make_optional(c10::Device(c10::kMeta)), /*pin_memory=*/std::nullopt);\\n // needs to handle wrapped numbers, so dtype promotion works properly.\\n if (tensor.unsafeGetTensorImpl()->is_wrapped_number()) {\\n out.unsafeGetTensorImpl()->set_wrapped_number(true);\\n }\\n return out;\\n}\\nstd::optional to_meta(const std::optional& tensor) {\\n if (tensor.has_value()) {\\n return to_meta(*tensor);\\n }\\n return std::nullopt;\\n}\\n\\nstd::vector to_meta(at::ITensorListRef t_list) {\\n std::vector outs;\\n outs.reserve(t_list.size());\\n for (const auto& tensor : t_list) {\\n outs.push_back(to_meta(tensor));\\n }\\n return outs;\\n}\\n\\\"\\\"\\\"\\n\\n\\nclass default_args:\\n node_base: str = \\\"Node\\\"\\n node_base_hdr: str | None = None\\n shape_inference_hdr: str = \\\"torch/csrc/lazy/core/shape_inference.h\\\"\\n tensor_class: str = \\\"torch::lazy::LazyTensor\\\"\\n tensor_class_hdr: str = \\\"torch/csrc/lazy/core/tensor.h\\\"\\n lazy_ir_generator: type[GenLazyIR] = GenLazyIR\\n native_func_definition_generator: type[\\n GenLazyNativeFuncDefinition\\n ] = GenLazyNativeFuncDefinition\\n backend_name: str = \\\"TorchScript\\\"\\n\\n\\ndef main() -> None:\\n parser = argparse.ArgumentParser(description=\\\"Generate Lazy Tensor backend files\\\")\\n parser.add_argument(\\n \\\"-s\\\",\\n \\\"--source-yaml\\\",\\n \\\"--source_yaml\\\",\\n help=\\\"path to source yaml file containing operator external definitions\\\",\\n )\\n parser.add_argument(\\\"-o\\\", \\\"--output-dir\\\", \\\"--output_dir\\\", help=\\\"output directory\\\")\\n parser.add_argument(\\n \\\"--dry-run\\\", \\\"--dry_run\\\", type=bool, default=False, help=\\\"output directory\\\"\\n )\\n parser.add_argument(\\n \\\"--impl-path\\\",\\n \\\"--impl_path\\\",\\n type=str,\\n default=None,\\n help=\\\"path to the source C++ file containing kernel definitions\\\",\\n )\\n parser.add_argument(\\n \\\"--gen-ts-lowerings\\\",\\n \\\"--gen_ts_lowerings\\\",\\n action=\\\"store_true\\\",\\n help=\\\"Generate TorchScript lowerings in addition to Lazy IR and NativeFunctions\\\",\\n )\\n parser.add_argument(\\n \\\"--node-base\\\",\\n \\\"--node_base\\\",\\n type=str,\\n default=default_args.node_base,\\n help=\\\"Name of backend specific custom Lazy IR Node base class\\\",\\n )\\n parser.add_argument(\\n \\\"--node-base-hdr\\\",\\n \\\"--node_base_hdr\\\",\\n type=str,\\n default=default_args.node_base_hdr,\\n help=\\\"Path to header file defining custom Lazy IR Node base class\\\",\\n )\\n parser.add_argument(\\n \\\"--shape-inference-hdr\\\",\\n \\\"--shape_inference_hdr\\\",\\n type=str,\\n default=default_args.shape_inference_hdr,\\n help=\\\"Path to header file defining custom Lazy shape inference functions\\\",\\n )\\n parser.add_argument(\\n \\\"--tensor-class\\\",\\n \\\"--tensor_class\\\",\\n type=str,\\n default=default_args.tensor_class,\\n help=\\\"Name of backend specific custom Lazy Tensor class\\\",\\n )\\n parser.add_argument(\\n \\\"--tensor-class-hdr\\\",\\n \\\"--tensor_class_hdr\\\",\\n type=str,\\n default=default_args.tensor_class_hdr,\\n help=\\\"Path to header file defining custom Lazy Tensor class\\\",\\n )\\n parser.add_argument(\\n \\\"--backend-name\\\",\\n \\\"--backend_name\\\",\\n type=str,\\n default=default_args.backend_name,\\n help=\\\"Name of the backend to generate\\\",\\n )\\n options = parser.parse_args()\\n\\n # Assumes that this file lives at PYTORCH_ROOT/torchgen/gen_backend_stubs.py\\n torch_root = Path(__file__).parent.parent.parent.absolute()\\n aten_path = str(torch_root / \\\"aten\\\" / \\\"src\\\" / \\\"ATen\\\")\\n lazy_ir_generator: type[GenLazyIR] = default_args.lazy_ir_generator\\n if options.gen_ts_lowerings:\\n lazy_ir_generator = GenTSLazyIR\\n native_func_definition_generator: type[\\n GenLazyNativeFuncDefinition\\n ] = default_args.native_func_definition_generator\\n\\n run_gen_lazy_tensor(\\n aten_path,\\n options.source_yaml,\\n options.output_dir,\\n options.dry_run,\\n options.impl_path,\\n options.node_base,\\n options.node_base_hdr,\\n options.tensor_class,\\n options.tensor_class_hdr,\\n options.shape_inference_hdr,\\n lazy_ir_generator,\\n native_func_definition_generator,\\n options.backend_name,\\n )\\n\\n\\ndef run_gen_lazy_tensor(\\n aten_path: str,\\n source_yaml: str,\\n output_dir: str,\\n dry_run: bool,\\n impl_path: str | None,\\n node_base: str = default_args.node_base,\\n node_base_hdr: str | None = default_args.node_base_hdr,\\n tensor_class: str = default_args.tensor_class,\\n tensor_class_hdr: str = default_args.tensor_class_hdr,\\n shape_inference_hdr: str = default_args.shape_inference_hdr,\\n lazy_ir_generator: type[GenLazyIR] = default_args.lazy_ir_generator,\\n native_func_definition_generator: type[\\n GenLazyNativeFuncDefinition\\n ] = default_args.native_func_definition_generator,\\n # build_in_tree is true for TS backend and affects include paths\\n build_in_tree: bool = False,\\n # per_operator_headers changes whether ATen/Functions.h or individual operator headers are used\\n # it must match how ATen was built\\n per_operator_headers: bool = False,\\n backend_name: str = default_args.backend_name,\\n gen_forced_fallback_code: bool = False,\\n use_lazy_shape: bool = True,\\n # the following arguments are temporary customization points for xla backend migration.\\n # do not rely on them otherwise, they should be removed once migration is complete\\n backend_namespace: str = \\\"torch::lazy\\\",\\n get_tensorlist: str = \\\"GetTensorList\\\",\\n get_tensor_or_wrap_number: str = \\\"GetLtcTensorOrCreateForWrappedNumber\\\",\\n try_get_tensor: str = \\\"TryGetLtcTensor\\\",\\n metrics_counter: str = 'TORCH_LAZY_FN_COUNTER(\\\"lazy::\\\")',\\n create_tensor: str = \\\"LazyTensor::Create\\\",\\n create_from_first_tensor: bool = False,\\n create_aten_from_ltc_tensor: str = \\\"torch::lazy::CreateAtenFromLtcTensor\\\",\\n tuple_aten_from_ltc_tensors: str = \\\"torch::lazy::TupleAtenFromLtcTensors\\\",\\n lazy_value_class: str = \\\"torch::lazy::Value\\\",\\n lazy_tensor_ptr: str = \\\"LazyTensorPtr\\\",\\n get_device_fn: str = \\\"torch::lazy::GetBackendDevice\\\",\\n) -> None:\\n lv_tokens = lazy_value_class.split(\\\"::\\\")\\n lv_class = lv_tokens[-1]\\n lv_ns = \\\"::\\\".join(lv_tokens[:-1])\\n setValueT(BaseCppType(lv_ns, lv_class))\\n template_dir = os.path.join(aten_path, \\\"templates\\\")\\n\\n def make_file_manager(install_dir: str) -> FileManager:\\n return FileManager(\\n install_dir=install_dir, template_dir=template_dir, dry_run=dry_run\\n )\\n\\n fm = make_file_manager(output_dir)\\n\\n native_yaml_path = os.path.join(aten_path, \\\"native/native_functions.yaml\\\")\\n tags_yaml_path = os.path.join(aten_path, \\\"native/tags.yaml\\\")\\n parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path)\\n native_functions, backend_indices = (\\n parsed_yaml.native_functions,\\n parsed_yaml.backend_indices,\\n )\\n grouped_native_functions = get_grouped_native_functions(native_functions)\\n\\n def sort_native_function(f: NativeFunctionsGroup | NativeFunction) -> str:\\n \\\"\\\"\\\"\\n We sort the native function because of the note in concat_map_codegen.\\n TODO(alanwaketan): Remove this sorting hack once all ops are grouped properly.\\n \\\"\\\"\\\"\\n func = f.functional.func if isinstance(f, NativeFunctionsGroup) else f.func\\n return str(func.name.name)\\n\\n grouped_native_functions = sorted(\\n grouped_native_functions, key=sort_native_function\\n )\\n\\n parsed_backend_yaml = parse_backend_yaml(\\n source_yaml, grouped_native_functions, backend_indices\\n )\\n backend_key = parsed_backend_yaml.backend_key\\n autograd_key = parsed_backend_yaml.autograd_key\\n cpp_namespace = parsed_backend_yaml.cpp_namespace\\n backend_indices = parsed_backend_yaml.backend_indices\\n # the following 3 keys are all processed differently\\n # for full_codegen, we generate IR, kernels, etc\\n # for ir_gen, we generate only IR\\n # non_native is used to register kernels not declared in\\n # native_functions.yaml\\n full_codegen, non_native, ir_gen = parse_native_functions_keys(\\n source_yaml, grouped_native_functions\\n )\\n\\n def concat_map_codegen(\\n func: Callable[[NativeFunction], Sequence[str]],\\n xs: Iterable[NativeFunctionsGroup | NativeFunction],\\n ops_list: list[OperatorName] = full_codegen,\\n ) -> Iterator[str]:\\n \\\"\\\"\\\"\\n We code-gen for the functional variant, which is all we need for IR classes/lowerings/shape inferences, but we\\n only code-gen additional entries for the inplace variant for the native functions.\\n \\\"\\\"\\\"\\n\\n for x in xs:\\n fs = list(x.functions()) if isinstance(x, NativeFunctionsGroup) else [x]\\n for f in fs:\\n if f.func.name in ops_list:\\n yield from func(f)\\n\\n selector = SelectiveBuilder.get_nop_selector()\\n\\n assert backend_key is not None\\n class_name = backend_indices[backend_key].native_function_class_name()\\n\\n if impl_path is not None:\\n error_on_missing_kernels(\\n native_functions,\\n backend_indices,\\n backend_key,\\n autograd_key,\\n class_name,\\n impl_path,\\n full_codegen,\\n )\\n\\n \\\"\\\"\\\" Validate Shape Inference Definitions\\n\\n Generated lazy native functions all perform shape inference, by first using a meta:: kernel\\n if available for that op, and otherwise using a 'compute_shape_{op}' function instead. The generator\\n knows the call signature for compute_shape_{op} because it matches the nativefunction (and meta::) signature,\\n so it just has to check whether the op is structured and generate a call for one or the other. It's up to the dev\\n to supply the missing compute_shape_{op} function, but the codegen at least warns you about this and provides\\n the expected signature which can be copy-pasted into shape_inference.h.\\n\\n compute_shape_{op} functions are handwritten and should be replaced over time as ops get ported\\n to structured kernels.\\n\\n See torch/csrc/lazy/core/shape_inference.cpp #READ THIS! for more information.\\n \\\"\\\"\\\"\\n if shape_inference_hdr is not None:\\n expected_shape_infr_decls = list(\\n concat_map_codegen(\\n dest.GenLazyShapeInferenceDefinition(\\n backend_indices[backend_key], tensor_class\\n ),\\n grouped_native_functions,\\n )\\n )\\n\\n validate_shape_inference_header(shape_inference_hdr, expected_shape_infr_decls)\\n assert class_name is not None\\n\\n # Generate nativefunction declarations\\n # Note, eager registrations is set to False for the lazy TS backend as another LTC backend\\n # may want to register their own lazy kernels instead of registering the TS ones.\\n # The registration will lazily happen when init_ts_backend is called.\\n gen_dispatchkey_nativefunc_headers(\\n fm,\\n class_name,\\n cpp_namespace,\\n backend_indices,\\n grouped_native_functions,\\n backend_key,\\n autograd_key,\\n backend_name,\\n )\\n\\n # Generate Dispatcher registrations which hook up the nativefunctions\\n for dispatch_key in (\\n [backend_key] if autograd_key is None else [backend_key, autograd_key]\\n ):\\n gen_dispatcher_registrations(\\n fm,\\n output_dir,\\n class_name,\\n backend_indices,\\n grouped_native_functions,\\n backend_key,\\n dispatch_key,\\n selector,\\n build_in_tree=build_in_tree,\\n per_operator_headers=per_operator_headers,\\n backend_name=backend_name,\\n eager_registration=False,\\n )\\n\\n # Generate native function impls that build IR nodes\\n ns_helper = NamespaceHelper(cpp_namespace)\\n fm.write_with_template(\\n f\\\"{backend_key}NativeFunctions.cpp\\\",\\n \\\"DispatchKeyNativeFunctions.cpp\\\",\\n lambda: {\\n \\\"includes\\\": [\\n f\\\"#include <{path}>\\\"\\n for path in [\\n tensor_class_hdr,\\n shape_inference_hdr,\\n \\\"ATen/Functions.h\\\",\\n \\\"ATen/native/TensorConversions.h\\\",\\n \\\"ATen/NativeFunctions.h\\\",\\n \\\"ATen/CompositeExplicitAutogradNonFunctionalFunctions.h\\\",\\n \\\"ATen/MetaFunctions.h\\\",\\n \\\"ATen/Operators.h\\\",\\n \\\"ATen/native/CPUFallback.h\\\",\\n \\\"torch/csrc/lazy/core/ir_builder.h\\\",\\n \\\"torch/csrc/lazy/core/lazy_graph_executor.h\\\",\\n \\\"torch/csrc/lazy/core/metrics.h\\\",\\n \\\"torch/csrc/lazy/core/shape.h\\\",\\n f\\\"{output_dir}/{backend_key}NativeFunctions.h\\\",\\n f\\\"{output_dir}/LazyIr.h\\\",\\n ]\\n + (\\n [\\\"torch/csrc/lazy/ts_backend/ts_eager_fallback.h\\\"]\\n if gen_forced_fallback_code\\n else []\\n )\\n ],\\n \\\"helper_fns\\\": get_ltc_helper_fns(),\\n \\\"native_functions_include\\\": \\\"\\\",\\n \\\"namespace_prologue\\\": ns_helper.prologue,\\n \\\"namespace_epilogue\\\": ns_helper.epilogue,\\n \\\"native_function_definitions\\\": list(\\n concat_map_codegen(\\n native_func_definition_generator(\\n f\\\"{backend_key}NativeFunctions\\\",\\n backend_indices[backend_key],\\n tensor_class,\\n gen_forced_fallback_code,\\n backend_namespace,\\n get_tensorlist,\\n get_tensor_or_wrap_number,\\n try_get_tensor,\\n metrics_counter,\\n create_tensor,\\n create_from_first_tensor,\\n create_aten_from_ltc_tensor,\\n tuple_aten_from_ltc_tensors,\\n lazy_tensor_ptr,\\n get_device_fn,\\n ),\\n grouped_native_functions,\\n )\\n ),\\n },\\n )\\n # Generate IR node classes\\n lazy_ir_obj = lazy_ir_generator(\\n backend_indices[backend_key], backend_name, node_base, use_lazy_shape\\n )\\n\\n fm.write_with_template(\\n \\\"LazyIr.h\\\",\\n \\\"LazyIr.h\\\",\\n lambda: {\\n \\\"lazy_ir_sysinc\\\": [\\n f\\\"#include <{path}>\\\"\\n for path in [\\n \\\"ATen/core/Formatting.h\\\",\\n \\\"c10/core/ScalarType.h\\\",\\n \\\"torch/csrc/lazy/core/hash.h\\\",\\n \\\"torch/csrc/lazy/core/ir.h\\\",\\n \\\"torch/csrc/lazy/core/shape.h\\\",\\n \\\"optional\\\",\\n \\\"vector\\\",\\n ]\\n ],\\n \\\"lazy_ir_inc\\\": [f'#include \\\"{node_base_hdr}\\\"']\\n if node_base_hdr is not None\\n else [],\\n \\\"ir_declarations\\\": list(\\n concat_map_codegen(\\n lazy_ir_obj, grouped_native_functions, full_codegen + ir_gen\\n )\\n ),\\n \\\"namespace_prologue\\\": ns_helper.prologue,\\n \\\"namespace_epilogue\\\": ns_helper.epilogue,\\n },\\n )\\n\\n # Generate Non Native IR Node classes\\n fm.write_with_template(\\n \\\"LazyNonNativeIr.h\\\",\\n \\\"LazyNonNativeIr.h\\\",\\n lambda: {\\n \\\"lazy_non_native_ir_inc\\\": [\\n f\\\"#include <{path}>\\\"\\n for path in [\\n \\\"torch/csrc/lazy/core/ir.h\\\",\\n \\\"torch/csrc/lazy/core/ir_builder.h\\\",\\n \\\"torch/csrc/lazy/core/internal_ops/ltc_ops.h\\\",\\n \\\"torch/csrc/lazy/core/shape_inference.h\\\",\\n ]\\n + ([node_base_hdr] if node_base_hdr else [])\\n if path\\n ],\\n \\\"non_native_ir_nodes\\\": dest.generate_non_native_lazy_ir_nodes(\\n non_native, lazy_ir_obj\\n ),\\n \\\"namespace_prologue\\\": ns_helper.prologue,\\n \\\"namespace_epilogue\\\": ns_helper.epilogue,\\n },\\n )\\n\\n\\nif __name__ == \\\"__main__\\\":\\n main()\\n\\n\\nfrom __future__ import annotations\\n\\nimport dataclasses\\nimport itertools\\nimport re\\nfrom dataclasses import dataclass\\nfrom enum import auto, Enum\\nfrom typing import Callable, Iterator, Sequence\\n\\nfrom torchgen.utils import assert_never, NamespaceHelper, OrderedSet\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# DATA MODEL\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Some general principles for our data model.\\n#\\n# - Stop using C++ data types as the internal data representation\\n# format. Instead, the internal data structures are centered\\n# around JIT schema representation. This avoid a big problem\\n# with the old codegen where we read in all the types from\\n# native_functions.yaml and then immediately had to retranslate\\n# them into C++ types.\\n#\\n# - More semantic data representation. Instead of representing\\n# everything as dicts and strings, we define dataclasses for\\n# every interesting entity the code generation has to deal with.\\n# These dataclasses have strong semantic invariants: for example,\\n# we generally require them to roundtrip losslessly into the\\n# form they were parsed from. These structures are immutable\\n# and you're expected to populate information once during\\n# construction.\\n\\n\\n# Represent a source location; used for better error reporting\\n@dataclass(frozen=True)\\nclass Location:\\n file: str\\n line: int\\n\\n def __str__(self) -> str:\\n return f\\\"{self.file}:{self.line}\\\"\\n\\n\\n# Valid values of the 'variants' field in native_functions.yaml\\nclass Variant(Enum):\\n function = auto()\\n method = auto()\\n\\n\\n# Default kernel namespace\\nDEFAULT_KERNEL_NAMESPACE = \\\"at::native\\\"\\n\\n# NOTE: Keep the list in sync with `DispatchKey` in c10/core/DispatchKey.h\\nBACKEND_COMPONENTS = \\\"CPU CUDA HIP XLA MTIA MPS IPU XPU HPU VE Lazy Meta PrivateUse1 PrivateUse2 PrivateUse3\\\".split()\\nFUNCTIONALITY_KEYS = [\\n \\\"\\\",\\n \\\"Quantized\\\",\\n \\\"Sparse\\\",\\n \\\"SparseCsr\\\",\\n \\\"NestedTensor\\\",\\n \\\"Autograd\\\",\\n]\\n\\n# This list guards dispatches that can be used in derivatives.yaml\\n# For now we omit AutogradFunctionality and AutogradOther\\nAUTOGRAD_KEYS = [\\\"AutogradNestedTensor\\\"] + [\\n \\\"Autograd\\\" + component for component in BACKEND_COMPONENTS\\n]\\n\\nFRAGMENT_NAMESPACES = {\\\"quantized\\\", \\\"quantized_decomposed\\\"}\\n\\n\\n# This doesn't have to be in sync with the header, it only needs to contain\\n# entries that we actually use in the codegen or want pyi entries for\\nclass DispatchKey(Enum):\\n Undefined = 0\\n CatchAll = Undefined\\n\\n FPGA = auto()\\n MAIA = auto()\\n Vulkan = auto()\\n Metal = auto()\\n MKLDNN = auto()\\n OpenGL = auto()\\n OpenCL = auto()\\n IDEEP = auto()\\n CustomRNGKeyId = auto()\\n MkldnnCPU = auto()\\n Sparse = auto()\\n SparseCsr = auto()\\n NestedTensor = auto()\\n Dense = auto()\\n\\n PythonTLSSnapshot = auto()\\n PreDispatch = auto()\\n PythonDispatcher = auto()\\n Python = auto()\\n FuncTorchDynamicLayerBackMode = auto()\\n ZeroTensor = auto()\\n Conjugate = auto()\\n Negative = auto()\\n BackendSelect = auto()\\n Named = auto()\\n AutogradOther = auto()\\n AutogradFunctionality = auto()\\n AutogradNestedTensor = auto()\\n Tracer = auto()\\n Autocast = auto()\\n AutocastCPU = auto()\\n AutocastCUDA = auto()\\n Batched = auto()\\n VmapMode = auto()\\n FuncTorchGradWrapper = auto()\\n FuncTorchBatched = auto()\\n BatchedNestedTensor = auto()\\n FuncTorchVmapMode = auto()\\n FuncTorchDynamicLayerFrontMode = auto()\\n Functionalize = auto()\\n TESTING_ONLY_GenericWrapper = auto()\\n TESTING_ONLY_GenericMode = auto()\\n\\n ADInplaceOrView = auto()\\n Autograd = auto()\\n CompositeImplicitAutograd = auto()\\n CompositeImplicitAutogradNestedTensor = auto()\\n CompositeExplicitAutograd = auto()\\n CompositeExplicitAutogradNonFunctional = auto()\\n FuncTorchBatchedDecomposition = auto()\\n\\n # BEGIN autogenerated\\n CPU = auto()\\n CUDA = auto()\\n HIP = auto()\\n XLA = auto()\\n MTIA = auto()\\n MPS = auto()\\n IPU = auto()\\n XPU = auto()\\n HPU = auto()\\n VE = auto()\\n Lazy = auto()\\n Meta = auto()\\n PrivateUse1 = auto()\\n PrivateUse2 = auto()\\n PrivateUse3 = auto()\\n QuantizedCPU = auto()\\n QuantizedCUDA = auto()\\n QuantizedHIP = auto()\\n QuantizedXLA = auto()\\n QuantizedMTIA = auto()\\n QuantizedMPS = auto()\\n QuantizedIPU = auto()\\n QuantizedXPU = auto()\\n QuantizedHPU = auto()\\n QuantizedVE = auto()\\n QuantizedLazy = auto()\\n QuantizedMeta = auto()\\n QuantizedPrivateUse1 = auto()\\n QuantizedPrivateUse2 = auto()\\n QuantizedPrivateUse3 = auto()\\n SparseCPU = auto()\\n SparseCUDA = auto()\\n SparseHIP = auto()\\n SparseXLA = auto()\\n SparseMTIA = auto()\\n SparseMPS = auto()\\n SparseIPU = auto()\\n SparseXPU = auto()\\n SparseHPU = auto()\\n SparseVE = auto()\\n SparseLazy = auto()\\n SparseMeta = auto()\\n SparsePrivateUse1 = auto()\\n SparsePrivateUse2 = auto()\\n SparsePrivateUse3 = auto()\\n SparseCsrCPU = auto()\\n SparseCsrCUDA = auto()\\n SparseCsrHIP = auto()\\n SparseCsrXLA = auto()\\n SparseCsrMTIA = auto()\\n SparseCsrMPS = auto()\\n SparseCsrIPU = auto()\\n SparseCsrXPU = auto()\\n SparseCsrHPU = auto()\\n SparseCsrVE = auto()\\n SparseCsrLazy = auto()\\n SparseCsrMeta = auto()\\n SparseCsrPrivateUse1 = auto()\\n SparseCsrPrivateUse2 = auto()\\n SparseCsrPrivateUse3 = auto()\\n NestedTensorCPU = auto()\\n NestedTensorCUDA = auto()\\n NestedTensorHIP = auto()\\n NestedTensorXLA = auto()\\n NestedTensorMTIA = auto()\\n NestedTensorMPS = auto()\\n NestedTensorIPU = auto()\\n NestedTensorXPU = auto()\\n NestedTensorHPU = auto()\\n NestedTensorVE = auto()\\n NestedTensorLazy = auto()\\n NestedTensorMeta = auto()\\n NestedTensorPrivateUse1 = auto()\\n NestedTensorPrivateUse2 = auto()\\n NestedTensorPrivateUse3 = auto()\\n AutogradCPU = auto()\\n AutogradCUDA = auto()\\n AutogradHIP = auto()\\n AutogradXLA = auto()\\n AutogradMTIA = auto()\\n AutogradMPS = auto()\\n AutogradIPU = auto()\\n AutogradXPU = auto()\\n AutogradHPU = auto()\\n AutogradVE = auto()\\n AutogradLazy = auto()\\n AutogradMeta = auto()\\n AutogradPrivateUse1 = auto()\\n AutogradPrivateUse2 = auto()\\n AutogradPrivateUse3 = auto()\\n # END autogenerated\\n\\n def __str__(self) -> str:\\n return self.name\\n\\n def lower(self) -> str:\\n return str(self).lower()\\n\\n @staticmethod\\n def parse(value: str) -> DispatchKey:\\n for k, v in DispatchKey.__members__.items():\\n if k == value:\\n return v\\n raise AssertionError(f\\\"unknown dispatch key {value}\\\")\\n\\n\\nclass _TorchDispatchModeKey(Enum):\\n FAKE = auto()\\n PROXY = auto()\\n FUNCTIONAL = auto()\\n\\n\\ndef codegen_per_backend_entries() -> str:\\n r = []\\n for fk in FUNCTIONALITY_KEYS:\\n for bc in BACKEND_COMPONENTS:\\n r.append(f\\\" {fk}{bc} = auto()\\\")\\n return \\\"\\\\n\\\".join(r)\\n\\n\\nfor fk in FUNCTIONALITY_KEYS:\\n for bc in BACKEND_COMPONENTS:\\n if not hasattr(DispatchKey, fk + bc):\\n r = codegen_per_backend_entries()\\n print(r)\\n raise RuntimeError(\\n f\\\"Missing {fk}{bc} from DispatchKey enum. Here is the autogenerated list we expect to have:\\\\n\\\\n{r}\\\"\\n )\\n\\n\\nSTRUCTURED_DISPATCH_KEYS = {\\n DispatchKey.MPS,\\n DispatchKey.CUDA,\\n DispatchKey.CPU,\\n DispatchKey.XPU,\\n}\\nUFUNC_DISPATCH_KEYS = {DispatchKey.CUDA, DispatchKey.CPU}\\n\\n# Set of supported dispatch keys\\ndispatch_keys = [\\n DispatchKey.CPU,\\n DispatchKey.SparseCPU,\\n DispatchKey.SparseCsrCPU,\\n DispatchKey.MkldnnCPU,\\n DispatchKey.CUDA,\\n DispatchKey.MPS,\\n DispatchKey.XPU,\\n DispatchKey.SparseCUDA,\\n DispatchKey.SparseCsrCUDA,\\n DispatchKey.QuantizedCPU,\\n DispatchKey.QuantizedCUDA,\\n DispatchKey.CompositeImplicitAutograd,\\n DispatchKey.CompositeImplicitAutogradNestedTensor,\\n DispatchKey.CompositeExplicitAutograd,\\n DispatchKey.CompositeExplicitAutogradNonFunctional,\\n DispatchKey.NestedTensorCPU,\\n DispatchKey.NestedTensorCUDA,\\n # Meta is a magic key: it is automatically generated for structured\\n # kernels\\n DispatchKey.Meta,\\n DispatchKey.SparseMeta,\\n DispatchKey.SparseCsrMeta,\\n DispatchKey.QuantizedMeta,\\n DispatchKey.NestedTensorMeta,\\n DispatchKey.ZeroTensor,\\n]\\n\\n\\n# Dispatch keys that \\\"support all backends\\\". These codegen slightly differently\\n# then backend specific keys.\\ndef is_generic_dispatch_key(dk: DispatchKey) -> bool:\\n return dk in {\\n DispatchKey.CompositeExplicitAutograd,\\n DispatchKey.CompositeExplicitAutogradNonFunctional,\\n DispatchKey.CompositeImplicitAutograd,\\n DispatchKey.CompositeImplicitAutogradNestedTensor,\\n }\\n\\n\\n# CUDA specific dispatch keys\\ndef is_cuda_dispatch_key(dk: DispatchKey) -> bool:\\n return dk in {\\n DispatchKey.CUDA,\\n DispatchKey.QuantizedCUDA,\\n DispatchKey.SparseCUDA,\\n DispatchKey.SparseCsrCUDA,\\n DispatchKey.NestedTensorCUDA,\\n DispatchKey.AutogradCUDA,\\n }\\n\\n\\n# XPU specific dispatcy keys\\ndef is_xpu_dispatch_key(dk: DispatchKey) -> bool:\\n return dk in {\\n DispatchKey.XPU,\\n DispatchKey.QuantizedXPU,\\n DispatchKey.SparseXPU,\\n DispatchKey.SparseCsrXPU,\\n DispatchKey.NestedTensorXPU,\\n DispatchKey.AutogradXPU,\\n }\\n\\n\\n# Structured kernel generation is only supported for certain key types;\\n# otherwise use old-style\\ndef is_structured_dispatch_key(dk: DispatchKey) -> bool:\\n return dk in STRUCTURED_DISPATCH_KEYS\\n\\n\\ndef is_ufunc_dispatch_key(dk: DispatchKey) -> bool:\\n # For now, ufunc dispatch keys coincide with structured keys\\n return dk in UFUNC_DISPATCH_KEYS\\n\\n\\n# This is oddly named ScalarType and not DType for symmetry with C++\\nclass ScalarType(Enum):\\n Byte = auto()\\n Char = auto()\\n Short = auto()\\n Int = auto()\\n Long = auto()\\n Half = auto()\\n Float = auto()\\n Double = auto()\\n ComplexHalf = auto()\\n ComplexFloat = auto()\\n ComplexDouble = auto()\\n Bool = auto()\\n BFloat16 = auto()\\n Float8_e5m2 = auto()\\n Float8_e5m2fnuz = auto()\\n Float8_e4m3fn = auto()\\n Float8_e4m3fnuz = auto()\\n\\n def __str__(self) -> str:\\n return self.name\\n\\n @staticmethod\\n def maybe_parse(value: str) -> ScalarType | None:\\n for k, v in ScalarType.__members__.items():\\n if k == value:\\n return v\\n return None\\n\\n @staticmethod\\n def parse(value: str) -> ScalarType:\\n mb_r = ScalarType.maybe_parse(value)\\n assert mb_r is not None, f\\\"unknown dtype {value}\\\"\\n return mb_r\\n\\n @staticmethod\\n def parse_set(values: str) -> OrderedSet[ScalarType]:\\n dtypes: OrderedSet[ScalarType] = OrderedSet()\\n for value in values.split(\\\", \\\"):\\n if value in DTYPE_CLASSES:\\n dtypes.update(DTYPE_CLASSES[value])\\n else:\\n dtypes.add(ScalarType.parse(value))\\n return dtypes\\n\\n\\nDTYPE_CLASSES: dict[str, OrderedSet[ScalarType]] = {}\\n# NB: Integral doesn't include boolean\\nDTYPE_CLASSES[\\\"Integral\\\"] = OrderedSet(\\n [\\n ScalarType.Byte,\\n ScalarType.Char,\\n ScalarType.Int,\\n ScalarType.Long,\\n ScalarType.Short,\\n ]\\n)\\n# NB: Floating doesn't include low precision types\\nDTYPE_CLASSES[\\\"Floating\\\"] = OrderedSet([ScalarType.Float, ScalarType.Double])\\nDTYPE_CLASSES[\\\"Complex\\\"] = OrderedSet(\\n [ScalarType.ComplexFloat, ScalarType.ComplexDouble]\\n)\\nDTYPE_CLASSES[\\\"All\\\"] = DTYPE_CLASSES[\\\"Integral\\\"] | DTYPE_CLASSES[\\\"Floating\\\"]\\nDTYPE_CLASSES[\\\"AllAndComplex\\\"] = DTYPE_CLASSES[\\\"All\\\"] | DTYPE_CLASSES[\\\"Complex\\\"]\\nDTYPE_CLASSES[\\\"FloatingAndComplex\\\"] = (\\n DTYPE_CLASSES[\\\"Floating\\\"] | DTYPE_CLASSES[\\\"Complex\\\"]\\n)\\n\\n\\n# Represents the valid entries for ufunc_inner_loop in native_functions.yaml.\\n# NB: if you add a new UfuncKey, you will teach torchgen.dest.ufunc how\\n# to process it. Most logic will ignore keys they don't understand, so your\\n# new key will get silently ignored until you hook in logic to deal with it.\\nclass UfuncKey(Enum):\\n # These are low level keys that represent exactly one particular\\n # instantiation of the kernel produced by codegen\\n CUDAFunctor = auto()\\n CUDAFunctorOnOther = auto()\\n CUDAFunctorOnSelf = auto()\\n\\n CPUScalar = auto()\\n CPUVector = auto()\\n\\n # These are the ones users will usually specify, and\\n # implicitly \\\"fill in\\\" the low level keys\\n ScalarOnly = auto() # CUDA*, CPUScalar\\n Generic = auto() # CUDA*, CPU*\\n\\n def __str__(self) -> str:\\n return self.name\\n\\n @staticmethod\\n def parse(value: str) -> UfuncKey:\\n for k, v in UfuncKey.__members__.items():\\n if k == value:\\n return v\\n raise AssertionError(f\\\"unknown ufunc key {value}\\\")\\n\\n\\nclass DeviceCheckType(Enum):\\n NoCheck = 0\\n ExactSame = 1\\n\\n\\nclass ViewSchemaKind(Enum):\\n aliasing = auto()\\n aliasing_inplace = auto()\\n non_aliasing = auto()\\n\\n\\n# The basic input to the code generation is native_functions.yaml.\\n# The name \\\"native\\\", BTW, comes from the distinction between native\\n# functions and legacy TH functions. The legacy TH functions are gone,\\n# but the \\\"native\\\" descriptor has stuck.\\n#\\n# NativeFunction models a single entry in native_functions.yaml. Its\\n# fields roughly correspond to what you would see in the YAML itself,\\n# but after canonicalization and parsing has occurred.\\n#\\n# You can see some of the overall design patterns for how we setup\\n# dataclasses in this class, but we will defer a complete discussion\\n# of this at FunctionSchema.\\n@dataclass(frozen=True)\\nclass NativeFunction:\\n # The namespace for this operator. For example, if we have \\\"at::add\\\"\\n # then the namespace would be \\\"at\\\". This enables ops to be registered\\n # through the same DSL with a custom namespace. If not specified, the\\n # default namespace would be \\\"at\\\".\\n namespace: str\\n\\n # The function schema of the operator in question. This schema\\n # has been parsed; see FunctionSchema for more about its structure.\\n # (This type is quoted as we are forward referencing a type\\n # defined later in the file. I opted for this ordering of the\\n # classes for expository clarity.)\\n func: FunctionSchema\\n\\n # Whether or not to generate mutable tensor arguments like regular\\n # ones\\n use_const_ref_for_mutable_tensors: bool\\n\\n # Whether or not to omit automatic generation of a DeviceGuard\\n device_guard: bool\\n\\n # How to emit automatic generation of device check\\n device_check: DeviceCheckType\\n\\n # What python module to put the function in\\n python_module: str | None\\n\\n # TODO: figure out what this does\\n category_override: str | None\\n\\n # If no variants are specified in native_functions.yaml, this is\\n # assumed to be {'function'}.\\n variants: set[Variant]\\n\\n # Whether or not we should skip generating registrations for\\n # this kernel. This is a bit of a double-edged sword, as manual\\n # registrations don't participate in codegen-based selective build!\\n manual_kernel_registration: bool\\n\\n # Whether or not to skip generating TensorMethod/Functions bindings\\n # for this kernel. Technically, this doesn't actually skip generating\\n # the binding; instead, the binding gets generated to __dispatch_{funcname}\\n # so you can make use of the normal binding if you need it.\\n manual_cpp_binding: bool\\n\\n # The location in the YAML file were this native function entry was\\n # defined. This is for conveniently reporting error messages!\\n loc: Location\\n\\n # A list of operators that are expected to be auto-generated for this NativeFunction.\\n # Note: This list isn't actually directly used by the codegen to generate anything.\\n # Instead, the codegen figures out what operators to generate purely based off of\\n # function schema, and uses the autogen declarations to error check.\\n # We expect every NativeFunction that gets auto-generated be explicitly called out\\n # in native_functions.yaml\\n autogen: list[OperatorName]\\n\\n # If non-empty, this kernel is subject to ufunc codegen.\\n # Sorted by ufunc_key\\n ufunc_inner_loop: dict[UfuncKey, UfuncInnerLoop]\\n\\n # Whether or not this out functions is a \\\"structured kernel\\\". Structured\\n # kernels are defined a little differently from normal kernels; in\\n # particular, their shape checking logic is defined separately from\\n # the kernel. Only out functions can be structured; other functions\\n # delegate to the out function using the structured_delegate keyword.\\n # Every structured kernel must have at least an out and a functional\\n # variant.\\n structured: bool\\n\\n # Whether or not this non-out function is a structured kernel, defined\\n # in terms of the out kernel referenced by the string here.\\n structured_delegate: OperatorName | None\\n\\n # Only valid for structured kernels. Specifies alternative of what\\n # to inherit from when defining the meta class for the structured\\n # operator. This will usually be TensorIteratorBase. This also\\n # changes the semantics of set_output to call the parent class.\\n structured_inherits: str | None\\n\\n # Structured kernels can declare elements as \\\"precomputed\\\". These elements\\n # are returned by the meta function in one struct and passed to the impl\\n # function in lieu of certain kernel arguments that these precomputed\\n # elements supersede. Information about the names and types of these\\n # precomputed elements and how they correspond to kernel arguments is stored\\n # in this member, if applicable.\\n precomputed: Precompute | None\\n\\n # Argument names whose default should be excluded from the C++ interface.\\n # Intended for resolving overload ambiguities between signatures.\\n cpp_no_default_args: set[str]\\n\\n # Note [Abstract ATen methods]\\n # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n # An abstract ATen method is one whose dispatch differs between\\n # types. These are implemented in derived types (with a\\n # standard (throwing) definition in Type). A concrete ATen\\n # method is one which has the same dispatch for all types;\\n # we just implement it in the base Type. This is exposed\\n # in Declarations.yaml via a field named 'abstract'.\\n is_abstract: bool\\n\\n # Whether or not the NativeFunction contains a backend-agnostic kernel\\n has_composite_implicit_autograd_kernel: bool\\n has_composite_implicit_autograd_nested_tensor_kernel: bool\\n has_composite_explicit_autograd_kernel: bool\\n has_composite_explicit_autograd_non_functional_kernel: bool\\n\\n # Tags are used to describe semantic information about (groups of) operators,\\n # That aren't easily inferrable directly from the operator's schema.\\n tags: set[str]\\n\\n # NB: The benefit of defining a dataclass is that we automatically get\\n # a constructor defined for all the fields we specify. No need\\n # to explicitly write it out.\\n\\n # We parse both the NativeFunction + backend-specific information about it, which it stored in a corresponding BackendIndex.\\n @staticmethod\\n def from_yaml(\\n ei: dict[str, object],\\n loc: Location,\\n valid_tags: set[str],\\n ignore_keys: set[DispatchKey] | None = None,\\n ) -> tuple[NativeFunction, dict[DispatchKey, dict[OperatorName, BackendMetadata]]]:\\n \\\"\\\"\\\"\\n Parse a NativeFunction from a dictionary as directly parsed\\n from native_functions.yaml\\n \\\"\\\"\\\"\\n e = ei.copy()\\n\\n funcs = e.pop(\\\"func\\\")\\n assert isinstance(funcs, str), f\\\"not a str: {funcs}\\\"\\n # only support one level of namespace. E.g., aten::add\\n namespace_helper = NamespaceHelper.from_namespaced_entity(\\n namespaced_entity=funcs, max_level=1\\n )\\n namespace = namespace_helper.get_cpp_namespace(default=\\\"aten\\\")\\n func = FunctionSchema.parse(namespace_helper.entity_name)\\n\\n cpp_no_default_args_list = e.pop(\\\"cpp_no_default_args\\\", [])\\n assert isinstance(cpp_no_default_args_list, list)\\n cpp_no_default_args = set(cpp_no_default_args_list)\\n\\n use_const_ref_for_mutable_tensors = e.pop(\\n \\\"use_const_ref_for_mutable_tensors\\\", False\\n )\\n assert isinstance(use_const_ref_for_mutable_tensors, bool)\\n\\n variants_s = e.pop(\\\"variants\\\", \\\"function\\\")\\n assert isinstance(variants_s, str)\\n variants: set[Variant] = set()\\n for v in variants_s.split(\\\", \\\"):\\n if v == \\\"function\\\":\\n variants.add(Variant.function)\\n elif v == \\\"method\\\":\\n variants.add(Variant.method)\\n else:\\n raise AssertionError(f\\\"illegal variant {v}\\\")\\n\\n manual_kernel_registration = e.pop(\\\"manual_kernel_registration\\\", False)\\n assert isinstance(\\n manual_kernel_registration, bool\\n ), f\\\"not a bool: {manual_kernel_registration}\\\"\\n\\n manual_cpp_binding = e.pop(\\\"manual_cpp_binding\\\", False)\\n assert isinstance(manual_cpp_binding, bool), f\\\"not a bool: {manual_cpp_binding}\\\"\\n\\n device_guard = e.pop(\\\"device_guard\\\", True)\\n assert isinstance(device_guard, bool), f\\\"not a bool: {device_guard}\\\"\\n\\n device_check_s = e.pop(\\\"device_check\\\", None)\\n assert device_check_s is None or isinstance(\\n device_check_s, str\\n ), f\\\"not a str: {device_check_s}\\\"\\n assert (\\n device_check_s is None or device_check_s in DeviceCheckType.__members__\\n ), f\\\"illegal device_check: {device_check_s}\\\"\\n device_check: DeviceCheckType\\n if device_check_s is None:\\n device_check = DeviceCheckType.ExactSame\\n else:\\n device_check = DeviceCheckType[device_check_s]\\n\\n structured = e.pop(\\\"structured\\\", False)\\n assert isinstance(structured, bool), f\\\"not a bool: {structured}\\\"\\n\\n structured_delegate_s = e.pop(\\\"structured_delegate\\\", None)\\n assert structured_delegate_s is None or isinstance(\\n structured_delegate_s, str\\n ), f\\\"not a str: {structured_delegate_s}\\\"\\n assert structured_delegate_s is None or \\\"::\\\" not in structured_delegate_s, (\\n \\\"namespace is not supported in structured delegate,\\\"\\n \\\" using the same namespace as the native function\\\"\\n )\\n structured_delegate: OperatorName | None = None\\n if structured_delegate_s is not None:\\n structured_delegate = OperatorName.parse(structured_delegate_s)\\n\\n structured_inherits = e.pop(\\\"structured_inherits\\\", None)\\n assert structured_inherits is None or isinstance(\\n structured_inherits, str\\n ), f\\\"not a str: {structured_inherits}\\\"\\n assert structured_inherits is None or \\\"::\\\" not in structured_inherits, (\\n \\\"namespace is not supported in structured inherits,\\\"\\n \\\" using the same namespace as the native function\\\"\\n )\\n\\n python_module = e.pop(\\\"python_module\\\", None)\\n assert python_module is None or isinstance(\\n python_module, str\\n ), f\\\"not a str: {python_module}\\\"\\n assert (\\n python_module is None or Variant.method not in variants\\n ), \\\"functions in modules cannot be methods\\\"\\n\\n category_override = e.pop(\\\"category_override\\\", None)\\n assert category_override is None or isinstance(\\n category_override, str\\n ), f\\\"not a str: {category_override}\\\"\\n\\n precomputed_dict = e.pop(\\\"precomputed\\\", None)\\n assert precomputed_dict is None or structured is True\\n precomputed = Precompute.parse(precomputed_dict) if precomputed_dict else None\\n\\n tags_inp = e.pop(\\\"tags\\\", [])\\n if isinstance(tags_inp, str):\\n tags_inp = [tags_inp]\\n assert isinstance(tags_inp, list)\\n\\n # All aten ops generated by torchgen receive the pt2_compliant tag.\\n if namespace == \\\"aten\\\" and \\\"pt2_compliant_tag\\\" in valid_tags:\\n tags_inp.append(\\\"pt2_compliant_tag\\\")\\n\\n tags: set[str] = set()\\n for t in tags_inp:\\n assert len(valid_tags) > 0\\n # TODO: verify that the tag is valid and has an entry in tags.yaml\\n if t in valid_tags:\\n tags.add(t)\\n else:\\n raise AssertionError(f\\\"illegal tag {t}\\\")\\n\\n from torchgen.api import cpp\\n\\n raw_dispatch = e.pop(\\\"dispatch\\\", None)\\n assert raw_dispatch is None or isinstance(raw_dispatch, dict), e\\n dispatch: dict[DispatchKey, BackendMetadata] = {}\\n num_dispatch_keys: int = 0\\n if raw_dispatch is not None:\\n assert not manual_kernel_registration, (\\n \\\"cannot specify both manual_kernel_registration and dispatch; with \\\"\\n \\\"manual registration, dispatch has no effect!\\\"\\n )\\n redundant_composite_implicit_autograd = False\\n for ks, v in raw_dispatch.items():\\n if ks == \\\"__line__\\\":\\n continue # not worth tracking line numbers for dispatch entries\\n assert isinstance(\\n ks, str\\n ), f\\\"illegal dispatch key '{ks}' in {raw_dispatch}\\\"\\n assert isinstance(\\n v, str\\n ), f\\\"illegal dispatch value '{v}' in {raw_dispatch}\\\"\\n for k in ks.split(\\\",\\\"):\\n dispatch_key = DispatchKey.parse(k.strip())\\n num_dispatch_keys += 1\\n\\n if ignore_keys and dispatch_key in ignore_keys:\\n continue\\n assert dispatch_key in dispatch_keys, (\\n f\\\"Dispatch key {dispatch_key} of kernel {v} \\\"\\n \\\"is not a supported dispatch key.\\\"\\n )\\n # We only allow at most 3 levels of namespace for kernels.\\n # We will append \\\"native\\\" to a custom kernel namespace.\\n namespace_helper = NamespaceHelper.from_namespaced_entity(\\n v, max_level=3\\n )\\n kernel_namespace = namespace_helper.get_cpp_namespace(default=\\\"at\\\")\\n # Why is 'structured' included? External backends (e.g.\\n # XLA) opt into which ops are structured independently\\n # of which in-tree ops are structured\\n dispatch[dispatch_key] = BackendMetadata(\\n kernel=namespace_helper.entity_name,\\n structured=structured\\n and is_structured_dispatch_key(dispatch_key),\\n cpp_namespace=(kernel_namespace + \\\"::native\\\"),\\n )\\n if (\\n dispatch_key is DispatchKey.CompositeImplicitAutograd\\n and v == cpp.name(func)\\n ):\\n redundant_composite_implicit_autograd = True\\n\\n # We count the number of dispatch keys which have not been ignored to prevent a dispatch table\\n # in which all backend keys are ignored but necessarily kept, remaining compositeimplicit,\\n # from being treated as redundant.\\n assert not (\\n num_dispatch_keys == 1 and redundant_composite_implicit_autograd\\n ), (\\n \\\"unnecessary dispatch table for this function; just delete the dispatch \\\"\\n \\\"key entirely\\\"\\n )\\n # if a function is a structured delegate, deleting the dispatch\\n # table is NOT semantics preserving\\n assert (\\n structured_delegate\\n or dispatch.keys() != {DispatchKey.CompositeImplicitAutograd}\\n or dispatch[DispatchKey.CompositeImplicitAutograd].supports_symint()\\n or num_dispatch_keys != 1\\n ), (\\n f\\\"unexpected name for singleton CompositeImplicitAutograd dispatch entry: expected {cpp.name(func)} \\\"\\n f\\\"but got {dispatch[DispatchKey.CompositeImplicitAutograd]}. Rename your implementation to the expected \\\"\\n \\\"name, then delete the dispatch table\\\"\\n )\\n elif not structured and structured_delegate is None:\\n name = str(func.name.name)\\n assert not (\\n name.startswith(\\\"new_\\\")\\n or name.endswith(\\\"_like\\\")\\n # TODO: maybe it's better to test the return\\n or (\\n func.arguments.tensor_options\\n and not func.arguments.has_tensor_arg()\\n )\\n ), (\\n f\\\"expected {name} to have a CompositeExplicitAutograd \\\"\\n \\\"dispatch entry, but there was no dispatch table. Factory functions \\\"\\n \\\"should not have implicit dispatch as they should not be decomposed \\\"\\n \\\"for __torch_dispatch__\\\"\\n )\\n dispatch[DispatchKey.CompositeImplicitAutograd] = BackendMetadata(\\n cpp.name(func), structured=False, cpp_namespace=DEFAULT_KERNEL_NAMESPACE\\n )\\n\\n composites_in_dispatch = [\\n d\\n for d in dispatch\\n if d == DispatchKey.CompositeExplicitAutograd\\n or d == DispatchKey.CompositeExplicitAutogradNonFunctional\\n or d == DispatchKey.CompositeImplicitAutograd\\n or d == DispatchKey.CompositeImplicitAutogradNestedTensor\\n ]\\n\\n assert len(composites_in_dispatch) <= 1 or (\\n len(composites_in_dispatch) == 2\\n and (\\n DispatchKey.CompositeExplicitAutogradNonFunctional\\n not in composites_in_dispatch\\n )\\n and (\\n DispatchKey.CompositeImplicitAutogradNestedTensor\\n in composites_in_dispatch\\n )\\n ), (\\n \\\"cannot specify more than one of CompositeExplicitAutograd, CompositeExplicitAutogradNonFunctional, \\\"\\n \\\"or CompositeImplicitAutograd on a single kernel; each \\\"\\n \\\"strictly subsumes the other. If you wanted to provide an explicit autograd \\\"\\n \\\"implementation, specify CompositeExplicitAutograd; otherwise specify CompositeImplicitAutograd only\\\"\\n )\\n\\n autogen_str = e.pop(\\\"autogen\\\", \\\"\\\")\\n assert isinstance(autogen_str, str)\\n autogen = (\\n []\\n if autogen_str == \\\"\\\"\\n else [OperatorName.parse(x) for x in autogen_str.split(\\\", \\\")]\\n )\\n\\n raw_ufunc_inner_loop = e.pop(\\\"ufunc_inner_loop\\\", {})\\n ufunc_inner_loop = {}\\n if isinstance(raw_ufunc_inner_loop, str):\\n ufunc_inner_loop[UfuncKey.Generic] = UfuncInnerLoop.parse(\\n raw_ufunc_inner_loop, UfuncKey.Generic\\n )\\n elif isinstance(raw_ufunc_inner_loop, dict):\\n for k, vo in raw_ufunc_inner_loop.items():\\n if k == \\\"__line__\\\":\\n continue\\n assert isinstance(k, str), f\\\"ufunc_inner_loop key is not a str: {k}\\\"\\n assert isinstance(vo, str), f\\\"ufunc_inner_loop value is not a str: {v}\\\"\\n ufunc_key = UfuncKey.parse(k)\\n ufunc_inner_loop[ufunc_key] = UfuncInnerLoop.parse(vo, ufunc_key)\\n else:\\n raise AssertionError(\\n f\\\"ufunc_inner_loop not str or dict: {raw_ufunc_inner_loop}\\\"\\n )\\n # Program the BackendIndex for the implicit dispatch entry from ufunc\\n if ufunc_inner_loop:\\n assert structured, \\\"ufunc must be structured\\\"\\n\\n # Delay import ufunc here to avoid circular import issue\\n # See: https://github.com/pytorch/pytorch/issues/81294\\n import torchgen.api.ufunc as ufunc\\n\\n for dispatch_key in UFUNC_DISPATCH_KEYS:\\n assert (\\n dispatch_key not in dispatch\\n ), f\\\"ufunc should not have explicit dispatch entry for {dispatch_key}\\\"\\n dispatch[dispatch_key] = BackendMetadata(\\n kernel=ufunc.schema_kernel_name(func, dispatch_key),\\n structured=True,\\n cpp_namespace=DEFAULT_KERNEL_NAMESPACE,\\n )\\n\\n if structured_delegate:\\n # Structured functions MUST have a dispatch table\\n is_abstract = True\\n else:\\n is_abstract = (\\n dispatch.keys() != {DispatchKey.CompositeImplicitAutograd}\\n and dispatch.keys()\\n != {DispatchKey.CompositeImplicitAutogradNestedTensor}\\n and dispatch.keys()\\n != {\\n DispatchKey.CompositeImplicitAutograd,\\n DispatchKey.CompositeImplicitAutogradNestedTensor,\\n }\\n )\\n\\n has_composite_implicit_autograd_kernel = (\\n DispatchKey.CompositeImplicitAutograd in dispatch\\n )\\n has_composite_implicit_autograd_nested_tensor_kernel = (\\n DispatchKey.CompositeImplicitAutogradNestedTensor in dispatch\\n )\\n has_composite_explicit_autograd_kernel = (\\n DispatchKey.CompositeExplicitAutograd in dispatch\\n )\\n has_composite_explicit_autograd_non_functional_kernel = (\\n DispatchKey.CompositeExplicitAutogradNonFunctional in dispatch\\n )\\n\\n # We aren't going to store dispatch metadata inline in NativeFunctions;\\n # instead it is separately indexed by backend (so other backends can\\n # add more dispatch entries after the fact). Reindex the individual\\n # metadata by OperatorName!\\n backend_metadata = {k: {func.name: v} for k, v in dispatch.items()}\\n\\n # don't care if it exists or not; make it easier to use this function\\n # with other yaml parsers that aren't setting __line__ in the dict\\n e.pop(\\\"__line__\\\", None)\\n assert not e, f\\\"leftover entries: {e}\\\"\\n\\n # Asserts that we can't do in post_init, because they rely on backend-specific info\\n if structured_delegate is not None:\\n for key in STRUCTURED_DISPATCH_KEYS:\\n assert key not in dispatch, (\\n f\\\"if structured_delegate, then must not have {key} in dispatch dictionary \\\"\\n \\\"(it is delegated!)\\\"\\n )\\n\\n return (\\n NativeFunction(\\n func=func,\\n use_const_ref_for_mutable_tensors=use_const_ref_for_mutable_tensors,\\n variants=variants,\\n structured=structured,\\n structured_delegate=structured_delegate,\\n structured_inherits=structured_inherits,\\n precomputed=precomputed,\\n autogen=autogen,\\n ufunc_inner_loop=ufunc_inner_loop,\\n manual_kernel_registration=manual_kernel_registration,\\n manual_cpp_binding=manual_cpp_binding,\\n python_module=python_module,\\n category_override=category_override,\\n device_guard=device_guard,\\n device_check=device_check,\\n loc=loc,\\n cpp_no_default_args=cpp_no_default_args,\\n is_abstract=is_abstract,\\n has_composite_implicit_autograd_kernel=has_composite_implicit_autograd_kernel,\\n has_composite_implicit_autograd_nested_tensor_kernel=has_composite_implicit_autograd_nested_tensor_kernel,\\n has_composite_explicit_autograd_kernel=has_composite_explicit_autograd_kernel,\\n has_composite_explicit_autograd_non_functional_kernel=has_composite_explicit_autograd_non_functional_kernel,\\n tags=tags,\\n namespace=namespace,\\n ),\\n backend_metadata,\\n )\\n\\n def validate_unstructured(self) -> None:\\n # TODO: probably better to accumulate these errors and report them all\\n # at once\\n assert not self.structured, (\\n \\\"This function is structured, but there was \\\"\\n \\\"no valid functional variant of it.\\\"\\n )\\n assert self.structured_delegate, (\\n \\\"This function delegates to another structured out function, \\\"\\n \\\"but no valid function was found (the delegate may not exist, or it has the wrong type)\\\"\\n )\\n\\n # __post_init__ functions in dataclasses can be used to do extra\\n # validation after construction.\\n #\\n # Notice that we don't do any type validation here. In fact, we\\n # rely exclusively on mypy to check if you've done types correctly!\\n # Validation is for nontrivial invariants that cannot be (conveniently)\\n # encoded in the type system.\\n def __post_init__(self) -> None:\\n if self.func.arguments.out:\\n assert self.variants == {Variant.function}, (\\n \\\"Native functions with out arguments MUST \\\"\\n \\\"be declared with only function variant; e.g., variants: function; \\\"\\n \\\"otherwise you will tickle a Python argument binding bug \\\"\\n \\\"(which usually manifests itself as the result variable being undefined.)\\\"\\n )\\n if self.structured:\\n assert self.func.kind() == SchemaKind.out, (\\n \\\"Put structured field on the out= \\\"\\n \\\"variant of a function; did you mean structured_delegate?\\\"\\n )\\n assert (\\n self.device_guard\\n ), \\\"device_guard: False is not respected by structured kernels\\\"\\n if self.structured_delegate:\\n assert self.func.kind() != SchemaKind.out, (\\n \\\"structured_delegate field not allowed \\\"\\n \\\"on out= functions; did you mean structured?\\\"\\n )\\n assert (\\n self.device_guard\\n ), \\\"device_guard: False is not respected by structured kernels\\\"\\n # Technically, with the asserts above, this assert is impossible to\\n # happen\\n assert not (\\n self.structured and self.structured_delegate\\n ), \\\"Cannot have both structured and structured_delegate on function\\\"\\n defaulted_arguments = {\\n a.name for a in self.func.schema_order_arguments() if a.default is not None\\n }\\n invalid_args = set.difference(self.cpp_no_default_args, defaulted_arguments)\\n assert len(invalid_args) == 0, f\\\"Invalid cpp_no_default_args: {invalid_args}\\\"\\n if self.structured_inherits is not None:\\n assert (\\n self.structured\\n ), \\\"structured_inherits must also imply structured: True\\\"\\n if str(self.func.name).startswith(\\\"_foreach\\\"):\\n assert self.device_check == DeviceCheckType.NoCheck, (\\n \\\"foreach kernels fall back to slow path when tensor are on different devices, \\\"\\n \\\"device_check not allowed to be enabled\\\"\\n )\\n\\n # NB: if your function accidentally has rand/dropout/... in its name\\n # but is not actually random, feel free to amend this to special case\\n if (\\n \\\"rand\\\" in str(self.func.name)\\n or (\\n (\\n \\\"dropout\\\" in str(self.func.name)\\n or any(\\n \\\"dropout\\\" in arg.name for arg in self.func.arguments.flat_all\\n )\\n )\\n # Backwards of dropout is typically deterministic\\n and \\\"backward\\\" not in str(self.func.name)\\n and str(self.func.name.name) not in [\\\"_cudnn_init_dropout_state\\\"]\\n )\\n or self.func.arguments.has_generator_arg()\\n ):\\n assert \\\"nondeterministic_seeded\\\" in self.tags, str(self.func.name)\\n\\n @property\\n def has_composite_kernel(self) -> bool:\\n return (\\n self.has_composite_implicit_autograd_kernel\\n or self.has_composite_explicit_autograd_kernel\\n or self.has_composite_explicit_autograd_non_functional_kernel\\n ) or (\\n self.has_composite_implicit_autograd_kernel\\n and self.has_composite_implicit_autograd_nested_tensor_kernel\\n )\\n\\n @property\\n def is_view_op(self) -> bool:\\n rets = self.func.returns\\n is_non_mutating_view = len(rets) > 0 and any(\\n r.annotation is not None and not r.annotation.is_write for r in rets\\n )\\n # See Note [resize_ in Functionalization] for more dtails\\n is_inplace_view = (\\n \\\"inplace_view\\\" in self.tags\\n and str(self.func.name) != \\\"resize_\\\"\\n and str(self.func.name) != \\\"resize_as_\\\"\\n )\\n is_wildcard_view = any(\\n inp.annotation is not None and \\\"*\\\" in inp.annotation.alias_set_after\\n for inp in self.func.schema_order_arguments()\\n )\\n return is_non_mutating_view or is_inplace_view or is_wildcard_view\\n\\n @property\\n def view_schema_kind(self) -> ViewSchemaKind:\\n if self.is_view_op and self.func.name.name.inplace:\\n assert \\\"inplace_view\\\" in self.tags\\n return ViewSchemaKind.aliasing_inplace\\n if self.is_view_op:\\n return ViewSchemaKind.aliasing\\n else:\\n return ViewSchemaKind.non_aliasing\\n\\n @property\\n def root_name(self) -> str:\\n return self.func.name.name.base\\n\\n @property\\n def part_of_structured_group(self) -> bool:\\n return self.structured or self.structured_delegate is not None\\n\\n\\nclass SchemaKind(Enum):\\n functional = auto()\\n inplace = auto()\\n out = auto()\\n mutable = auto()\\n scratch = auto()\\n\\n\\n# A structured kernel is guaranteed to have a functional and out variant, and\\n# optionally an inplace variant.\\n#\\n# NB: we create NativeFunctionsGroup *even if* the function is not\\n# actually annotated structured. Test the structured boolean to see if it\\n# actually is structured or not.\\n@dataclass(frozen=True)\\nclass NativeFunctionsGroup:\\n functional: NativeFunction\\n inplace: NativeFunction | None\\n mutable: NativeFunction | None\\n out: NativeFunction\\n\\n @property\\n def structured(self) -> bool:\\n # Whether or not the operator has a meta() function. This information is backend-agnostic.\\n return self.out.structured\\n\\n def __post_init__(self) -> None:\\n test_sig: FunctionSchema = self.functional.func.signature()\\n for f in self.functions():\\n if test_sig != f.func.signature():\\n raise AssertionError(\\n \\\"NativeFunctionsGroup constructed from two NativeFunctions \\\"\\n f\\\"that don't have matching signatures: {test_sig} != {f.func.signature()}\\\"\\n )\\n\\n if self.structured != f.part_of_structured_group:\\n raise AssertionError(\\n \\\"NativeFunctionsGroup constructed from structured and unstructured \\\"\\n f\\\"functions: {self.out.func.name} and {f.func.name}\\\"\\n )\\n assert self.functional.func.kind() == SchemaKind.functional\\n assert self.out.func.kind() == SchemaKind.out\\n assert self.functional.namespace == self.out.namespace\\n if self.inplace is not None:\\n assert self.inplace.func.kind() == SchemaKind.inplace\\n assert self.inplace.namespace == self.functional.namespace\\n\\n if self.mutable is not None:\\n assert self.mutable.func.kind() == SchemaKind.mutable\\n assert self.mutable.namespace == self.functional.namespace\\n # See Note [Overload Ambiguity With Functional Variants]\\n assert self.functional.func.name.name.functional_overload\\n\\n if self.structured:\\n # For now, structured composite kernels are not supported (need some\\n # design work to figure out how to make the composite case work)\\n assert (\\n not self.out.has_composite_implicit_autograd_kernel\\n and not self.out.has_composite_implicit_autograd_nested_tensor_kernel\\n )\\n\\n assert self.functional.structured_delegate == self.out.func.name, (\\n f\\\"{self.functional.func.name} delegates to {self.functional.structured_delegate} \\\"\\n f\\\"but its actual delegate is {self.out.func.name}\\\"\\n )\\n if self.inplace is not None:\\n assert self.inplace.structured_delegate == self.out.func.name\\n\\n generated_fns = sorted(\\n [str(f.func.name) for f in self.functions() if \\\"generated\\\" in f.tags]\\n )\\n generated_fns_str = \\\", \\\".join(str(x) for x in generated_fns)\\n expected_generated_fns: set[str] = set()\\n for f in self.functions():\\n expected_generated_fns.update(str(op) for op in f.autogen)\\n expected_generated_fns_str = \\\", \\\".join(\\n str(x) for x in sorted(expected_generated_fns)\\n )\\n if len(expected_generated_fns) == 0 and len(generated_fns) > 0:\\n raise RuntimeError(\\n f\\\"The codegen expects to be able to generate '{generated_fns_str}'.\\\"\\n \\\" In order to generate them however, we expect them to be called out explicitly in the yaml.\\\"\\n f\\\" Please add an 'autogen: {generated_fns_str}' line to the entry for {str(f.func.name)}\\\"\\n )\\n if expected_generated_fns_str != generated_fns_str:\\n raise RuntimeError(\\n f\\\"The codegen expects to be able to generate '{generated_fns_str}'.\\\"\\n f\\\" To do so, it expects a line: 'autogen: {generated_fns_str}'.\\\"\\n f\\\" Instead, it found 'autogen: {expected_generated_fns_str}'\\\"\\n )\\n\\n def signature(self) -> FunctionSchema:\\n return self.out.func.signature()\\n\\n def functions(self) -> Iterator[NativeFunction]:\\n yield self.functional\\n yield self.out\\n if self.inplace is not None:\\n yield self.inplace\\n if self.mutable is not None:\\n yield self.mutable\\n\\n @property\\n def root_name(self) -> str:\\n return self.functional.root_name\\n\\n @staticmethod\\n def from_dict(d: dict[SchemaKind, NativeFunction]) -> NativeFunctionsGroup | None:\\n assert d\\n if len(d) == 1:\\n return None\\n d = dict(d) # non-destructive updates please\\n functional = d.pop(SchemaKind.functional, None)\\n inplace = d.pop(SchemaKind.inplace, None)\\n mutable = d.pop(SchemaKind.mutable, None)\\n out = d.pop(SchemaKind.out, None)\\n assert not d\\n assert functional is not None\\n # There are a few operators which only have functional/inplace variants;\\n # these don't count as structured for our purposes here\\n if out is None:\\n return None\\n # assuming all variants have the same namespace\\n return NativeFunctionsGroup(\\n functional=functional,\\n inplace=inplace,\\n mutable=mutable,\\n out=out,\\n )\\n\\n\\n@dataclass(frozen=True)\\nclass BackendMetadata:\\n # The name of the backend kernel, for a given operator\\n # for in-tree backends. These names come directly from the 'dispatch\\\" field\\n # in native_functions.yaml. The dispatch entry is optional; in that\\n # case, that is equivalent to having written:\\n #\\n # dispatch:\\n # CompositeImplicitAutograd: $operator_name\\n kernel: str\\n # Whether or not the operator has a structured kernel implemented, for this particular backend.\\n # For in-tree backends, they all have the same value for structured- this is listed\\n # in native_functions.yaml.\\n # However, external backends like XLA can indendently toggle which ops are structured.\\n structured: bool\\n\\n # The namespace for kernels, default value: DEFAULT_KERNEL_NAMESPACE\\n cpp_namespace: str\\n\\n def supports_symint(self) -> bool:\\n return \\\"_symint\\\" in self.kernel\\n\\n\\n@dataclass(frozen=True)\\nclass UfuncInnerLoop:\\n name: str\\n supported_dtypes: OrderedSet[ScalarType]\\n # key is stored here because it affects the semantics of name,\\n # so its helpful to have them together for further processing\\n ufunc_key: UfuncKey\\n\\n @staticmethod\\n def parse(value: str, ufunc_key: UfuncKey) -> UfuncInnerLoop:\\n name, supported_dtypes_str = value.split(\\\" \\\", 1)\\n assert supported_dtypes_str[0] == \\\"(\\\"\\n assert supported_dtypes_str[-1] == \\\")\\\"\\n supported_dtypes: OrderedSet[ScalarType] = OrderedSet()\\n for k in supported_dtypes_str[1:-1].split(\\\", \\\"):\\n supported_dtypes |= ScalarType.parse_set(k)\\n return UfuncInnerLoop(\\n name=name, supported_dtypes=supported_dtypes, ufunc_key=ufunc_key\\n )\\n\\n\\n# BackendIndex represents a backend.\\n# The BackendIndex encodes per-operator information that is potentially different\\n# for each backend. The most obvious example is the name of the kernel\\n# (the 'dispatch' entry in native_functions.yaml).\\n# However, there can be other examples of different backends having different information.\\n# External backends can choose to opt their kernels to be structured independently from in-tree backends,\\n# which means that this information isn't inherently tied to a NativeFunction- it's different per backend.\\n@dataclass(frozen=True)\\nclass BackendIndex:\\n dispatch_key: DispatchKey\\n # Mainly important for structured kernels, this determines which variant in the operator group is used to implement the others.\\n # All in-tree ops use out kernels, while XLA uses functional kernels.\\n use_out_as_primary: bool\\n # Whether the backend requires a device guard, and device checks.\\n # For in-tree backends, this is currently just CUDA/HIP\\n # For out-of-tree backends, this is currently just Intel XPU\\n device_guard: bool\\n # Whether the backend is in-tree (CPU/CUDA) or out-of-tree (XLA)\\n external: bool\\n # Other backend-specific information that is on a per-operator basis\\n index: dict[OperatorName, BackendMetadata]\\n\\n @staticmethod\\n def grow_index(\\n parent_index: dict[DispatchKey, dict[OperatorName, BackendMetadata]],\\n child_index: dict[DispatchKey, dict[OperatorName, BackendMetadata]],\\n ) -> None:\\n for k, v in child_index.items():\\n for op_name, metadata in v.items():\\n assert (\\n op_name not in parent_index[k]\\n ), f\\\"duplicate operator {op_name} for dispatch key {k}\\\"\\n parent_index[k][op_name] = metadata\\n\\n def primary(self, g: NativeFunctionsGroup) -> NativeFunction:\\n if self.use_out_as_primary:\\n return g.out\\n else:\\n return g.functional\\n\\n def has_kernel(self, g: NativeFunction | NativeFunctionsGroup) -> bool:\\n m = self.get_kernel(g)\\n return m is not None\\n\\n def get_kernel(\\n self, g: NativeFunction | NativeFunctionsGroup\\n ) -> BackendMetadata | None:\\n if isinstance(g, NativeFunction):\\n f = g\\n elif isinstance(g, NativeFunctionsGroup):\\n f = self.primary(g)\\n else:\\n assert_never(g)\\n if f.func.name not in self.index:\\n return None\\n return self.index[f.func.name]\\n\\n def native_function_class_name(self) -> str | None:\\n if self.external:\\n return f\\\"{str(self.dispatch_key)}NativeFunctions\\\"\\n else:\\n # TODO: This discrepancy isn't required; we could also generated\\n # a class for in-tree kernels. It'll just require carefully\\n # updating every kernel definition + callsite of every in-tree aten kernel.\\n return None\\n\\n\\n# The function schema is undoubtedly the most important data structure\\n# in all of the codegen, as it defines the type signature for operators,\\n# and most of the code generation we do is type directed (e.g., look at\\n# the types, decide what to do. Think about how we code generate\\n# C++ function stubs!)\\n#\\n# We will also see in this class the general structure for how we model\\n# data in this code generation. A few notable properties to point out\\n# ahead of time:\\n#\\n# - These dataclasses are a *lossless* representation of the strings\\n# they are parsed from. In fact, we assert that given the\\n# information stored in the dataclass, we can exactly reconstruct\\n# the string we parsed from (and assert this inside the parse\\n# definition). There are a few reasons for this:\\n#\\n# - If you find that it is difficult to reconstruct the string\\n# given a dataclass, that is a clue that you are data\\n# representation is wrong.\\n#\\n# - It helps ensure that all relevant information is present\\n# in the dataclass, so that downstream users aren't tempted\\n# to reparse the original string to get some information\\n# that was omitted.\\n#\\n# - It forces you to represent the data in-memory in the same way\\n# it is recorded textually, which makes the dataclasses easier\\n# to understand for someone who is familiar with the\\n# textual format. (As a tradeoff, it means you have to model\\n# the syntax, even when it is inconvenient. But maybe that means\\n# the syntax is bad!) If you don't understand the internal\\n# representation, go look at the printing code to see how\\n# it maps onto the surface syntax!\\n#\\n# - It makes it easy to test the parsing code, as parsing code\\n# that is inconsistent with the string code will fail early\\n# and loudly. (As a tradeoff, it makes the parsing code a bit\\n# brittle (in particular, with trivial whitespace changes you\\n# are likely to trigger an assert error).\\n#\\n# In general, try to make the __str__ code as simple as possible\\n# (even at the cost of more complex parsing logic.) Additionally,\\n# try to minimize redundancy in data representation. (Precomputed\\n# fields are OK though: they are defined as a simple function on\\n# the canonical representation in question.)\\n#\\n# - These dataclasses are all frozen; once constructed their\\n# values never change. This makes it easy to tell where any\\n# given data came from: just look to the constructor. As a\\n# tradeoff, you can't easily \\\"decorate\\\" a schema with extra\\n# information from a post-facto analysis. We impose this\\n# restriction to make these structures more understandable.\\n#\\n@dataclass(frozen=True)\\nclass FunctionSchema:\\n # The name of the operator this function schema describes.\\n name: OperatorName\\n\\n arguments: Arguments\\n\\n # TODO: Need to handle collisions with argument names at some point\\n returns: tuple[Return, ...]\\n\\n @property\\n def is_mutable(self) -> bool:\\n def is_write(arg: Argument) -> bool:\\n if arg.annotation is None:\\n return False\\n return arg.annotation.is_write\\n\\n # Corresponds to torch._C._FunctionSchema.is_mutable\\n # See aten/src/ATen/core/function_schema.h (keep these in sync)\\n return any(is_write(a) for a in self.arguments.flat_all)\\n\\n def schema_order_arguments(self) -> Iterator[Argument]:\\n return itertools.chain(\\n self.arguments.flat_positional,\\n self.arguments.flat_kwarg_only,\\n self.arguments.out,\\n )\\n\\n decl_re = re.compile(r\\\"(?P[^\\\\(]+)\\\\((?P.*)\\\\) -> (?P.*)\\\")\\n\\n @staticmethod\\n def parse(func: str) -> FunctionSchema:\\n # We should probably get a proper parser here\\n decls = FunctionSchema.decl_re.findall(func)\\n assert len(decls) == 1, f\\\"Invalid function schema: {func}\\\"\\n ops, args, return_decl = decls[0]\\n name = OperatorName.parse(ops)\\n arguments = Arguments.parse(args)\\n returns = parse_returns(return_decl)\\n r = FunctionSchema(name=name, arguments=arguments, returns=returns)\\n assert str(r) == func, f\\\"{str(r)} != {func}\\\"\\n return r\\n\\n def returns_are_aliased(self) -> bool:\\n # We assert earlier that schemas can't have a mix of aliased and non-aliased returns\\n return any(\\n r\\n for r in self.returns\\n if r.annotation is not None and r.annotation.is_write\\n )\\n\\n def __post_init__(self) -> None:\\n for arg, ret in zip(self.arguments.out, self.returns):\\n assert arg.annotation == ret.annotation, (\\n \\\"Out arguments must have matching return Tensor; furthermore, \\\"\\n \\\"the ith-argument needs to correspond to the ith return\\\"\\n )\\n # We also enforce that if you have any mutable, positional args, then they are not returned.\\n # This makes it easier to group these functions properly with their functional/out= counterparts.\\n for a in self.arguments.post_self_positional_mutable:\\n assert not any(\\n a.annotation == r.annotation for r in self.returns\\n ), f\\\"If you have a schema with mutable positional args, we expect them to not be returned. schema: {str(self)}\\\"\\n # Invariant: we expect out arguments to appear as keyword arguments in the schema.\\n # This means that all mutable returns should be aliased to a keyword argument\\n # (except for \\\"self\\\", which we explicitly don't treat as an out argument because of its use in methods)\\n # See Note [is_out_fn]\\n out_and_self = list(self.arguments.out) + [\\n arg for arg in self.arguments.flat_positional if arg.name == \\\"self\\\"\\n ]\\n mutable_returns = [\\n ret\\n for ret in self.returns\\n if ret.annotation is not None and ret.annotation.is_write\\n ]\\n immutable_returns = [\\n ret\\n for ret in self.returns\\n if ret.annotation is None or not ret.annotation.is_write\\n ]\\n # Some assertions: We don't want any functions with a return type of \\\"-> (Tensor(a!), Tensor)\\\",\\n # because:\\n # (1) It's more annoying to handle properly\\n # (2) It's unnecessary - you can't method-chain on the first (mutated) output because it's part of a tuple.\\n # Instead, we expect the (a!) argument to not be returned.\\n assert (\\n len(mutable_returns) == 0 or len(immutable_returns) == 0\\n ), f\\\"NativeFunctions must have either only mutable returns, or only immutable returns. Found: {str(self)}\\\"\\n for ret in mutable_returns:\\n assert any(ret.annotation == arg.annotation for arg in out_and_self), (\\n 'All mutable returns must be aliased either to a keyword argument, or to \\\"self\\\". '\\n \\\"Did you forget to mark an out argument as keyword-only?\\\"\\n )\\n if self.arguments.out:\\n # out= ops that return their mutable inputs are only really useful for method chaining.\\n # And method chaining is only really useful if the thing you're returning is a plain Tensor.\\n # So ideally, we'd enforce that out= ops with a single plain mutable tensor should return the tensor,\\n # and all other types of out= op schemas should return void.\\n # There are a bunch of existing out= ops that return tuples of tensors though, so we're stuck with allowing that.\\n if any(a.type != BaseType(BaseTy.Tensor) for a in self.arguments.out):\\n assert (\\n len(self.returns) == 0\\n ), \\\"out= ops that accept tensor lists as out arguments \\\"\\n \\\"are expected to have no return type (since you can't do method chaining on them)\\\"\\n else:\\n # mutable keyword arguments whose name has _scratch_ prefix are\\n # scratch tensors for memory planning and should not be returned\\n assert len(\\n [\\n arg\\n for arg in self.arguments.out\\n if not arg.name.startswith(\\\"_scratch_\\\")\\n ]\\n ) == len(\\n self.returns\\n ), \\\"Must return as many arguments as there are out arguments, or no return at all\\\"\\n\\n if self.name.name.inplace:\\n self_a = self.arguments.self_arg\\n assert (\\n self_a\\n and self_a.argument.annotation\\n and self_a.argument.annotation.is_write\\n )\\n if self_a.argument.type == BaseType(BaseTy.Tensor):\\n # All inplace ops with an ordinary `Tensor self` argument should return self,\\n # to allow for method chaining.\\n assert (\\n len(self.returns) == 1\\n and self.returns[0].annotation == self_a.argument.annotation\\n )\\n else:\\n # You can't method chain on non-tensor self arguments though (like a List[Tensor])\\n # so in all other cases we expect the return type to be none.\\n assert len(self.returns) == 0\\n\\n if self.arguments.tensor_options is not None:\\n assert self.kind() == SchemaKind.functional, (\\n \\\"Found an operator that is not functional or out variant, but has tensor options arguments.\\\"\\n \\\"This is not allowed- tensor options arguments are only allowed for factory functions.\\\"\\n f\\\"schema: {str(self)}\\\"\\n )\\n if self.is_functional_fn():\\n assert self.kind() == SchemaKind.functional, (\\n \\\"Found an operator that is not functional, but its overload contains the string 'functional'.\\\"\\n \\\"This is a special keyword in the codegen, please use a different overload name.\\\"\\n f\\\"schema: {str(self)}\\\"\\n )\\n\\n def is_functional_fn(self) -> bool:\\n return \\\"functional\\\" in self.name.overload_name\\n\\n def is_out_fn(self) -> bool:\\n # Note [is_out_fn]\\n #\\n # out functions are the variants which take an explicit out= argument\\n # to populate into. We need to know if a schema corresponds to an\\n # out function for several reasons:\\n #\\n # - They codegen differently in C++ API\\n # - codegen to at::add_out rather than at::add\\n # - out argument is moved to front of C++ argument list\\n #\\n # out functions are DEFINED to be any function with a keyword-only\\n # argument that is mutable. In principle, this could lead to a\\n # false positive if you define a function that mutates a\\n # kwarg only argument, but this isn't the \\\"true\\\" output of this\\n # function. A more robust definition that would work in this\\n # case would also look at:\\n #\\n # - The output types. Out functions take in the arguments\\n # they mutate and then return them again; this is sort\\n # of \\\"definitionally\\\" what makes something an out function.\\n # Historically, we DO check this for consistency.\\n # - Correspondence with pure variant. An out function\\n # should have a signature equivalent to its pure variant,\\n # but just with extra kwargs for the output elements. This\\n # is difficult to actually check for and historically\\n # we only do this check in tools/\\n return bool(self.arguments.out)\\n\\n def kind(self) -> SchemaKind:\\n \\\"\\\"\\\"\\n What kind of schema is this? A functional schema is one\\n that returns a newly allocated output; an inplace schema\\n modifies the self argument inplace; an out schema writes\\n the result into an explicitly provided out argument.\\n \\\"\\\"\\\"\\n is_out = bool(self.arguments.out)\\n is_scratch = bool(\\n [arg for arg in self.arguments.out if arg.name.startswith(\\\"_scratch_\\\")]\\n )\\n is_inplace = self.name.name.inplace\\n is_mutable = any(\\n a.annotation is not None and a.annotation.is_write\\n for a in self.arguments.post_self_positional\\n )\\n assert not (is_out and is_inplace)\\n # out= and inplace schemas can also have post_self_positional mutable args,\\n # but we give precedence to out= and inplace when deciding the schema kind.\\n # Tradeoff: we probably don't want to have to teach codegen that looks at inplace ops\\n # to also worry about mutable post_self_positional arguments,\\n # but it seems like a much bigger lift to classify them has having a new schema kind.\\n # The number of ops that fit in this strange category is small enough that\\n # we can probably manually write code for them instead of forcing the codegen to handle them.\\n if is_inplace:\\n return SchemaKind.inplace\\n elif is_scratch:\\n assert (\\n is_out\\n ), \\\"invariant: all scratch operators are expected to be out= operators too\\\"\\n return SchemaKind.scratch\\n elif is_out:\\n assert (\\n not is_scratch\\n ), \\\"We should not categorize a scratch op as an out variant. Check if the order of if statements are expected!\\\"\\n return SchemaKind.out\\n elif is_mutable:\\n return SchemaKind.mutable\\n else:\\n return SchemaKind.functional\\n\\n # For every return:\\n # - If the return aliases an input, we return the input name\\n # - Otherwise, we return None.\\n # If return names were enforced to be consistent with aliasing information, then we wouldn't need this.\\n def aliased_return_names(self) -> list[str | None]:\\n outs: list[str | None] = []\\n for r in self.returns:\\n aliased_args = [\\n a\\n for a in self.arguments.flat_all\\n if a.annotation is not None and a.annotation == r.annotation\\n ]\\n if len(aliased_args) == 0:\\n outs.append(None)\\n elif len(aliased_args) == 1:\\n outs.append(aliased_args[0].name)\\n else:\\n aliased_names = \\\", \\\".join(a.name for a in aliased_args)\\n raise AssertionError(\\n f\\\"Found a return ({r.name})that aliases multiple inputs ({aliased_names})\\\"\\n )\\n return outs\\n\\n def signature(\\n self,\\n *,\\n strip_default: bool = False,\\n strip_view_copy_name: bool = False,\\n keep_return_names: bool = False,\\n ) -> FunctionSchema:\\n \\\"\\\"\\\"\\n Certain schemas are 'related', in that they are simply\\n inplace/out/functional versions of the same function. This method\\n factors these schemas into the \\\"core\\\" functional signature which\\n is equal across all versions.\\n\\n Here is what normalization happens to the schema to convert\\n it to a signature:\\n - The overload name is stripped (name is retained, since\\n it expresses semantic content about what the function does)\\n - Inplace is set False\\n - Out arguments are stripped\\n - Mutable post_self_positional args are converted to returns\\n - Mutability annotations are stripped (this is sound\\n because you cannot overload on mutability annotation)\\n - Return names are stripped since they are not overloadable and\\n some variants have return names but some not\\n - TensorOptions are dropped\\n because out= variants of factory functions don't include them\\n (and we want to be able to pair up factory functions with their out variants)\\n\\n Finally, we want to be able to pair up related \\\"view\\\" and their\\n corresponding \\\"view_copy\\\" operators. We do this by optionally\\n stripping the trailing \\\"_copy\\\" from the base name.\\n\\n Example of a mutable op before and after:\\n\\n f.func (Mutable operator):\\n _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask) # noqa: B950\\n\\n f.func (Corresponding functional operator):\\n _fused_moving_avg_obs_fq_helper.functional(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor running_min, Tensor running_max, Tensor scale, Tensor zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask, Tensor running_min_out, Tensor running_max_out, Tensor scale_out, Tensor zero_point_out) # noqa: B950\\n\\n f.func.signature() output:\\n _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor running_min, Tensor running_max, Tensor scale, Tensor zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) # noqa: B950\\n \\\"\\\"\\\"\\n\\n def strip_ret_annotation(r: Return) -> Return:\\n return Return(\\n name=r.name if keep_return_names else None,\\n type=r.type,\\n annotation=None,\\n )\\n\\n base_name = self.name.name.base\\n if strip_view_copy_name:\\n if base_name.endswith(\\\"_copy\\\"):\\n base_name = base_name.replace(\\\"_copy\\\", \\\"\\\")\\n elif base_name.endswith(\\\"_scatter\\\"):\\n base_name = base_name.replace(\\\"scatter\\\", \\\"inverse\\\")\\n\\n # find mutable inputs that are not originally returned, and convert them to returns\\n returns_from_mutable_inputs = tuple(\\n # When we're grouping functions we strip the return names,\\n # but when we're generating the actual functional variants then we follow\\n # a convention for what to name the returns\\n Return(\\n name=f\\\"{a.name}_out\\\" if keep_return_names else None,\\n type=a.type,\\n annotation=None,\\n )\\n for a in itertools.chain(\\n # Order is important here (otherwise e.g. inplace with mutable args\\n # and out= with mutable args won't have the same signature)\\n [self.arguments.self_arg.argument]\\n if self.arguments.self_arg is not None\\n else [],\\n self.arguments.out,\\n self.arguments.post_self_positional,\\n )\\n if a.annotation is not None\\n and a.annotation.is_write\\n and not any(a.annotation == r.annotation for r in self.returns)\\n )\\n original_returns = tuple(map(strip_ret_annotation, self.returns))\\n # Ordering is important here. We expect the \\\"mutable input\\\" returns to come last.\\n returns = original_returns + returns_from_mutable_inputs\\n\\n args_sig = self.arguments.signature(strip_default=strip_default)\\n # See Note [bernoulli.p schema]\\n if str(self.name) == \\\"bernoulli.p\\\":\\n args_sig = Arguments.parse(str(args_sig).replace(\\\"float p\\\", \\\"float p=0.5\\\"))\\n\\n return FunctionSchema(\\n name=OperatorName(\\n name=BaseOperatorName(\\n base=base_name,\\n inplace=False,\\n dunder_method=self.name.name.dunder_method,\\n ),\\n overload_name=\\\"\\\", # stripped\\n ),\\n arguments=args_sig,\\n returns=returns,\\n )\\n\\n def view_signature(self) -> FunctionSchema:\\n return self.signature(strip_view_copy_name=True)\\n\\n def with_name(self, name: OperatorName) -> FunctionSchema:\\n return FunctionSchema(\\n name=name,\\n arguments=self.arguments,\\n returns=self.returns,\\n )\\n\\n @property\\n def modifies_arguments(self) -> bool:\\n return self.kind() in [SchemaKind.inplace, SchemaKind.out, SchemaKind.mutable]\\n\\n def has_symint(self) -> bool:\\n return self.arguments.has_symint_arg()\\n\\n def __str__(self) -> str:\\n all_arguments_str = str(self.arguments)\\n if len(self.returns) == 1:\\n returns = str(self.returns[0]) # omit parentheses\\n else:\\n returns = \\\"(\\\" + \\\", \\\".join(map(str, self.returns)) + \\\")\\\"\\n return f\\\"{self.name}({all_arguments_str}) -> {returns}\\\"\\n\\n\\n# Here is the rest of the data model, described more briefly.\\n\\n\\n# Simplified version for what actually shows up in built-ins.\\n# Look at alias_info.h for expanded syntax. If you need the structure,\\n# you also need to make this structure recursive so it can be lined\\n# up with the type components too. For primitives this isn't really\\n# necessary\\n@dataclass(frozen=True)\\nclass Annotation:\\n # Typically only has one element. Not actually a set so\\n # we can conveniently assume it is canonically ordered\\n alias_set: tuple[str, ...]\\n is_write: bool\\n alias_set_after: tuple[str, ...]\\n\\n @staticmethod\\n def parse(ann: str) -> Annotation:\\n # TODO: implement a proper parser if this gets more ugly\\n # Regex Explanation:\\n # Example: \\\"a! -> a|b\\\"\\n # Group #1: alias before optional '|', required. Matches the first\\n # character 'a' in the example\\n # Group #2: optional alias set after optional '|', matches empty string\\n # in the example\\n # Group #3: optional \\\"is write\\\" flag, matches '!' in the example.\\n # Group #4: optional section containing arrow, matches \\\" -> a|b\\\" in the\\n # example.\\n # Group #5: optional alias after set, supports wildcard, matches \\\"a|b\\\"\\n # in the example.\\n # Group #6: optional sub-section of alias after set, matches \\\"|b\\\" in the\\n # example.\\n m = re.match(r\\\"^([a-z])(\\\\|[a-z])*(!?)( -> (\\\\*|[a-z](\\\\|[a-z])*))?$\\\", ann)\\n\\n assert m is not None, f\\\"unrecognized alias annotation {ann}\\\"\\n before_alias = m.group(1) + (m.group(2) if m.group(2) else \\\"\\\")\\n alias_set = tuple(before_alias.split(\\\"|\\\"))\\n is_write = m.group(3) == \\\"!\\\"\\n assert not (\\n is_write and len(alias_set) > 1\\n ), f\\\"alias set larger than 1 is not mutable, got {ann} instead.\\\"\\n after_set = tuple(m.group(5).split(\\\"|\\\")) if m.group(5) else ()\\n assert not (\\n len(before_alias) > 1 and len(after_set) > 1\\n ), f\\\"before alias set and after alias set cannot be larger than 1 at the same time, got {ann} instead.\\\"\\n r = Annotation(\\n alias_set=alias_set, is_write=is_write, alias_set_after=after_set\\n )\\n assert str(r) == ann, f\\\"{r} != {ann}\\\"\\n return r\\n\\n def __str__(self) -> str:\\n alias_set = \\\"|\\\".join(self.alias_set)\\n if self.is_write:\\n alias_set = f\\\"{alias_set}!\\\"\\n alias_set_after = \\\"|\\\".join(self.alias_set_after)\\n if alias_set_after:\\n alias_set = f'{alias_set}{\\\" -> \\\"}{alias_set_after}'\\n return alias_set\\n\\n\\n# The base class for the type system. This is also loosely modeled\\n# off of jit_type.h, but we've simplified the hierarchy to focus\\n# in on the aspects of the type system that matter for code generation\\n# (for example, there's no SingleElementType subclass anymore).\\n# You never actually construct a Type; usually it's going to be one\\n# of the subclasses. If Python had ADTs this would be one!\\n@dataclass(frozen=True)\\nclass Type:\\n @staticmethod\\n def parse(t: str) -> Type:\\n r = Type._parse(t)\\n assert str(r) == t, f\\\"{r} != {t}\\\"\\n return r\\n\\n @staticmethod\\n def _parse(t: str) -> Type:\\n m = re.match(r\\\"^(.+)\\\\?$\\\", t)\\n if m is not None:\\n return OptionalType(Type.parse(m.group(1)))\\n m = re.match(r\\\"^(.+)\\\\[([0-9]+)?\\\\]$\\\", t)\\n if m is not None:\\n size = int(m.group(2)) if m.group(2) is not None else None\\n return ListType(elem=Type.parse(m.group(1)), size=size)\\n\\n # '__torch__.torch.classes.' is the prefix for custom class\\n m = re.match(r\\\"^__torch__\\\\.torch\\\\.classes\\\\.([a-zA-Z0-9_.]+)$\\\", t)\\n if m is not None:\\n return CustomClassType(m.group(1))\\n try:\\n return BaseType(BaseTy[t])\\n except KeyError as e:\\n raise RuntimeError(f\\\"unrecognized type {t}\\\") from e\\n\\n def __str__(self) -> str:\\n raise NotImplementedError\\n\\n # WARNING: These concepts are not very well-defined. For example,\\n # is \\\"int?\\\" nullable? How about \\\"int?[]\\\". They are defined\\n # so we can conveniently generate legacy Declarations.yaml but\\n # really we should probably just remove these at some point\\n\\n def is_base_ty_like(self, base_ty: BaseTy) -> bool:\\n raise NotImplementedError\\n\\n def is_tensor_like(self) -> bool:\\n return self.is_base_ty_like(BaseTy.Tensor)\\n\\n def is_generator_like(self) -> bool:\\n return self.is_base_ty_like(BaseTy.Generator)\\n\\n def is_symint_like(self) -> bool:\\n return self.is_base_ty_like(BaseTy.SymInt)\\n\\n def is_nullable(self) -> bool:\\n raise NotImplementedError\\n\\n def is_list_like(self) -> ListType | None:\\n raise NotImplementedError\\n\\n\\n# Base types are simple, atomic types with no further structure\\nclass BaseTy(Enum):\\n Generator = auto()\\n ScalarType = auto()\\n Tensor = auto()\\n int = auto()\\n Dimname = auto()\\n DimVector = auto()\\n float = auto()\\n str = auto()\\n bool = auto()\\n Layout = auto()\\n Device = auto()\\n DeviceIndex = auto()\\n Scalar = auto()\\n MemoryFormat = auto()\\n QScheme = auto()\\n Storage = auto()\\n Stream = auto()\\n SymInt = auto()\\n SymBool = auto()\\n ConstQuantizerPtr = auto() # TODO: rename\\n GraphModule = auto()\\n\\n\\n@dataclass(frozen=True)\\nclass BaseType(Type):\\n name: BaseTy\\n\\n def __str__(self) -> str:\\n return f\\\"{self.name.name}\\\"\\n\\n def is_base_ty_like(self, base_ty: BaseTy) -> bool:\\n return self.name == base_ty\\n\\n def is_nullable(self) -> bool:\\n return False\\n\\n def is_list_like(self) -> ListType | None:\\n return None\\n\\n def is_symint_like(self) -> bool:\\n return self.name == BaseTy.SymInt\\n\\n\\n# Optional types may be specified, or may also be validly given None\\n@dataclass(frozen=True)\\nclass OptionalType(Type):\\n elem: Type\\n\\n def __str__(self) -> str:\\n return f\\\"{self.elem}?\\\"\\n\\n def is_base_ty_like(self, base_ty: BaseTy) -> bool:\\n return self.elem.is_base_ty_like(base_ty)\\n\\n def is_symint_like(self) -> bool:\\n return self.elem.is_symint_like()\\n\\n def is_nullable(self) -> bool:\\n return True\\n\\n def is_list_like(self) -> ListType | None:\\n return self.elem.is_list_like()\\n\\n\\n# A type representing a PyTorch custom class\\n@dataclass(frozen=True)\\nclass CustomClassType(Type):\\n class_name: str\\n\\n def __str__(self) -> str:\\n \\\"\\\"\\\"\\n Return the class name will prefix __torch__.torch.classes\\n \\\"\\\"\\\"\\n return f\\\"__torch__.torch.classes.{self.class_name}\\\"\\n\\n def is_base_ty_like(self, base_ty: BaseTy) -> bool:\\n return False\\n\\n def is_symint_like(self) -> bool:\\n return False\\n\\n def is_nullable(self) -> bool:\\n \\\"\\\"\\\"\\n Assume a custom class is not nullable.\\n \\\"\\\"\\\"\\n return False\\n\\n def is_list_like(self) -> ListType | None:\\n return None\\n\\n\\n# List types specify that we may have multiples of an element. We\\n# also support explicit sizes on list types, but these have\\n# some nontrivial semantics! (However, for C++ API purposes, explicit\\n# sizes are mostly erased from the type system.)\\n#\\n# DANGER WILL ROBINSON: C++ elaboration depends on elem type; e.g.,\\n# int[] elaborates differently than bool[3]!\\n@dataclass(frozen=True)\\nclass ListType(Type):\\n elem: Type\\n size: int | None\\n\\n def __str__(self) -> str:\\n size = f\\\"{self.size}\\\" if self.size else \\\"\\\"\\n return f\\\"{self.elem}[{size}]\\\"\\n\\n def is_base_ty_like(self, base_ty: BaseTy) -> bool:\\n return self.elem.is_base_ty_like(base_ty)\\n\\n def is_symint_like(self) -> bool:\\n return self.elem.is_symint_like()\\n\\n def is_nullable(self) -> bool:\\n return self.elem.is_nullable()\\n\\n def is_list_like(self) -> ListType | None:\\n return self\\n\\n\\n@dataclass(frozen=True)\\nclass Argument:\\n # NB: I didn't put kwarg_only as a boolean field here, unlike\\n # c10::Argument, so that printing works correctly\\n\\n name: str\\n type: Type\\n default: str | None\\n\\n # The semantics of the annotation field are a little strange.\\n #\\n # Alias annotations parametrize Tensors (since Tensors are the only things\\n # that can alias.) This motivates why I write Tensor(a!)? (and not, for\\n # example, Tensor?(a!)), because the (a!) describes aliasing on the tensor,\\n # which may be optional (i.e., the alias annotation should bind first to\\n # Tensor, before the optional postfix annotation).\\n #\\n # However, despite being a property of Tensor, we (and c10::Argument)\\n # store the annotation at the top level of the Argument, rather than\\n # inside the embedded Tensor type. In the C++ version of this\\n # class, we then go through great lengths to mimic the type\\n # structure in the annotation structure so we can correlate\\n # annotations with types.\\n #\\n # Now, it turns out, in all applications in code generation, the\\n # structure of annotated types is very simple. So we just hard\\n # code it here. But if we ever do get anything more complex, this\\n # model will have to change!\\n annotation: Annotation | None\\n\\n @property\\n def alias_info(self) -> Annotation | None:\\n return self.annotation\\n\\n @staticmethod\\n def parse(arg: str) -> Argument:\\n name: str\\n default: str | None\\n assert \\\" \\\" in arg, f\\\"illegal argument '{arg}'\\\"\\n if \\\"=\\\" in arg:\\n assert arg.count(\\\"=\\\") == 1, f\\\"illegal argument with default value: '{arg}'\\\"\\n type_and_annot_and_name, default = arg.split(\\\"=\\\")\\n type_and_annot, name = type_and_annot_and_name.rsplit(\\\" \\\", 1)\\n name_and_default = f\\\"{name}={default}\\\"\\n else:\\n type_and_annot, name_and_default = arg.rsplit(\\\" \\\", 1)\\n name = name_and_default\\n default = None\\n # TODO: deduplicate annotation matching with Return\\n match = re.match(r\\\"Tensor\\\\((.+)\\\\)(.*)\\\", type_and_annot)\\n annotation: Annotation | None\\n if match:\\n # If you update this, make sure the __str__ still works too\\n assert match.group(2) in [\\n \\\"\\\",\\n \\\"?\\\",\\n \\\"[]\\\",\\n ], \\\"unrecognized alias analysis form with Tensor\\\"\\n type_s = \\\"Tensor\\\" + match.group(2)\\n annotation = Annotation.parse(match.group(1))\\n else:\\n type_s = type_and_annot\\n annotation = None\\n type = Type.parse(type_s)\\n r = Argument(\\n name=name,\\n type=type,\\n default=default,\\n annotation=annotation,\\n )\\n assert str(r) == arg, f\\\"{str(r)} != {arg}\\\"\\n return r\\n\\n @property\\n def is_write(self) -> bool:\\n return self.annotation is not None and self.annotation.is_write\\n\\n def __str__(self) -> str:\\n type = f\\\"{self.type}\\\"\\n if self.annotation:\\n assert type in [\\\"Tensor\\\", \\\"Tensor?\\\", \\\"Tensor[]\\\"]\\n type = type.replace(\\\"Tensor\\\", f\\\"Tensor({self.annotation})\\\")\\n if self.name is None:\\n return type\\n else:\\n mb_default = \\\"\\\"\\n if self.default:\\n mb_default = f\\\"={self.default}\\\"\\n return f\\\"{type} {self.name}{mb_default}\\\"\\n\\n\\n@dataclass(frozen=True)\\nclass Return:\\n name: str | None\\n type: Type\\n annotation: Annotation | None\\n\\n @property\\n def alias_info(self) -> Annotation | None:\\n return self.annotation\\n\\n @staticmethod\\n def parse(arg: str) -> Return:\\n name: str | None\\n if \\\" \\\" in arg:\\n type_and_annot, name = arg.rsplit(\\\" \\\", 1)\\n else:\\n type_and_annot = arg\\n name = None\\n match = re.match(r\\\"Tensor\\\\((.+)\\\\)(.*)\\\", type_and_annot)\\n annotation: Annotation | None\\n if match:\\n # If you update this, make sure the __str__ still works too\\n assert match.group(2) in [\\n \\\"\\\",\\n \\\"?\\\",\\n \\\"[]\\\",\\n ], \\\"unrecognized alias analysis form with Tensor\\\"\\n type_s = \\\"Tensor\\\" + match.group(2)\\n annotation = Annotation.parse(match.group(1))\\n else:\\n type_s = type_and_annot\\n annotation = None\\n type = Type.parse(type_s)\\n r = Return(\\n name=name,\\n type=type,\\n annotation=annotation,\\n )\\n assert str(r) == arg, f\\\"{str(r)} != {arg}\\\"\\n return r\\n\\n @property\\n def is_write(self) -> bool:\\n return self.annotation is not None and self.annotation.is_write\\n\\n def __str__(self) -> str:\\n type = f\\\"{self.type}\\\"\\n if self.annotation:\\n assert type in [\\\"Tensor\\\", \\\"Tensor?\\\", \\\"Tensor[]\\\"]\\n type = type.replace(\\\"Tensor\\\", f\\\"Tensor({self.annotation})\\\")\\n if self.name is None:\\n return type\\n else:\\n return f\\\"{type} {self.name}\\\"\\n\\n\\n# Represents the self argument for functions that may be methods\\n@dataclass(frozen=True)\\nclass SelfArgument:\\n argument: Argument\\n\\n\\n# Bundle of arguments that represent a TensorOptions. This is mostly\\n# relevant for the public C++ API but we bake it into the core data\\n# model because other APIs often have to interact with it\\n@dataclass(frozen=True)\\nclass TensorOptionsArguments:\\n dtype: Argument\\n layout: Argument\\n device: Argument\\n pin_memory: Argument\\n\\n def all(self) -> Sequence[Argument]:\\n return [self.dtype, self.layout, self.device, self.pin_memory]\\n\\n\\n@dataclass(frozen=True)\\nclass Arguments:\\n # pre_self_positional is usually empty, but is notably non-empty\\n # for where.self, where the condition argument comes before the\\n # self argument\\n pre_self_positional: tuple[Argument, ...]\\n self_arg: SelfArgument | None\\n post_self_positional: tuple[Argument, ...]\\n\\n pre_tensor_options_kwarg_only: tuple[Argument, ...]\\n tensor_options: TensorOptionsArguments | None\\n # post_tensor_options is typically memory format, which should be\\n # part of tensor options but isn't right now, and is usually\\n # placed after the tensor options arguments\\n post_tensor_options_kwarg_only: tuple[Argument, ...]\\n\\n # Unlike in the previous codegen, we have factored out 'out' arguments\\n # in the canonical representation, removing them from kwarg\\n # arguments. This choice is justified by numerous downstream\\n # transformations which treat out arguments specially; additionally,\\n # you can see that canonicity is not violated!\\n out: tuple[Argument, ...] # these are also kwarg-only\\n\\n @property\\n def flat_non_out(self) -> Sequence[Argument]:\\n ret: list[Argument] = []\\n ret.extend(self.flat_positional)\\n ret.extend(self.flat_kwarg_only)\\n return ret\\n\\n @property\\n def flat_positional(self) -> Sequence[Argument]:\\n ret: list[Argument] = []\\n ret.extend(self.pre_self_positional)\\n if self.self_arg is not None:\\n ret.append(self.self_arg.argument)\\n ret.extend(self.post_self_positional)\\n return ret\\n\\n @property\\n def post_self_positional_mutable(self) -> Sequence[Argument]:\\n return [a for a in self.post_self_positional if a.is_write]\\n\\n # NB: doesn't contain out arguments\\n @property\\n def flat_kwarg_only(self) -> Sequence[Argument]:\\n ret: list[Argument] = []\\n ret.extend(self.pre_tensor_options_kwarg_only)\\n if self.tensor_options is not None:\\n ret.extend(self.tensor_options.all())\\n ret.extend(self.post_tensor_options_kwarg_only)\\n return ret\\n\\n @property\\n def flat_all(self) -> Sequence[Argument]:\\n ret: list[Argument] = []\\n ret.extend(self.flat_positional)\\n ret.extend(self.flat_kwarg_only)\\n ret.extend(self.out)\\n return ret\\n\\n @property\\n def non_out(\\n self,\\n ) -> Sequence[Argument | SelfArgument | TensorOptionsArguments]:\\n ret: list[Argument | SelfArgument | TensorOptionsArguments] = []\\n ret.extend(self.positional)\\n ret.extend(self.kwarg_only)\\n return ret\\n\\n @property\\n def positional(self) -> Sequence[Argument | SelfArgument]:\\n ret: list[Argument | SelfArgument] = []\\n ret.extend(self.pre_self_positional)\\n if self.self_arg is not None:\\n ret.append(self.self_arg)\\n ret.extend(self.post_self_positional)\\n return ret\\n\\n @property\\n def kwarg_only(self) -> Sequence[Argument | TensorOptionsArguments]:\\n ret: list[Argument | TensorOptionsArguments] = []\\n ret.extend(self.pre_tensor_options_kwarg_only)\\n if self.tensor_options is not None:\\n ret.append(self.tensor_options)\\n ret.extend(self.post_tensor_options_kwarg_only)\\n return ret\\n\\n @property\\n def all(self) -> Sequence[Argument | SelfArgument | TensorOptionsArguments]:\\n ret: list[Argument | SelfArgument | TensorOptionsArguments] = []\\n ret.extend(self.positional)\\n ret.extend(self.kwarg_only)\\n ret.extend(self.out)\\n return ret\\n\\n def mutable_arg_names(self) -> list[str]:\\n return [\\n a.name\\n for a in self.flat_all\\n if a.annotation is not None and a.annotation.is_write\\n ]\\n\\n def has_tensor_arg(self) -> bool:\\n return any(a.type.is_tensor_like() for a in self.flat_non_out)\\n\\n def has_symint_arg(self) -> bool:\\n return any(a.type.is_symint_like() for a in self.flat_non_out)\\n\\n def has_generator_arg(self) -> bool:\\n return any(a.type.is_generator_like() for a in self.flat_non_out)\\n\\n def signature(self, *, strip_default: bool = False) -> Arguments:\\n # dataclasses.replace could be used here, but it is less\\n # type safe so for now I've opted to type everything out\\n def strip_arg_annotation(a: Argument) -> Argument:\\n return Argument(\\n name=a.name,\\n type=a.type,\\n default=a.default if not strip_default else None,\\n annotation=None,\\n )\\n\\n return Arguments(\\n pre_self_positional=tuple(\\n map(strip_arg_annotation, self.pre_self_positional)\\n ),\\n self_arg=SelfArgument(strip_arg_annotation(self.self_arg.argument))\\n if self.self_arg is not None\\n else None,\\n post_self_positional=tuple(\\n map(strip_arg_annotation, self.post_self_positional)\\n ),\\n # Since TensorOptions are dropped, the post_tensor_options_kwargs are\\n # converted to pre_tensor_options_kwargs\\n pre_tensor_options_kwarg_only=tuple(\\n map(strip_arg_annotation, self.pre_tensor_options_kwarg_only)\\n )\\n + tuple(map(strip_arg_annotation, self.post_tensor_options_kwarg_only)),\\n # TensorOptions are dropped in signature,\\n # so we can pair factory functions with their out= variants.\\n tensor_options=None,\\n post_tensor_options_kwarg_only=(),\\n # out arguments are dropped in signature\\n out=(),\\n )\\n\\n def remove_self_annotation(self) -> Arguments:\\n assert self.self_arg is not None\\n return dataclasses.replace(\\n self,\\n self_arg=SelfArgument(\\n dataclasses.replace(self.self_arg.argument, annotation=None)\\n ),\\n )\\n\\n def with_out_args(self, outs: list[Argument]) -> Arguments:\\n assert len(self.out) == 0\\n return dataclasses.replace(\\n self,\\n out=tuple(outs),\\n )\\n\\n @staticmethod\\n def _preparse(args: str) -> tuple[list[Argument], list[Argument], list[Argument]]:\\n positional: list[Argument] = []\\n kwarg_only: list[Argument] = []\\n out: list[Argument] = []\\n arguments_acc = positional\\n\\n # TODO: Use a real parser here; this will get bamboozled\\n # by signatures that contain things like std::array (note the space)\\n for arg in args.split(\\\", \\\"):\\n if not arg:\\n continue\\n if arg == \\\"*\\\":\\n assert (\\n arguments_acc is positional\\n ), \\\"invalid syntax: kwarg-only specifier * can only occur once\\\"\\n arguments_acc = kwarg_only\\n continue\\n parg = Argument.parse(arg)\\n # Currently, we rely directly on the invariant that there are NO\\n # kwarg-only mutating arguments. If you want to relax this,\\n # we will need a more semantic way of matching that takes\\n # into account return arguments. In that case, you will have\\n # to manage out computation a level up, in FunctionSchema. See Note\\n # [is_out_fn]\\n if parg.annotation is not None and parg.annotation.is_write:\\n if arguments_acc is positional:\\n pass # do nothing\\n elif arguments_acc is kwarg_only:\\n arguments_acc = out\\n else:\\n assert arguments_acc is not out\\n arguments_acc.append(parg)\\n\\n return positional, kwarg_only, out\\n\\n @staticmethod\\n def parse(args: str) -> Arguments:\\n \\\"\\\"\\\"\\n Input: 'int x, int y, int z'\\n \\\"\\\"\\\"\\n\\n # We do this in two phases. First we parse into three\\n # main categories: positional, kwarg_only, out.\\n # Then, we reparse positional and kwarg_only to separate\\n # out the self argument and tensor options arguments.\\n\\n positional, kwarg_only, out = Arguments._preparse(args)\\n\\n # Split self argument\\n self_ix = None\\n for i, a in enumerate(positional):\\n if a.name == \\\"self\\\":\\n self_ix = i\\n break\\n pre_self_positional: list[Argument]\\n self_arg: SelfArgument | None\\n post_self_positional: list[Argument]\\n if self_ix is not None:\\n pre_self_positional = positional[:self_ix]\\n self_arg = SelfArgument(positional[self_ix])\\n post_self_positional = positional[self_ix + 1 :]\\n else:\\n pre_self_positional = []\\n self_arg = None\\n post_self_positional = positional\\n\\n # Group tensor options arguments\\n pre_tensor_options_kwarg_only: list[Argument] = []\\n tensor_options: TensorOptionsArguments | None = None\\n post_tensor_options_kwarg_only: list[Argument] = []\\n kwarg_only_acc = pre_tensor_options_kwarg_only\\n\\n def pred(name: str, ty: Type) -> Callable[[Argument], bool]:\\n return lambda a: a.name == name and a.type in [ty, OptionalType(ty)]\\n\\n predicates = [ # order matters\\n pred(\\\"dtype\\\", Type.parse(\\\"ScalarType\\\")),\\n pred(\\\"layout\\\", Type.parse(\\\"Layout\\\")),\\n pred(\\\"device\\\", Type.parse(\\\"Device\\\")),\\n pred(\\\"pin_memory\\\", Type.parse(\\\"bool\\\")),\\n ]\\n\\n i = 0\\n while i < len(kwarg_only):\\n # If there is enough space...\\n if i <= len(kwarg_only) - len(predicates):\\n # And the next len(predicates) arguments look like TensorOptions arguments\\n if all(\\n p(a)\\n for p, a in zip(predicates, kwarg_only[i : i + len(predicates)])\\n ):\\n assert kwarg_only_acc is pre_tensor_options_kwarg_only\\n # Group them together as one argument\\n tensor_options = TensorOptionsArguments(\\n dtype=kwarg_only[i],\\n layout=kwarg_only[i + 1],\\n device=kwarg_only[i + 2],\\n pin_memory=kwarg_only[i + 3],\\n )\\n i += len(predicates)\\n kwarg_only_acc = post_tensor_options_kwarg_only\\n continue\\n kwarg_only_acc.append(kwarg_only[i])\\n i += 1\\n\\n return Arguments(\\n pre_self_positional=tuple(pre_self_positional),\\n self_arg=self_arg,\\n post_self_positional=tuple(post_self_positional),\\n pre_tensor_options_kwarg_only=tuple(pre_tensor_options_kwarg_only),\\n tensor_options=tensor_options,\\n post_tensor_options_kwarg_only=tuple(post_tensor_options_kwarg_only),\\n out=tuple(out),\\n )\\n\\n def __str__(self) -> str:\\n all_arguments: list[str] = []\\n all_arguments.extend(map(str, self.flat_positional))\\n if self.flat_kwarg_only or self.out:\\n all_arguments.append(\\\"*\\\")\\n all_arguments.extend(map(str, self.flat_kwarg_only))\\n all_arguments.extend(map(str, self.out))\\n return \\\", \\\".join(all_arguments)\\n\\n def __post_init__(self) -> None:\\n # TODO: These invariants are weirdly asymmetric?\\n # TODO: Fancier types?\\n if self.self_arg is None:\\n assert not self.pre_self_positional\\n if self.tensor_options is None:\\n assert not self.post_tensor_options_kwarg_only\\n\\n # We don't allow any of the following to have argument annotations,\\n # to keep things simple.\\n mutable_pre_self_positionals = [\\n a\\n for a in self.pre_self_positional\\n if a.annotation is not None and a.annotation.is_write\\n ]\\n assert (\\n len(mutable_pre_self_positionals) == 0\\n ), \\\"mutable pre_self_positional arguments are not currently supported in the schema\\\"\\n\\n\\n# Names that validly are __iXXX__ indicating inplace operations.\\n# Taken from https://www.python.org/dev/peps/pep-0203/#new-methods\\n# NB: PyTorch hasn't actually implemented all of these\\nAUGMENTED_ASSIGNMENT_NAMES = [\\n \\\"add\\\",\\n \\\"sub\\\",\\n \\\"mul\\\",\\n \\\"div\\\",\\n \\\"mod\\\",\\n \\\"pow\\\",\\n \\\"lshift\\\",\\n \\\"rshift\\\",\\n \\\"and\\\",\\n \\\"xor\\\",\\n \\\"or\\\",\\n]\\n\\n\\n# A BaseOperatorName is what we think of the operator name, without\\n# the overload name. Unusually, we don't represent this as just a\\n# string; instead, we directly represent a few important semantic\\n# bits of information we derive from the string: namely whether\\n# or not it's inplace (add_) and whether or not it's a double-underscore\\n# method (__add__)\\n@dataclass(frozen=True)\\nclass BaseOperatorName:\\n base: str\\n inplace: bool\\n dunder_method: bool\\n # Note [Overload Ambiguity With Functional Variants]\\n # A handful of operators have both a \\\"mutable\\\" and a \\\"functional\\\" variant.\\n # (native_batch_norm is a good example, although this isn't the case today).\\n # For those operators, the mutable and functional variant take in the same set of\\n # arguments, but have different alias annotations.\\n # this makes it ambiguous when you try to resolve an OverloadPacket into an overload,\\n # given a set of input arguments.\\n #\\n # So instead of making the \\\"functional\\\" variant in this case a real overload, e.g:\\n # native_batch_norm (mutable variant)\\n # native_batch_norm.functional (functional variant)\\n # we make it a new base operator,\\n # native_batch_norm_functional (functional variant)\\n #\\n # In an ideal world, we would probably invert this so the operators were:\\n # native_batch_norm.mutable (mutable variant)\\n # native_batch_norm (functional variant)\\n #\\n # Doing that is BC-breaking though, so we're stuck with the above modeling.\\n functional_overload: bool = False\\n\\n @staticmethod\\n def parse(op: str) -> BaseOperatorName:\\n assert op != \\\"\\\"\\n assert not op.endswith(\\\"_out\\\"), (\\n \\\"_out suffix is reserved and not permitted for operator names; \\\"\\n \\\"did you mean to specify an out overload name instead?\\\"\\n )\\n m = re.match(r\\\"^__([^_]+)__$\\\", op)\\n if m is not None:\\n dunder_method = True\\n base = m.group(1)\\n if any(base == f\\\"i{n}\\\" for n in AUGMENTED_ASSIGNMENT_NAMES):\\n inplace = True\\n base = base[1:]\\n else:\\n inplace = False\\n # temporary, this is not intrinsically true but\\n # has been historically true for dunder methods\\n # we support (but, if we ever got, say, __int__, this would\\n # be wrong!)\\n assert base[0] != \\\"i\\\"\\n else:\\n dunder_method = False\\n base = op\\n if base[-1] == \\\"_\\\":\\n inplace = True\\n base = base[:-1]\\n else:\\n inplace = False\\n\\n # See Note [Overload Ambiguity With Functional Variants]\\n functional_suffix = \\\"_functional\\\"\\n if base.endswith(functional_suffix):\\n functional_overload = True\\n base = base[: -len(functional_suffix)]\\n # This seems complicated and unnecessary, so banning dunder methods\\n # for now on ops that have a functional + mutable variant (like native_batch_norm).\\n assert not dunder_method and not inplace\\n else:\\n functional_overload = False\\n\\n r = BaseOperatorName(\\n base=base,\\n inplace=inplace,\\n dunder_method=dunder_method,\\n functional_overload=functional_overload,\\n )\\n assert str(r) == op, f\\\"{str(r)} != {op}\\\"\\n return r\\n\\n def __str__(self) -> str:\\n if self.dunder_method:\\n i = \\\"i\\\" if self.inplace else \\\"\\\"\\n return f\\\"__{i}{self.base}__\\\"\\n else:\\n i = (\\n \\\"_\\\"\\n if self.inplace\\n else \\\"_functional\\\"\\n if self.functional_overload\\n else \\\"\\\"\\n )\\n return f\\\"{self.base}{i}\\\"\\n\\n\\n# Operator name is the base operator name along with the (typically not\\n# user visible) overload string.\\n@dataclass(frozen=True)\\nclass OperatorName:\\n name: BaseOperatorName\\n overload_name: str\\n\\n @staticmethod\\n def parse(op_name: str) -> OperatorName:\\n if \\\".\\\" in op_name:\\n name, overload_name = op_name.split(\\\".\\\", 1)\\n else:\\n name = op_name\\n overload_name = \\\"\\\"\\n r = OperatorName(name=BaseOperatorName.parse(name), overload_name=overload_name)\\n assert str(r) == op_name, f\\\"{str(r)} != {op_name}\\\"\\n return r\\n\\n def __str__(self) -> str:\\n if self.overload_name:\\n return f\\\"{self.name}.{self.overload_name}\\\"\\n else:\\n return f\\\"{self.name}\\\"\\n\\n # NB: This must be synchronized with the naming scheme in\\n # aten/src/ATen/templates/Operators.h\\n # Given a function schema \\\"aten::op.overload(...)\\\",\\n # If there is no overload name, this returns f\\\"{op}\\\"\\n # If there is an overload name, this returns f\\\"{op}_{overload}\\\"\\n def unambiguous_name(self) -> str:\\n if self.overload_name:\\n return f\\\"{self.name}_{self.overload_name}\\\"\\n else:\\n return f\\\"{self.name}\\\"\\n\\n def remove_inplace(self) -> OperatorName:\\n return OperatorName(\\n name=BaseOperatorName(\\n base=self.name.base,\\n inplace=False,\\n dunder_method=self.name.dunder_method,\\n ),\\n overload_name=self.overload_name,\\n )\\n\\n def with_overload(self, overload: str) -> OperatorName:\\n return OperatorName(\\n name=BaseOperatorName(\\n base=self.name.base,\\n inplace=False,\\n dunder_method=self.name.dunder_method,\\n ),\\n overload_name=overload,\\n )\\n\\n\\ndef gets_generated_out_inplace_wrapper(\\n f: NativeFunction, g: NativeFunctionsGroup, b: BackendIndex\\n) -> bool:\\n return (\\n f.func.kind() is not SchemaKind.functional\\n and not b.has_kernel(f)\\n and b.has_kernel(g.functional)\\n )\\n\\n\\n# NativeFunction objects that are views (f.is_view_op returns True)\\n# are added into a `NativeFunctionsViewGroup`, which we can use to\\n# easily access the generated (optional) view_copy NativeFunction.\\n# It's convenient to group them together, so we pair them up in NativeFunctionsViewGroup.\\n# See Note [Codegen'd {view}_copy Operators]\\n#\\n# One property of this representation is that in order for a view-like op to be part of\\n# a NativeFunctionsViewGroup, the \\\"aliasing\\\" version of that view op must exist.\\n# There's one case where that doesn't happen: we have a non-aliasing `narrow_copy.out` op,\\n# but don't have corresponding aliasing `narrow.out` op.\\n# This means that `narrow_copy.out` won't appear as a NativeFunctionsViewGroup.\\n@dataclass(frozen=True)\\nclass NativeFunctionsViewGroup:\\n view: NativeFunction\\n # Note: the {view}_copy operator is optional because we currently don't generate copy variants\\n # for all view ops. Notably, we don't generate them for CompositeImplicitAutograd views\\n # (we already get them \\\"for free\\\" through decomposition)\\n view_copy: NativeFunction | None\\n # view_inplace ops are also optional, but every view_inplace op should have out-of-place variant.\\n view_inplace: NativeFunction | None\\n\\n def __post_init__(self) -> None:\\n assert self.view.is_view_op\\n if self.view_copy is None:\\n assert not gets_generated_view_copy(self.view), (\\n f\\\"{str(self.view.func.name)} appears to be a new operator that aliases its inputs.\\\"\\n \\\" The codegen expects you to add a corresponding operator to native_functions.yaml:\\\"\\n f\\\" {get_view_copy_name(self.view)!s}.\\\"\\n \\\" See Note [view_copy NativeFunctions] for details.\\\"\\n )\\n else:\\n assert self.view_copy.func.name.name.base.endswith((\\\"_copy\\\", \\\"_scatter\\\"))\\n assert self.view.func.signature() == self.view_copy.func.signature(\\n strip_view_copy_name=True,\\n )\\n assert \\\"view_copy\\\" in self.view_copy.tags, (\\n f\\\"{str(self.view_copy.func.name), str(self.view.tags)} appears to be a view_copy operator. The codegen expects\\\"\\n \\\" view_copy operators to be annotated with the 'view_copy' tag in native_functions.yaml.\\\"\\n \\\" See Note [view_copy NativeFunction] for details.\\\"\\n )\\n if self.view_inplace is not None:\\n assert self.view.func.signature() == self.view_inplace.func.signature()\\n\\n if self.view.has_composite_implicit_autograd_kernel:\\n if self.view_inplace is not None:\\n assert self.view_inplace.has_composite_implicit_autograd_kernel, (\\n f\\\"{str(self.view.func.name)} and {str(self.view_inplace.func.name)} must either\\\"\\n \\\" both have CompositeImplicitAutograd kernels, or both not have composite kernels.\\\"\\n )\\n if self.view.has_composite_implicit_autograd_nested_tensor_kernel:\\n if self.view_inplace is not None:\\n assert (\\n self.view_inplace.has_composite_implicit_autograd_nested_tensor_kernel\\n ), (\\n f\\\"{str(self.view.func.name)} and {str(self.view_inplace.func.name)} must either\\\"\\n \\\" both have CompositeImplicitAutogradNestedTensor kernels, or both not have composite kernels.\\\"\\n )\\n\\n def functions(self, *, include_copy: bool = True) -> Iterator[NativeFunction]:\\n yield self.view\\n if self.view_inplace is not None:\\n yield self.view_inplace\\n if self.view_copy is not None and include_copy:\\n yield self.view_copy\\n\\n @property\\n def root_name(self) -> str:\\n return self.view.root_name\\n\\n @property\\n def composite(self) -> bool:\\n # We currently assert that the \\\"group\\\" is consistent.\\n # If the view op is composite, then its view_inplace op is too.\\n return self.view.has_composite_implicit_autograd_kernel\\n\\n\\ndef gets_generated_view_copy(f: NativeFunction) -> bool:\\n # Only aliasing (view) operators get a copy variant.\\n if not f.is_view_op:\\n return False\\n # We don't need to bother generating copy variants for CompositeImplicitAutograd ops,\\n # because we can let them decompose into base view ops.\\n if f.has_composite_implicit_autograd_kernel:\\n return False\\n # We also don't need to generate copy variants for inplace views.\\n if \\\"inplace_view\\\" in f.tags:\\n return False\\n # Assume ops ending in _inverse have manually-defined copy variants\\n # (e.g. slice_inverse() has the copy variant slice_scatter()).\\n # We -could- probably generate these as well, but the codegen will be\\n # slightly different, and hand-writing these few kernels keeps codegen\\n # complexity lower.\\n if f.func.name.name.base.endswith(\\\"_inverse\\\"):\\n return False\\n return True\\n\\n\\n# Given a NativeFunction that corresponds to a view op,\\n# returns the OperatorName of the corresponding \\\"copy\\\" variant of the op.\\ndef get_view_copy_name(f: NativeFunction) -> OperatorName:\\n # Right now, when asking for a view op's corresponding \\\"view_copy\\\" name\\n # we assert for sanity that the op is allowed to have a generated view_copy variant.\\n # (We can do this because \\\"gets_generated_view_copy()\\\" tell us which ops get a generated view_copy op).\\n # However, narrow_copy() already exists as an op directly in native_functions.yaml.\\n # I'm hardcoding narrow_copy here for now to maintain the assert,\\n # But we could also just get rid of the assert.\\n list_of_ops_with_explicit_view_copy_operators = [\\\"narrow\\\"]\\n if str(f.func.name) not in list_of_ops_with_explicit_view_copy_operators:\\n assert gets_generated_view_copy(f)\\n\\n base_name = f\\\"{f.func.name.name.base}_copy\\\"\\n view_copy_name = OperatorName(\\n name=BaseOperatorName(\\n base=base_name, inplace=False, dunder_method=f.func.name.name.dunder_method\\n ),\\n overload_name=f.func.name.overload_name,\\n )\\n return view_copy_name\\n\\n\\n# Helper functions for parsing argument lists (both inputs and returns)\\n\\n\\ndef parse_returns(return_decl: str) -> tuple[Return, ...]:\\n \\\"\\\"\\\"\\n Input: '()'\\n Output: []\\n \\\"\\\"\\\"\\n if return_decl == \\\"()\\\":\\n return ()\\n if return_decl[0] == \\\"(\\\" and return_decl[-1] == \\\")\\\":\\n return_decl = return_decl[1:-1]\\n return tuple(Return.parse(arg) for arg in return_decl.split(\\\", \\\"))\\n\\n\\n# A Precompute instance consists of a map from kernel argument name\\n# to the list of Argument instances that should replace that\\n# kernel argument in the impl function.\\n@dataclass(frozen=True)\\nclass Precompute:\\n # A map from kernel argument name -> a list of precomputed\\n # elements that replaces/supersedes it.\\n replace: dict[str, list[Argument]]\\n # List of precomputed args added without replacement\\n add: list[Argument]\\n\\n @staticmethod\\n def parse(src: object) -> Precompute:\\n assert isinstance(src, list)\\n\\n # src is a list of strings of the format:\\n # {kernel param name} -> {replacement decl}[, {replacement decl}, ...]\\n # [{add decl}[, {add decl}, ...]]\\n # The last line is optional and contains the precomputed parameters that are\\n # added without replacement.\\n # The other lines are parsed to get the names of which precomputed elements\\n # should replace which kernel arguments.\\n add_args = []\\n if \\\" -> \\\" not in src[-1]:\\n add_list = src[-1].split(\\\",\\\")\\n add_args = [Argument.parse(name.strip()) for name in add_list]\\n src = src[:-1]\\n\\n replace = {}\\n for raw_replace_item in src:\\n assert isinstance(raw_replace_item, str)\\n assert \\\" -> \\\" in raw_replace_item, (\\n \\\"precomputed parameters without replacement\\\"\\n \\\" are allowed only in the last line\\\"\\n )\\n\\n arg, with_list_raw = raw_replace_item.split(\\\" -> \\\")\\n assert (\\n \\\" \\\" not in arg\\n ), f\\\"illegal kernel param name '{arg}' in precomputed parameters'\\\"\\n with_list = with_list_raw.split(\\\",\\\")\\n with_list_args = [Argument.parse(name.strip()) for name in with_list]\\n replace[arg] = with_list_args\\n\\n r = Precompute(replace=replace, add=add_args)\\n assert r.to_list() == src, \\\"r.to_list() != src\\\"\\n return r\\n\\n def __post_init__(self) -> None:\\n # the template parameters are upper so if these are the\\n # same then it is ambiguous\\n for a in self.add:\\n assert a.name.upper() != a.name\\n for args in self.replace.values():\\n for a in args:\\n assert a.name.upper() != a.name\\n\\n def to_list(self) -> list[str]:\\n replace_list = []\\n for kernel_param, replacement_params in self.replace.items():\\n replacements = \\\", \\\".join(str(param) for param in replacement_params)\\n replace_list.append(f\\\"{kernel_param} -> {replacements}\\\")\\n\\n return replace_list\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\nfrom typing import Callable, TYPE_CHECKING\\n\\nfrom torchgen.api import cpp, dispatcher\\nfrom torchgen.api.translate import translate\\nfrom torchgen.api.types import (\\n BaseCType,\\n Binding,\\n CType,\\n DispatcherSignature,\\n FunctionalizationLambda,\\n iTensorListRefT,\\n NativeSignature,\\n OptionalCType,\\n optionalSymIntArrayRefT,\\n symIntArrayRefT,\\n SymIntT,\\n tensorListT,\\n tensorT,\\n VectorCType,\\n ViewInverseSignature,\\n)\\nfrom torchgen.context import (\\n method_with_native_function,\\n native_function_manager,\\n with_native_function,\\n with_native_function_and,\\n)\\nfrom torchgen.model import (\\n Argument,\\n BackendIndex,\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n ListType,\\n NativeFunction,\\n NativeFunctionsGroup,\\n NativeFunctionsViewGroup,\\n Return,\\n SchemaKind,\\n SelfArgument,\\n TensorOptionsArguments,\\n)\\nfrom torchgen.native_function_generation import (\\n INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY,\\n MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT,\\n OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY,\\n)\\nfrom torchgen.utils import dataclass_repr\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.selective_build.selector import SelectiveBuilder\\n\\n\\n# Note: [Mutable Ops Not Using Functionalization]\\n# Ops in this list currently do not work with functionalization and should be fixed.\\nMUTABLE_OPS_NOT_USING_FUNCTIONALIZATION = (\\n OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY\\n + MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT\\n + INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY\\n + [\\n # It will be BC-breaking, but we should fix their schemas.\\n # should be inplace?\\n \\\"record_stream\\\",\\n # See Note [resize_ in Functionalization]\\n \\\"resize_\\\",\\n \\\"resize_as_\\\",\\n # This function is used as for testing purposes only.\\n \\\"_fill_mem_eff_dropout_mask_\\\",\\n ]\\n)\\n\\n# This file contains codegen that relates to the functionalization pass.\\n# It includes:\\n# - gen_functionalization_definition\\n# Generates dispatcher kernel definitions for the functionalization pass.\\n# - gen_functionalization_registration\\n# Generates dispatcher kernel registrations for the functionalization pass.\\n# - gen_functionalization_view_inverse_declaration\\n# Generates a declaration for an \\\"inverse view\\\", for every view op\\n# that is needed in functionalization. We manually implement their definitions.\\n# - gen_composite_view_copy_kernel\\n# Generates view_copy() composite kernels for all view_copy operators.\\n\\n\\n# Generates the body of the default composite C++ kernel for a {view}_copy NativeFunction\\n# See Note [view_copy NativeFunctions]\\n@dataclass(frozen=True)\\nclass GenCompositeViewCopyKernel:\\n backend_index: BackendIndex\\n\\n @method_with_native_function\\n def __call__(self, g: NativeFunctionsViewGroup) -> str | None:\\n if g.view_copy is None:\\n return None\\n elif g.view_copy.func.name.name.base != f\\\"{g.view.func.name.name}_copy\\\":\\n # If the view_copy doesn't match the standard naming scheme of _copy,\\n # assume it already exists and doesn't need to be generated.\\n # Example: slice_inverse() with the copy variant named slice_scatter()\\n # instead of slice_inverse_copy()\\n return None\\n\\n metadata = self.backend_index.get_kernel(g.view_copy)\\n assert metadata is not None\\n\\n # We can make view_copy work in more cases by using reshape()\\n # when a normal view call would ordinarily fail.\\n # This also makes LTC more efficient, because they don't need to include\\n # clone() calls in their graph (which is normally needed by reshape).\\n if str(g.view_copy.func.name) == \\\"view_copy\\\":\\n assert metadata.kernel == \\\"view_copy_symint\\\"\\n return \\\"\\\"\\\"\\\\\\nat::Tensor view_copy_symint(const at::Tensor & self, at::SymIntArrayRef size) {\\n c10::SymDimVector shape = infer_size_dv(size, self.sym_numel());\\n if (!at::detail::computeStride(self.sym_sizes(), self.sym_strides(), shape).has_value()) {\\n return self.reshape_symint(size);\\n } else {\\n auto output = at::_ops::view::call(self, size);\\n return output.clone(/*memory_format=*/at::MemoryFormat::Contiguous);\\n }\\n}\\n\\\"\\\"\\\"\\n # view_copy is a native signature, since we're generating an at::native:: kernel\\n # Functionalization always operates on symints though\\n view_copy_sig = NativeSignature(\\n g.view_copy.func, symint=metadata.supports_symint()\\n )\\n\\n # view is a dispatcher signature, since we're calling into the at::_ops API\\n view_sig = DispatcherSignature(g.view.func)\\n\\n view_api_name = g.view.func.name.unambiguous_name()\\n exprs = \\\", \\\".join(\\n [e.expr for e in translate(view_copy_sig.arguments(), view_sig.arguments())]\\n )\\n\\n # view ops today always return either a Tensor or a list of Tensors\\n assert len(g.view.func.returns) == 1\\n assert g.view.func.returns[0].type == BaseType(\\n BaseTy.Tensor\\n ) or g.view.func.returns[0].type == ListType(BaseType(BaseTy.Tensor), None)\\n\\n if g.view.func.returns[0].type == BaseType(BaseTy.Tensor):\\n return_cloned_output = \\\"\\\"\\\"\\\\\\n return output.clone(/*memory_format=*/at::MemoryFormat::Contiguous);\\\"\\\"\\\"\\n else:\\n # If the return type is a list, we need to clone each tensor in the list.\\n return_cloned_output = f\\\"\\\"\\\"\\\\\\n {view_copy_sig.returns_type().cpp_type()} out_clone;\\n for (const auto i : c10::irange(output.size())) {{\\n out_clone.push_back(output[i].clone(/*memory_format=*/at::MemoryFormat::Contiguous));\\n }}\\n return out_clone;\\\"\\\"\\\"\\n\\n # The default generated composite kernel for {view}_copy() operators just clones\\n # the input tensor, and runs the underlying view on the clone.\\n return f\\\"\\\"\\\"\\n{view_copy_sig.defn(name=metadata.kernel)} {{\\n auto output = at::_ops::{view_api_name}::call({exprs});\\n {return_cloned_output}\\n}}\\n\\\"\\\"\\\"\\n\\n\\ndef return_str(rets: tuple[Return, ...], names: list[str]) -> str:\\n assert len(rets) == len(names)\\n if len(rets) == 0:\\n return \\\"\\\"\\n elif len(rets) == 1:\\n return f\\\"return {names[0]};\\\"\\n else:\\n return f\\\"return {dispatcher.returns_type(rets).cpp_type()}({', '.join(names)});\\\"\\n\\n\\ndef modifies_arguments(f: NativeFunction) -> bool:\\n return any(\\n a.annotation is not None and a.annotation.is_write\\n for a in f.func.arguments.flat_all\\n )\\n\\n\\ndef wrapper_name(func: FunctionSchema) -> str:\\n if func.name.overload_name:\\n return f\\\"{cpp.name(func)}_{func.name.overload_name}\\\"\\n else:\\n return cpp.name(func)\\n\\n\\ndef is_tensor_like(a: Argument | TensorOptionsArguments | SelfArgument) -> bool:\\n return isinstance(a, SelfArgument) or (\\n isinstance(a, Argument) and a.type.is_tensor_like()\\n )\\n\\n\\n# We need to wrap / unwrap various arguments from the op in the functionalization kernels.\\n# Some op schemas include non-owning types though (like TensorList),\\n# and when we unwrap them we expect to get out an owning type!.\\n# We also return a lambda that tells you how to conver the non-owning type argument into the owning type.\\ndef get_owning_type(t: CType) -> tuple[CType, Callable[[str], str]]:\\n if t == BaseCType(tensorListT):\\n return VectorCType(BaseCType(tensorT)), lambda x: f\\\"{x}.vec()\\\"\\n if t == BaseCType(iTensorListRefT):\\n return VectorCType(BaseCType(tensorT)), lambda x: f\\\"{{{x}.begin(), {x}.end()}}\\\"\\n # There are technically other non-owning types out there (like IntArrayRef),\\n # but functionalization only actually cares about the ones involving tensors.\\n return t, lambda x: x\\n\\n\\n# unwraps all tensor-like arguments, returning:\\n# (1) a string containing all of the logic that does the unwrapping\\n# (2) a context, to be used by translate(), with all of the relevant bindings.\\ndef unwrap_tensor_args(\\n sig: DispatcherSignature, *, is_view_op: bool\\n) -> tuple[str, list[Binding]]:\\n context: list[Binding] = []\\n unwrapped_tensor_args: list[str] = []\\n for arg in sig.arguments():\\n if is_tensor_like(arg.argument):\\n # for tensor inputs, we want to unwrap them before passing them into the redispatch calls.\\n unwrapped_name = f\\\"{arg.name}_\\\"\\n # For most ops, the functionalization needs to sync any pending updates on the input tensors\\n # before calling the operator, since otherwise the operator will act on stale data.\\n # For view ops though, we can continue to defer syncing until the tensor is used by\\n # a non-view operator.\\n maybe_sync_input = (\\n \\\"\\\" if is_view_op else f\\\"at::functionalization::impl::sync({arg.name});\\\"\\n )\\n unwrapped_type, conversion_fn = get_owning_type(\\n arg.nctype.remove_const_ref().type\\n )\\n unwrapped_tensor_args.append(\\n f\\\"\\\"\\\"\\n {unwrapped_type.cpp_type()} {unwrapped_name};\\n if (at::functionalization::impl::isFunctionalTensor({arg.name})) {{\\n {maybe_sync_input}\\n {unwrapped_name} = at::functionalization::impl::from_functional_tensor({arg.name});\\n }} else {{\\n {unwrapped_name} = {conversion_fn(arg.name)};\\n }}\\\"\\\"\\\"\\n )\\n context.append(arg.with_name(unwrapped_name))\\n else:\\n # for non-tensor inputs, we want to pass them directly into the redispatch calls.\\n context.append(arg)\\n unwrap_tensor_args_str = \\\"\\\\n \\\".join(unwrapped_tensor_args)\\n return unwrap_tensor_args_str, context\\n\\n\\n# converts all tensor-like arguments to meta tensors, which are used to compute stride info. Returns:\\n# (1) a string containing all of the logic that does the conversions.\\n# (2) a context, to be used by translate(), with all of the relevant bindings.\\ndef convert_to_meta_tensors(sig: DispatcherSignature) -> tuple[str, list[Binding]]:\\n context: list[Binding] = []\\n unwrapped_tensor_args: list[str] = []\\n for arg in sig.arguments():\\n if is_tensor_like(arg.argument):\\n # for tensor inputs, we want to unwrap them before passing them into the redispatch calls.\\n a_ = arg.name\\n unwrapped_name = f\\\"{arg.name}_meta\\\"\\n unwrapped_tensor_args.append(f\\\"auto {unwrapped_name} = to_meta({a_});\\\")\\n context.append(arg.with_name(unwrapped_name))\\n else:\\n # for non-tensor inputs, we want to pass them directly into the redispatch calls.\\n context.append(arg)\\n unwrap_tensor_args_str = \\\"\\\\n \\\".join(unwrapped_tensor_args)\\n return unwrap_tensor_args_str, context\\n\\n\\n# The functionalization codegen currently expects view op schemas to have this form:\\n# foo(Tensor(a), ...) -> Tensor(a) (e.g. transpose)\\n# foo(Tensor(a!), ...) -> Tensor(a!) (e.g. transpose_)\\ndef assert_view_op_properties(func: FunctionSchema) -> None:\\n def is_alias(a: Argument) -> bool:\\n return a.annotation is not None\\n\\n args = func.arguments.flat_non_out\\n # The first argument is a tensor with an alias semantics (annotations)\\n assert len(args) > 0 and args[0].type == BaseType(\\n BaseTy.Tensor\\n ), f\\\"\\\"\\\"In the functionalization codegen, we expect the first argument of every view operator to be a tensor,\\nbut found an argument of type {str(args[0].type)} for operator: {str(func.name)}.\\\"\\\"\\\"\\n # No other arguments have aliasing semantics\\n assert is_alias(args[0]) and not any(\\n is_alias(a) for a in args[1:]\\n ), \\\"\\\"\\\"In the functionalization codegen, we expect the first argument of every view operator to alias the output.\\nView operators with multiple aliasing inputs aren't supported yet. Found an operator that doesn't satisfy this constraint\\\"\\\"\\\"\\n\\n\\n# One-liner expression for checking if an expression expr of type type has any\\n# symbolic values.\\ndef emit_expr_has_symbolic_values(expr: str, type: CType) -> str:\\n if type == BaseCType(SymIntT):\\n return f\\\"{expr}.is_symbolic()\\\"\\n\\n if isinstance(type, OptionalCType):\\n innerexpr = f\\\"(*{expr})\\\"\\n return f\\\"{expr}.has_value() ? {emit_expr_has_symbolic_values(innerexpr, type.elem)} : false\\\"\\n\\n if type == BaseCType(optionalSymIntArrayRefT):\\n return emit_expr_has_symbolic_values(\\n expr, OptionalCType(BaseCType(symIntArrayRefT))\\n )\\n\\n if type in (BaseCType(symIntArrayRefT), VectorCType(BaseCType(SymIntT))):\\n argname = \\\"arg\\\"\\n lambda_check = emit_expr_has_symbolic_values(argname, BaseCType(SymIntT))\\n return (\\n \\\"std::any_of(\\\"\\n f\\\"{expr}.begin(), {expr}.end(), \\\"\\n f\\\"[=](auto& {argname}) {{ return {lambda_check}; }})\\\"\\n )\\n\\n raise ValueError(\\n \\\"unsupported type for has_symbolic_values check. \\\"\\n \\\"It should be a SymInt or a collection of those. \\\"\\n f\\\"Got: {type.cpp_type()}\\\"\\n )\\n\\n\\n# Detects whether any of the SymInt arguments are, in fact, symbolic values.\\n# This is used in the constructor of ViewMeta.\\ndef emit_has_symbolic_inputs(sig: DispatcherSignature) -> tuple[str, str]:\\n name = \\\"has_symbolic_inputs\\\"\\n statements = [\\n f\\\"{name} = {name} | ({emit_expr_has_symbolic_values(binding.name, binding.nctype.type)});\\\"\\n for binding in sig.arguments()\\n if (\\n isinstance(binding.argument, Argument)\\n and binding.argument.type.is_symint_like()\\n )\\n ]\\n body = \\\"\\\\n \\\".join(statements)\\n return (\\n name,\\n f\\\"\\\"\\\"\\n bool {name} = false;\\n {body}\\\"\\\"\\\",\\n )\\n\\n\\n# Generates the Functionalization kernel for:\\n# - ops that create aliases (e.g. transpose())\\n# - ops that are views AND mutations (e.g. transpose_())\\ndef emit_view_functionalization_body(\\n g: NativeFunctionsViewGroup, *, view_inplace: bool\\n) -> str:\\n if view_inplace:\\n # This op is both an inplace op AND a view op.\\n # See Note [Functionalization Pass - Inplace View Ops] for details.\\n # I currently have the view meta call into the out-of-place variant of the view, to avoid\\n # having to define an extra ~20 inplace {view}_inverse_ functions.\\n # Most view ops don't have NativeFunctionGroup's both, because we don't define out= variants for view ops.\\n # I'm assuming that every inplace-view op has a corresponding out-of-place view op,\\n # with the same name but the trailing underscore removed.\\n # This is currently asserted at parse time in gen.py (see error_check_native_functions).\\n assert g.view_inplace is not None\\n f = g.view_inplace\\n else:\\n f = g.view\\n\\n assert g.view_copy is not None\\n with native_function_manager(f):\\n call_sig = DispatcherSignature.from_schema(g.view_copy.func)\\n\\n # the \\\"view_copy\\\" op name that the functionalization kernels need to call\\n api_name = g.view_copy.func.name.unambiguous_name()\\n # Sometimes the functionalization pass needs to no-op (e.g. if it was passed non-functional tensors)\\n # \\\"no-op\\\"ing in this context is just redispatching to the original op.\\n noop_api_name = f.func.name.unambiguous_name()\\n\\n dispatcher_sig = DispatcherSignature.from_schema(f.func)\\n assert_view_op_properties(f.func)\\n view_tensor_name = dispatcher_sig.arguments()[0].name\\n\\n return_type = dispatcher_sig.returns_type().remove_const_ref().cpp_type()\\n\\n unwrap_tensor_args_str, unwrapped_args_ctx = unwrap_tensor_args(\\n dispatcher_sig, is_view_op=True\\n )\\n view_redispatch_args = [\\n e.expr\\n for e in translate(unwrapped_args_ctx, call_sig.arguments(), method=False)\\n ]\\n\\n forward_lambda = FunctionalizationLambda.from_func(g, is_reverse=False)\\n reverse_lambda = FunctionalizationLambda.from_func(g, is_reverse=True)\\n\\n # The meta API call should use the same arguments, but convert all tensors to meta tensors first.\\n meta_conversion_str, meta_call_ctx = convert_to_meta_tensors(dispatcher_sig)\\n meta_call_args = [\\n e.expr for e in translate(meta_call_ctx, call_sig.arguments(), method=False)\\n ]\\n\\n (\\n symbolic_inputs_varname,\\n symbolic_inputs_check,\\n ) = emit_has_symbolic_inputs(call_sig)\\n\\n if \\\"inplace_view\\\" in f.tags:\\n # See Note [Functionalization Pass - Inplace View Ops] for more details\\n return f\\\"\\\"\\\"\\n {dispatcher_sig.defn(name=wrapper_name(f.func), is_redispatching_fn=True)} {{\\n if (!at::functionalization::impl::isFunctionalTensor({view_tensor_name})) {{\\n // functionalization is re-entrant, but will no-op if it wasn't passed a FunctionalTensorWrapper.\\n {unwrap_tensor_args_str}\\n at::AutoDispatchSkipFunctionalize guard;\\n return at::_ops::{noop_api_name}::call({', '.join(view_redispatch_args)});\\n }}\\n auto reapply_views = at::functionalization::impl::getFunctionalizationReapplyViewsTLS();\\n auto inverse_return_mode = (\\n reapply_views ? at::functionalization::InverseReturnMode::ViewOrScatterInverse\\n : at::functionalization::InverseReturnMode::NeverView\\n );\\n {symbolic_inputs_check}\\n at::functionalization::ViewMeta view_meta = at::functionalization::ViewMeta(\\n {forward_lambda.decl()} {{\\n if (reapply_views) {{\\n return {forward_lambda.inner_call(reapply_views=True)}\\n }} else {{\\n return {forward_lambda.inner_call(reapply_views=False)}\\n }}\\n }},\\n {reverse_lambda.decl()} {{\\n return {reverse_lambda.inner_call()}\\n }},\\n /*has_symbolic_inputs=*/{symbolic_inputs_varname}\\n );\\n auto compute_reference_meta =\\n {view_tensor_name}.key_set().has_backend(c10::BackendComponent::XLABit) ||\\n {view_tensor_name}.key_set().has_backend(c10::BackendComponent::LazyBit);\\n {return_type} reference_tensor_output;\\n if (compute_reference_meta) {{\\n {meta_conversion_str}\\n at::AutoDispatchSkipFunctionalize func_guard;\\n c10::impl::ExcludeDispatchKeyGuard guard(exclude_keys_for_meta_dispatch);\\n reference_tensor_output = at::_ops::{noop_api_name}::call({', '.join(meta_call_args)});\\n }}\\n // This function adds the above view meta to the current tensor and replays them off the base,\\n // mutating the size/stride info of the current FunctionalTensorWrapper.\\n // Because of this, we need to make sure to run the reference shape function above,\\n // BEFORE doing this (otherwise we'll end up runnin the reference function using the wrong sizes/strides)\\n at::functionalization::impl::mutate_view_meta({view_tensor_name}, view_meta);\\n // See Note [Propagating strides in the functionalization pass]\\n // XLA/LTC don't implement the logic to propagate strides correctly, so we need to rely\\n // on a reference implementation here (instead of relying on the output from the forward lambda\\n // having the correct stride info)\\n if (compute_reference_meta) {{\\n at::functionalization::impl::set_sizes_strides_offset({view_tensor_name}, reference_tensor_output);\\n }}\\n return {view_tensor_name};\\n }}\\n\\\"\\\"\\\"\\n\\n else:\\n is_multi_output_view = isinstance(f.func.returns[0].type, ListType)\\n return f\\\"\\\"\\\"\\n {dispatcher_sig.defn(name=wrapper_name(f.func), is_redispatching_fn=True)} {{\\n {unwrap_tensor_args_str}\\n if (!at::functionalization::impl::isFunctionalTensor({view_tensor_name})) {{\\n // functionalization is re-entrant, but will no-op if it wasn't passed a FunctionalTensorWrapper.\\n at::AutoDispatchSkipFunctionalize guard;\\n return at::_ops::{noop_api_name}::call({', '.join(view_redispatch_args)});\\n }}\\n auto reapply_views = at::functionalization::impl::getFunctionalizationReapplyViewsTLS();\\n auto inverse_return_mode = (\\n reapply_views ? at::functionalization::InverseReturnMode::ViewOrScatterInverse\\n : at::functionalization::InverseReturnMode::NeverView\\n );\\n auto compute_reference_meta =\\n {view_tensor_name}.key_set().has_backend(c10::BackendComponent::XLABit) ||\\n {view_tensor_name}.key_set().has_backend(c10::BackendComponent::LazyBit);\\n {return_type} reference_tensor_output;\\n if (compute_reference_meta) {{\\n {meta_conversion_str}\\n at::AutoDispatchSkipFunctionalize func_guard;\\n c10::impl::ExcludeDispatchKeyGuard guard(exclude_keys_for_meta_dispatch);\\n reference_tensor_output = at::_ops::{noop_api_name}::call({', '.join(meta_call_args)});\\n }}\\n {return_type} tmp_output;\\n {{\\n at::AutoDispatchSkipFunctionalize guard;\\n if (reapply_views) {{\\n tmp_output = at::_ops::{noop_api_name}::call({', '.join(view_redispatch_args)});\\n }} else {{\\n tmp_output = at::_ops::{api_name}::call({', '.join(view_redispatch_args)});\\n }}\\n }}\\n {symbolic_inputs_check}\\n at::functionalization::ViewMeta view_meta = at::functionalization::ViewMeta(\\n {forward_lambda.decl()} {{\\n if (reapply_views) {{\\n return {forward_lambda.inner_call(reapply_views=True)}\\n }} else {{\\n return {forward_lambda.inner_call(reapply_views=False)}\\n }}\\n }},\\n {reverse_lambda.decl()} {{\\n return {reverse_lambda.inner_call()}\\n }},\\n /*has_symbolic_inputs=*/{symbolic_inputs_varname},\\n /*is_multi_output=*/{str(is_multi_output_view).lower()},\\n /*is_as_strided=*/{str(str(f.func.name) == 'as_strided').lower()}\\n );\\n auto out = at::functionalization::impl::create_functional_tensor_with_view_meta(tmp_output, {view_tensor_name}, view_meta);\\n // See Note [Propagating strides in the functionalization pass]\\n if (compute_reference_meta) {{\\n at::functionalization::impl::set_sizes_strides_offset(out, reference_tensor_output);\\n }}\\n return out;\\n }}\\n\\\"\\\"\\\"\\n\\n\\ndef maybe_create_output(f: NativeFunction, var_name: str) -> str:\\n if len(f.func.returns) == 0:\\n return \\\"\\\"\\n return_type = dispatcher.returns_type(f.func.returns).remove_const_ref().cpp_type()\\n return f\\\"{return_type} {var_name} = \\\"\\n\\n\\n# Given a NativeFunction, and a variable name corresponding to the output of redispatching on the function,\\n# this returns two lists of names, consisting of:\\n# - the names of returns corresponding to the original (mutable) inputs of the outer function\\n# - the names of returns corresponding to the (immutable) outputs of the inner redispatched function\\ndef get_mutable_redispatch_return_names(\\n f: NativeFunction, inner_return_var: str\\n) -> tuple[list[str], list[str]]:\\n aliased_returns = []\\n non_aliased_returns = []\\n for i, name in enumerate(f.func.aliased_return_names()):\\n if name is not None:\\n aliased_returns.append(name)\\n else:\\n non_aliased_returns.append(\\n inner_return_var\\n if len(f.func.returns) == 1\\n else f\\\"std::get<{i}>({inner_return_var})\\\"\\n )\\n return aliased_returns, non_aliased_returns\\n\\n\\n# When functionalization \\\"no-op's\\\" and redispatches on a mutable operator, we need to take care so that:\\n# - For fresh outputs, we return the result of the redispatch (without wrapping outputs)\\n# - For outputs that were aliased to inputs, we return the inputs directly (since some of them might have been wrapped)\\ndef return_from_mutable_noop_redispatch(\\n f: NativeFunction, inner_return_var: str\\n) -> str:\\n aliased, non_aliased = get_mutable_redispatch_return_names(f, inner_return_var)\\n # Just get all of the return names, and immediately return them\\n return return_str(f.func.returns, aliased + non_aliased)\\n\\n\\ndef wrap_propagate_mutations_and_return(\\n f: NativeFunction, functional_op: NativeFunction, inner_return_var: str\\n) -> str:\\n mutable_arg_names = f.func.arguments.mutable_arg_names()\\n (\\n aliased_outer_rets,\\n non_aliased_outer_rets,\\n ) = get_mutable_redispatch_return_names(f, inner_return_var)\\n _, non_aliased_inner_rets = get_mutable_redispatch_return_names(\\n functional_op, inner_return_var\\n )\\n # The outer function may have a mix of aliased and non-aliased outputs,\\n # But the inner functional op that we're transforming to should only have non-aliased outputs\\n assert len(mutable_arg_names) + len(non_aliased_outer_rets) == len(\\n non_aliased_inner_rets\\n )\\n\\n # First, take all of the newly created outputs from the inner call and wrap them into functional tensors\\n updates = []\\n non_aliased_wrapped_ret_names = []\\n for i, inner_ret in enumerate(\\n non_aliased_inner_rets[: len(non_aliased_outer_rets)]\\n ):\\n ret_name = f\\\"output_{i}\\\"\\n updates.append(\\n f\\\"\\\"\\\"\\\\\\n auto output_{i} = at::functionalization::impl::to_functional_tensor({inner_ret});\\\"\\\"\\\"\\n )\\n non_aliased_wrapped_ret_names.append(ret_name)\\n\\n # Next, take all of the mutated outputs from the inner call corresponding to mutated inputs,\\n # and propagate the mutations\\n for outer_arg, inner_ret in zip(\\n mutable_arg_names, non_aliased_inner_rets[len(non_aliased_outer_rets) :]\\n ):\\n updates.append(\\n f\\\"\\\"\\\"\\\\\\n auto {outer_arg}_inner = at::functionalization::impl::from_functional_tensor({outer_arg});\\n at::functionalization::impl::replace_({outer_arg}, {inner_ret});\\n at::functionalization::impl::commit_update({outer_arg});\\n at::functionalization::impl::sync({outer_arg});\\n auto {outer_arg}_inner_updated = at::functionalization::impl::from_functional_tensor({outer_arg});\\n at::functionalization::impl::propagate_xla_data_direct({outer_arg}_inner, {outer_arg}_inner_updated);\\\"\\\"\\\"\\n )\\n\\n # Finally, we return:\\n # - Any mutable arguments that also returns\\n # - Any immutable returns that were created wrapping the output from the inner call\\n returns_str = return_str(\\n f.func.returns, aliased_outer_rets + non_aliased_wrapped_ret_names\\n )\\n updates_str = \\\"\\\\n\\\".join(updates)\\n return f\\\"\\\"\\\"\\\\\\n{updates_str}\\n {returns_str}\\\"\\\"\\\"\\n\\n\\n# Generates the Functionalization kernel for:\\n# - mutation ops (inplace and out= ops)\\n@with_native_function_and\\ndef emit_inplace_functionalization_body(\\n f: NativeFunction, g: NativeFunctionsGroup\\n) -> str:\\n # mutation case\\n assert modifies_arguments(f)\\n\\n dispatcher_sig = DispatcherSignature.from_schema(f.func)\\n\\n unwrap_tensor_args_str, unwrapped_args_ctx = unwrap_tensor_args(\\n dispatcher_sig, is_view_op=False\\n )\\n\\n mutated_names = [\\n a.name\\n for a in f.func.arguments.flat_all\\n if a.type.is_tensor_like() and a.annotation is not None\\n ]\\n non_mutated_names = [\\n a.name\\n for a in f.func.arguments.flat_all\\n if a.type.is_tensor_like() and a.annotation is None\\n ]\\n non_mutated_tensor_names = [\\n a.name\\n for a in f.func.arguments.flat_all\\n if a.type == BaseType(BaseTy.Tensor) and a.annotation is None\\n ]\\n # all mutable inputs must be functional tensors in order to participate in functionalization\\n check_all_mutated_args_are_functional = \\\" && \\\".join(\\n [\\\"true\\\"]\\n + [\\n f\\\"at::functionalization::impl::isFunctionalTensor({a})\\\"\\n for a in mutated_names\\n ]\\n )\\n check_any_non_mutated_args_are_functional = \\\" || \\\".join(\\n [\\\"false\\\"]\\n + [\\n f\\\"at::functionalization::impl::isFunctionalTensor({a})\\\"\\n for a in non_mutated_names\\n ]\\n )\\n\\n check_any_non_mutated_tensors_are_xla = \\\" || \\\".join(\\n [\\\"false\\\"]\\n + [\\n f\\\"{a}.device().type() == c10::DeviceType::XLA\\\"\\n for a in non_mutated_tensor_names\\n ]\\n )\\n # These are used in the cases where we don't functionalize and redispatch to the inplace op\\n # case 1: we hit an inplace op that doesn't have an out-of-place equivalent\\n # case 2: we hit an inplace ops but our inputs are not functional tensors (in which case our kernel just no-ops)\\n inplace_exprs = [\\n e.expr\\n for e in translate(unwrapped_args_ctx, dispatcher_sig.arguments(), method=False)\\n ]\\n\\n # call the out-of-place variant of the op\\n return_type = (\\n dispatcher.returns_type(g.functional.func.returns).remove_const_ref().cpp_type()\\n )\\n functional_sig = DispatcherSignature.from_schema(g.functional.func)\\n functional_exprs = [\\n e.expr\\n for e in translate(unwrapped_args_ctx, functional_sig.arguments(), method=False)\\n ]\\n\\n if f.func.is_out_fn():\\n mutable_input_post_processing = \\\"\\\\n\\\".join(\\n [\\n f\\\"\\\"\\\"\\n at::functionalization::impl::replace_(\\n {a.name}, {'std::get<' + str(i) + '>(tmp_output)' if len(f.func.returns) > 1 else 'tmp_output'});\\n at::functionalization::impl::commit_update({a.name});\\\"\\\"\\\"\\n for (i, a) in enumerate(f.func.arguments.out)\\n if a.annotation and a.annotation.is_write and a.type.is_tensor_like()\\n ]\\n )\\n else:\\n mutable_input_post_processing = \\\"\\\\n\\\".join(\\n [\\n f\\\"\\\"\\\"\\n at::functionalization::impl::replace_({a.name}, tmp_output);\\n at::functionalization::impl::commit_update({a.name});\\\"\\\"\\\"\\n for a in f.func.arguments.flat_all\\n if a.annotation and a.annotation.is_write and a.type.is_tensor_like()\\n ]\\n )\\n\\n meta_conversion_str, meta_call_ctx = convert_to_meta_tensors(dispatcher_sig)\\n # We don't want to run the inplace meta func for ops like .set_(), because:\\n # (1) they're unnecessary: inplace meta checks are only useful for ops like add_(),\\n # where broadcasting will work for the out-of-place case but should fail on the inplace call\\n # (2) They'll also fail without adding extra infra: we'd need to convert the input storage argument\\n # into a meta storage\\n any_storage_args = any(\\n a.type == BaseType(BaseTy.Storage) for a in f.func.arguments.flat_all\\n )\\n\\n return f\\\"\\\"\\\"\\n {dispatcher_sig.defn(name=wrapper_name(f.func), is_redispatching_fn=True)} {{\\n if ({str(not any_storage_args and f.func.kind() == SchemaKind.inplace).lower()}) {{\\n // Before converting the mutable op to its functional variant, run meta tensors through the original op.\\n // This will help us catch shape errors that apply to inplace ops that wouldn't apply to their functional variants.\\n // (We can only do this for inplace ops today though, because they technically all support meta tensors).\\n {meta_conversion_str}\\n at::AutoDispatchSkipFunctionalize func_guard;\\n c10::impl::ExcludeDispatchKeyGuard guard(exclude_keys_for_meta_dispatch);\\n at::_ops::{f.func.name.unambiguous_name()}::call({', '.join(a.name for a in meta_call_ctx)});\\n }}\\n {unwrap_tensor_args_str}\\n if (!({check_all_mutated_args_are_functional})) {{\\n // We want to disable this check if there are any XLA tensors.\\n // cpu_tensor.copy_(xla_tensor) is valid code.\\n if (!({check_any_non_mutated_tensors_are_xla}) && ({check_any_non_mutated_args_are_functional})) {{\\n // case 1: trying to mutate a non functional tensor with a functional tensor is an error\\n TORCH_INTERNAL_ASSERT(false,\\n \\\"mutating a non-functional tensor with a functional tensor is not allowed.\\\",\\n \\\" Please ensure that all of your inputs are wrapped inside of a functionalize() call.\\\");\\n }} else {{\\n // case 2: arguments are not functional tensors, so we no-op and redispatch.\\n at::AutoDispatchSkipFunctionalize guard;\\n {maybe_create_output(f, 'tmp_output')}at::_ops::{f.func.name.unambiguous_name()}::call({', '.join(inplace_exprs)});\\n {return_from_mutable_noop_redispatch(f, 'tmp_output')}\\n }}\\n }} else {{\\n {return_type} tmp_output;\\n {{\\n at::AutoDispatchSkipFunctionalize guard;\\n tmp_output = at::_ops::{g.functional.func.name.unambiguous_name()}::call({', '.join(functional_exprs)});\\n }}\\n {wrap_propagate_mutations_and_return(f, g.functional, 'tmp_output')}\\n }}\\n }}\\\"\\\"\\\"\\n\\n\\n# The below functions generate RegisterFunctionalization.cpp\\n# These files provide the kernels that run the functionalization pass, which can be opted into\\n# per backend (e.g. XLA or Vulkan), or as a composable transform (functionalize() in functorch).\\n\\n\\n# See Note [Functionalization Pass: View Inverses].\\ndef gen_functionalization_view_inverse_declaration(\\n selector: SelectiveBuilder, g: NativeFunctionsViewGroup\\n) -> str | None:\\n # For every (non-composite) view op, we need a corresponding \\\"inverse view\\\" function.\\n # This generates the declarations so we get a good compiler error when someone adds a new view.\\n @with_native_function\\n def emit_decl_helper(g: NativeFunctionsViewGroup) -> str | None:\\n if g.view.has_composite_implicit_autograd_kernel:\\n return None\\n view_inverse_sig = ViewInverseSignature(g)\\n return view_inverse_sig.decl()\\n\\n return emit_decl_helper(g)\\n\\n\\ndef gen_functionalization_registration(\\n selector: SelectiveBuilder,\\n g: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup,\\n composite_implicit_autograd_index: BackendIndex,\\n) -> list[str]:\\n @with_native_function\\n def emit_registration_helper(f: NativeFunction) -> str:\\n assert not f.has_composite_implicit_autograd_kernel\\n registration_str = f\\\"TORCH_FN(functionalization::{wrapper_name(f.func)})\\\"\\n return f'm.impl(\\\"{f.func.name}\\\", {registration_str});'\\n\\n # Don't generate kernels in mobile build\\n if not selector.include_all_operators:\\n return []\\n\\n if isinstance(g, NativeFunctionsViewGroup):\\n # functionalization needs to register kernels for view + view_inplace ops\\n # See Note [Functionalization <> torch.Tensor constructor]\\n if str(g.view.func.name) == \\\"lift_fresh\\\":\\n return []\\n view_str = []\\n if not g.view.has_composite_implicit_autograd_kernel:\\n view_str.append(emit_registration_helper(g.view))\\n if (\\n g.view_inplace is not None\\n and not g.view_inplace.has_composite_implicit_autograd_kernel\\n ):\\n assert g.view_inplace.is_view_op\\n view_str.append(emit_registration_helper(g.view_inplace))\\n return view_str\\n\\n elif isinstance(g, NativeFunctionsGroup):\\n # Gets a hand-written functionalization kernel\\n if g.inplace is not None and str(g.inplace.func.name) == \\\"set_.source_Tensor\\\":\\n fns = []\\n else:\\n fns = list(g.functions())\\n else:\\n if str(g.func.name) in MUTABLE_OPS_NOT_USING_FUNCTIONALIZATION:\\n return []\\n fns = [g]\\n\\n registrations = []\\n for f in fns:\\n if f.has_composite_implicit_autograd_kernel:\\n continue\\n if str(f.func.name) == \\\"lift\\\":\\n # See Note [Functionalization <> torch.Tensor constructor]\\n return []\\n if str(f.func.name) == \\\"resize_\\\":\\n # See Note [resize_ in Functionalization]\\n return []\\n if str(f.func.name.name) != \\\"set_\\\":\\n assert not f.is_view_op\\n # functionalization needs to generate and register kernels for inplace ops.\\n # We *also* need to directly register CompositeImplicitAUtograd kernels\\n # so that they decompose properly before functioanlization.\\n if modifies_arguments(f):\\n registrations.append(emit_registration_helper(f))\\n return registrations\\n\\n\\ndef gen_functionalization_definition(\\n selector: SelectiveBuilder,\\n # Note: Ideally this code should never have to look at NativeFunction\\n # (and instead only need to operate on grouped NativeFunctions).\\n # The only reason currently is because we need to emit direct dispatch registrations\\n # For CompositeImplicitAutograd operators, which are potentially ungrouped.\\n g: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup,\\n) -> list[str]:\\n # Don't generate kernels in mobile build\\n if not selector.include_all_operators:\\n return []\\n\\n if isinstance(g, NativeFunctionsViewGroup):\\n # Case 1: emit view -> view_copy kernels for the functionalization pass\\n view_defs = []\\n if not g.composite:\\n # invariant: NativeFunctionsViewGroup's always have a view_copy operator\\n # if the view is not composite (implicit autograd)\\n assert g.view_copy is not None, dataclass_repr(g, indent=1)\\n view_defs.append(emit_view_functionalization_body(g, view_inplace=False))\\n if g.view_inplace is not None:\\n view_defs.append(emit_view_functionalization_body(g, view_inplace=True))\\n return view_defs\\n elif isinstance(g, NativeFunction):\\n # Invariant: all mutable operators that we need to handle in functionalization\\n # should have been properly grouped up.\\n # TODO: The below ops all have \\\"problematic\\\" schemas that prevent them from\\n # getting functionalized. Instead of bending over backwards to get things to work,\\n # I think we should either:\\n # (1) fix their schemas (BC-breaking)\\n # (2) hand-write their functionalization kernels\\n if (\\n str(g.func.name) not in MUTABLE_OPS_NOT_USING_FUNCTIONALIZATION\\n and str(g.func.name.name) not in MUTABLE_OPS_NOT_USING_FUNCTIONALIZATION\\n ):\\n assert g.has_composite_implicit_autograd_kernel or not modifies_arguments(g)\\n return []\\n else:\\n # Case 2: emit inplace -> out-of-place kernels for the functionalization pass\\n mutation_defs = []\\n mutation_defs.append(emit_inplace_functionalization_body(g.out, g))\\n if g.inplace is not None:\\n mutation_defs.append(emit_inplace_functionalization_body(g.inplace, g))\\n if g.mutable is not None:\\n mutation_defs.append(emit_inplace_functionalization_body(g.mutable, g))\\n return mutation_defs\\n return []\\n\\n\\nfrom __future__ import annotations\\n\\nimport re\\nfrom typing import Mapping, Sequence\\n\\n\\n# match $identifier or ${identifier} and replace with value in env\\n# If this identifier is at the beginning of whitespace on a line\\n# and its value is a list then it is treated as\\n# block substitution by indenting to that depth and putting each element\\n# of the list on its own line\\n# if the identifier is on a line starting with non-whitespace and a list\\n# then it is comma separated ${,foo} will insert a comma before the list\\n# if this list is not empty and ${foo,} will insert one after.\\n\\n\\nclass CodeTemplate:\\n substitution_str = r\\\"(^[^\\\\n\\\\S]*)?\\\\$([^\\\\d\\\\W]\\\\w*|\\\\{,?[^\\\\d\\\\W]\\\\w*\\\\,?})\\\"\\n substitution = re.compile(substitution_str, re.MULTILINE)\\n\\n pattern: str\\n filename: str\\n\\n @staticmethod\\n def from_file(filename: str) -> CodeTemplate:\\n with open(filename) as f:\\n return CodeTemplate(f.read(), filename)\\n\\n def __init__(self, pattern: str, filename: str = \\\"\\\") -> None:\\n self.pattern = pattern\\n self.filename = filename\\n\\n def substitute(\\n self, env: Mapping[str, object] | None = None, **kwargs: object\\n ) -> str:\\n if env is None:\\n env = {}\\n\\n def lookup(v: str) -> object:\\n assert env is not None\\n return kwargs[v] if v in kwargs else env[v]\\n\\n def indent_lines(indent: str, v: Sequence[object]) -> str:\\n return \\\"\\\".join(\\n [indent + l + \\\"\\\\n\\\" for e in v for l in str(e).splitlines()]\\n ).rstrip()\\n\\n def replace(match: re.Match[str]) -> str:\\n indent = match.group(1)\\n key = match.group(2)\\n comma_before = \\\"\\\"\\n comma_after = \\\"\\\"\\n if key[0] == \\\"{\\\":\\n key = key[1:-1]\\n if key[0] == \\\",\\\":\\n comma_before = \\\", \\\"\\n key = key[1:]\\n if key[-1] == \\\",\\\":\\n comma_after = \\\", \\\"\\n key = key[:-1]\\n v = lookup(key)\\n if indent is not None:\\n if not isinstance(v, list):\\n v = [v]\\n return indent_lines(indent, v)\\n elif isinstance(v, list):\\n middle = \\\", \\\".join([str(x) for x in v])\\n if len(v) == 0:\\n return middle\\n return comma_before + middle + comma_after\\n else:\\n return str(v)\\n\\n return self.substitution.sub(replace, self.pattern)\\n\\n\\nif __name__ == \\\"__main__\\\":\\n c = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n int foo($args) {\\n\\n $bar\\n $bar\\n $a+$b\\n }\\n int commatest(int a${,stuff})\\n int notest(int a${,empty,})\\n \\\"\\\"\\\"\\n )\\n print(\\n c.substitute(\\n args=[\\\"hi\\\", 8],\\n bar=[\\\"what\\\", 7],\\n a=3,\\n b=4,\\n stuff=[\\\"things...\\\", \\\"others\\\"],\\n empty=[],\\n )\\n )\\n\\n\\nfrom __future__ import annotations\\n\\nimport textwrap\\nfrom dataclasses import dataclass\\nfrom typing import Sequence\\n\\nfrom torchgen.api.translate import translate\\nfrom torchgen.api.types import DispatcherSignature\\nfrom torchgen.context import method_with_native_function\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n ListType,\\n NativeFunction,\\n OptionalType,\\n Return,\\n SchemaKind,\\n Type,\\n)\\nfrom torchgen.utils import mapMaybe\\n\\n\\ndef is_tensor(typ: Type) -> bool:\\n return isinstance(typ, BaseType) and typ.name == BaseTy.Tensor\\n\\n\\ndef is_optional_tensor(typ: Type) -> bool:\\n return isinstance(typ, OptionalType) and is_tensor(typ.elem)\\n\\n\\ndef is_tensor_list(typ: Type) -> bool:\\n return isinstance(typ, ListType) and is_tensor(typ.elem)\\n\\n\\ndef unwrap_tensor(name: str, cur_level_var: str) -> list[str]:\\n result = f\\\"\\\"\\\"\\\\\\n auto [{name}_value, {name}_bdim] = unwrapTensorAtLevel({name}, {cur_level_var});\\\"\\\"\\\"\\n return textwrap.dedent(result).split(\\\"\\\\n\\\")\\n\\n\\ndef unwrap_optional_tensor(name: str, cur_level_var: str) -> list[str]:\\n result = f\\\"\\\"\\\"\\\\\\n std::optional {name}_value;\\n std::optional {name}_bdim;\\n if ({name}) {{\\n std::tie({name}_value, {name}_bdim) = unwrapTensorAtLevel({name}.value(), {cur_level_var});\\n }}\\\"\\\"\\\"\\n return textwrap.dedent(result).split(\\\"\\\\n\\\")\\n\\n\\ndef gen_unwraps(\\n flat_arguments: Sequence[Argument], cur_level_var: str\\n) -> tuple[str, list[str]]:\\n arg_names = [a.name for a in flat_arguments]\\n arg_types = [a.type for a in flat_arguments]\\n\\n tensors = [name for typ, name in zip(arg_types, arg_names) if is_tensor(typ)]\\n optional_tensors = [\\n name for typ, name in zip(arg_types, arg_names) if is_optional_tensor(typ)\\n ]\\n\\n unwraps = []\\n for tensor in tensors:\\n unwraps += unwrap_tensor(tensor, cur_level_var)\\n\\n for opt_tensor in optional_tensors:\\n unwraps += unwrap_optional_tensor(opt_tensor, cur_level_var)\\n unwrap_code = \\\"\\\\n\\\".join(unwraps)\\n\\n unwrapped_arg_list = []\\n for arg in arg_names:\\n if arg in tensors or arg in optional_tensors:\\n unwrapped_arg_list += [f\\\"{arg}_value\\\", f\\\"{arg}_bdim\\\"]\\n else:\\n unwrapped_arg_list.append(arg)\\n return unwrap_code, unwrapped_arg_list\\n\\n\\ndef gen_case_where_all_bdims_are_none(\\n outer_sig: DispatcherSignature, schema: FunctionSchema, cur_level_var: str\\n) -> str:\\n conditions = []\\n flat_args = schema.arguments.flat_all\\n for arg in flat_args:\\n if not arg.type.is_tensor_like():\\n continue\\n conditions.append(f\\\"!isBatchedAtLevel({arg.name}, {cur_level_var})\\\")\\n\\n sig = DispatcherSignature.from_schema(schema)\\n translated_args = \\\", \\\".join(\\n e.expr for e in translate(outer_sig.arguments(), sig.arguments())\\n )\\n return f\\\"\\\"\\\"\\\\\\nif ({' && '.join(conditions)}) {{\\n return at::_ops::{sig.func.name.unambiguous_name()}::call({translated_args});\\n}}\\\"\\\"\\\"\\n\\n\\ndef gen_returns(\\n returns: tuple[Return, ...], cur_level_var: str, results_var: str\\n) -> str:\\n idx = 0\\n wrapped_returns = []\\n for ret in returns:\\n if is_tensor(ret.type):\\n wrapped_returns.append(\\n f\\\"makeBatched(std::get<{idx}>({results_var}), std::get<{idx + 1}>({results_var}), {cur_level_var})\\\"\\n )\\n idx += 2\\n elif is_tensor_list(ret.type):\\n wrapped_returns.append(\\n f\\\"makeBatchedVector(std::get<{idx}>({results_var}), std::get<{idx+1}>({results_var}), {cur_level_var})\\\"\\n )\\n idx += 2\\n else:\\n wrapped_returns.append(f\\\"std::get<{idx}>({results_var})\\\")\\n idx += 1\\n if len(wrapped_returns) == 1:\\n result = f\\\"return {wrapped_returns[0]};\\\"\\n else:\\n result = f'return std::make_tuple({\\\", \\\".join(wrapped_returns)});'\\n return result\\n\\n\\ndef accepts_at_least_one_tensor_input(schema: FunctionSchema) -> bool:\\n return any(a.type.is_tensor_like() for a in schema.arguments.flat_all)\\n\\n\\ndef is_mutated_arg(argument: Argument) -> bool:\\n return argument.annotation is not None and argument.annotation.is_write\\n\\n\\ndef gen_vmap_inplace_plumbing(native_function: NativeFunction) -> str | None:\\n # Assumptions:\\n # - only one argument is being modified in-place\\n # - the argument that is being modified in-place is the first argument\\n # - all returns are either Tensor, tuple of Tensor, or TensorList\\n schema = native_function.func\\n sig = DispatcherSignature.from_schema(schema)\\n returns = schema.returns\\n\\n # Check assumptions. If these are invalid we return None\\n # and punt the work to handle them to the future.\\n assert schema.kind() == SchemaKind.inplace\\n if not is_mutated_arg(schema.arguments.flat_all[0]):\\n return None\\n if not len([arg for arg in schema.arguments.flat_all if is_mutated_arg(arg)]) == 1:\\n return None\\n\\n # Only support cases where all returns are Tensors or vector\\n if len(returns) == 0:\\n return None\\n if not all(is_tensor(ret.type) or is_tensor_list(ret.type) for ret in returns):\\n return None\\n if not accepts_at_least_one_tensor_input(schema):\\n return None\\n\\n cur_level_var = \\\"cur_level\\\"\\n\\n unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var)\\n bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)\\n\\n return f\\\"\\\"\\\"\\\\\\ntemplate \\n{sig.decl(name=schema.name.unambiguous_name() + '_generated_plumbing')} {{\\n c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);\\n auto maybe_layer = maybeCurrentDynamicLayer();\\n vmap_check_escaped(maybe_layer, \\\"gen_vmap_inplace_plumbing\\\");\\n int64_t {cur_level_var} = maybe_layer->layerId();\\n{textwrap.indent(bdims_all_none_case, \\\" \\\")}\\n{textwrap.indent(unwraps, \\\" \\\")}\\n batch_rule({', '.join(unwrapped_arg_list)});\\n return {schema.arguments.flat_all[0].name};\\n}}\\\"\\\"\\\"\\n\\n\\ndef gen_vmap_plumbing_no_returns(native_function: NativeFunction) -> str:\\n schema = native_function.func\\n sig = DispatcherSignature.from_schema(schema)\\n cur_level_var = \\\"cur_level\\\"\\n\\n unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var)\\n bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)\\n\\n return f\\\"\\\"\\\"\\\\\\ntemplate \\n{sig.decl(name=schema.name.unambiguous_name() + '_generated_plumbing')} {{\\n c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);\\n auto maybe_layer = maybeCurrentDynamicLayer();\\n vmap_check_escaped(maybe_layer, \\\"gen_vmap_plumbing_no_returns\\\");\\n int64_t {cur_level_var} = maybe_layer->layerId();\\n{textwrap.indent(bdims_all_none_case, \\\" \\\")}\\n{textwrap.indent(unwraps, \\\" \\\")}\\n batch_rule({', '.join(unwrapped_arg_list)});\\n}}\\\"\\\"\\\"\\n\\n\\ndef gen_vmap_plumbing(native_function: NativeFunction) -> str | None:\\n schema = native_function.func\\n sig = DispatcherSignature.from_schema(schema)\\n returns = schema.returns\\n\\n # Only support cases where all returns are Tensors or vector\\n if not accepts_at_least_one_tensor_input(schema):\\n return None\\n if len(returns) == 0:\\n return gen_vmap_plumbing_no_returns(native_function)\\n return_symint_overrides = [\\n \\\"_scaled_dot_product_flash_attention\\\",\\n \\\"_scaled_dot_product_cudnn_attention\\\",\\n ]\\n if (\\n not all(ret.type.is_tensor_like() for ret in returns)\\n and schema.name.unambiguous_name() not in return_symint_overrides\\n ):\\n return None\\n # in-place views need special handling\\n if \\\"inplace_view\\\" in native_function.tags:\\n return None\\n\\n if schema.kind() == SchemaKind.inplace:\\n return gen_vmap_inplace_plumbing(native_function)\\n\\n # Don't support these (mutable, out, scratch)\\n if schema.kind() != SchemaKind.functional:\\n return None\\n\\n results_var = \\\"results\\\"\\n cur_level_var = \\\"cur_level\\\"\\n\\n unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var)\\n bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var)\\n\\n wrapped_returns = gen_returns(returns, cur_level_var, results_var)\\n return f\\\"\\\"\\\"\\\\\\ntemplate \\n{sig.decl(name=schema.name.unambiguous_name() + '_generated_plumbing')} {{\\n c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched);\\n auto maybe_layer = maybeCurrentDynamicLayer();\\n vmap_check_escaped(maybe_layer, \\\"gen_vmap_plumbing\\\");\\n int64_t {cur_level_var} = maybe_layer->layerId();\\n{textwrap.indent(bdims_all_none_case, \\\" \\\")}\\n{textwrap.indent(unwraps, \\\" \\\")}\\n auto {results_var} = batch_rule({', '.join(unwrapped_arg_list)});\\n {wrapped_returns}\\n}}\\\"\\\"\\\"\\n\\n\\n@dataclass(frozen=True)\\nclass ComputeBatchRulePlumbing:\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str | None:\\n result = gen_vmap_plumbing(f)\\n return result\\n\\n\\ndef gen_all_vmap_plumbing(native_functions: Sequence[NativeFunction]) -> str:\\n body = \\\"\\\\n\\\".join(list(mapMaybe(ComputeBatchRulePlumbing(), native_functions)))\\n return f\\\"\\\"\\\"\\n#pragma once\\n#include \\n#include \\n\\nnamespace at {{ namespace functorch {{\\n\\n{body}\\n\\n}}}} // namespace at::functorch\\n\\\"\\\"\\\"\\n\\n\\nfrom typing import Any, Optional, Tuple, Union\\n\\nfrom torchgen.model import (\\n Annotation,\\n Argument,\\n Arguments,\\n BaseOperatorName,\\n BaseTy,\\n BaseType,\\n CustomClassType,\\n FunctionSchema,\\n ListType,\\n OperatorName,\\n Return,\\n)\\n\\n\\n# Note: These aren't actually used in torchgen, they're some utilities for generating a schema\\n# from real arguments. For example, this is used to generate HigherOrderOperators' schema since\\n# their schemas can vary for different instances of the same HOP.\\n\\n\\nclass TypeGen:\\n convert_to_base_ty = {\\n int: BaseTy.int,\\n float: BaseTy.float,\\n str: BaseTy.str,\\n bool: BaseTy.bool,\\n }\\n\\n @staticmethod\\n def from_example(obj: Any) -> Union[BaseType, ListType, CustomClassType]:\\n import torch\\n\\n if isinstance(obj, torch.fx.GraphModule):\\n return BaseType(BaseTy.GraphModule)\\n elif isinstance(obj, torch.Tensor):\\n return BaseType(BaseTy.Tensor)\\n elif isinstance(obj, torch.SymInt):\\n return BaseType(BaseTy.SymInt)\\n elif isinstance(obj, torch.SymBool):\\n return BaseType(BaseTy.SymBool)\\n elif isinstance(obj, torch.ScriptObject):\\n return CustomClassType(obj._type().name()) # type: ignore[attr-defined]\\n elif isinstance(obj, (list, tuple)):\\n assert len(obj) > 0\\n all_base_tys = [TypeGen.from_example(x) for x in obj]\\n if len(set(all_base_tys)) > 1:\\n raise RuntimeError(\\n f\\\"Cannot generate schema for a seqeunce of args of heterogeneous types: {all_base_tys}. \\\"\\n \\\"Consider unpacking the argument and give proper names to them if possible \\\"\\n \\\"instead of using *args.\\\"\\n )\\n return ListType(all_base_tys[0], len(obj))\\n tp = type(obj)\\n if tp not in TypeGen.convert_to_base_ty:\\n raise RuntimeError(f\\\"unsupported type {tp}\\\")\\n return BaseType(TypeGen.convert_to_base_ty[tp])\\n\\n\\nclass ReturnGen:\\n @staticmethod\\n def from_example(\\n name: Optional[str], obj: Any, annotation: Optional[Annotation]\\n ) -> Return:\\n return Return(name, TypeGen.from_example(obj), annotation)\\n\\n\\nclass ArgumentGen:\\n @staticmethod\\n def from_example(\\n name: str, obj: Any, default: Optional[str], annotation: Optional[Annotation]\\n ) -> Argument:\\n return Argument(\\n name, TypeGen.from_example(obj), default=default, annotation=annotation\\n )\\n\\n\\nclass FunctionSchemaGen:\\n @staticmethod\\n def from_example(\\n op_name: str,\\n example_inputs: Tuple[Tuple[str, Any], ...],\\n example_outputs: Tuple[Any, ...],\\n ) -> FunctionSchema:\\n args = []\\n for name, inp in example_inputs:\\n args.append(ArgumentGen.from_example(name, inp, None, None))\\n # ignore the annotations and other attributes for now, we could add more when needed.\\n arguments = Arguments(\\n tuple(), None, tuple(args), tuple(), None, tuple(), tuple()\\n )\\n returns = tuple(\\n ReturnGen.from_example(None, out, None) for out in example_outputs\\n )\\n op_name = OperatorName(BaseOperatorName(op_name, False, False, False), \\\"\\\")\\n return FunctionSchema(op_name, arguments, returns)\\n\\n\\nfrom __future__ import annotations\\n\\nimport contextlib\\nimport functools\\nimport hashlib\\nimport os\\nimport re\\nimport sys\\nimport textwrap\\nfrom dataclasses import fields, is_dataclass\\nfrom enum import auto, Enum\\nfrom pathlib import Path\\nfrom typing import (\\n Any,\\n Callable,\\n Generic,\\n Iterable,\\n Iterator,\\n Literal,\\n NoReturn,\\n Sequence,\\n TYPE_CHECKING,\\n TypeVar,\\n)\\nfrom typing_extensions import Self\\n\\nfrom torchgen.code_template import CodeTemplate\\n\\n\\nif TYPE_CHECKING:\\n from argparse import Namespace\\n\\n\\nREPO_ROOT = Path(__file__).absolute().parent.parent\\n\\n\\n# Many of these functions share logic for defining both the definition\\n# and declaration (for example, the function signature is the same), so\\n# we organize them into one function that takes a Target to say which\\n# code we want.\\n#\\n# This is an OPEN enum (we may add more cases to it in the future), so be sure\\n# to explicitly specify with Literal[Target.XXX] or Literal[Target.XXX, Target.YYY]\\n# what targets are valid for your use.\\nclass Target(Enum):\\n # top level namespace (not including at)\\n DEFINITION = auto()\\n DECLARATION = auto()\\n # TORCH_LIBRARY(...) { ... }\\n REGISTRATION = auto()\\n # namespace { ... }\\n ANONYMOUS_DEFINITION = auto()\\n # namespace cpu { ... }\\n NAMESPACED_DEFINITION = auto()\\n NAMESPACED_DECLARATION = auto()\\n\\n\\n# Matches \\\"foo\\\" in \\\"foo, bar\\\" but not \\\"foobar\\\". Used to search for the\\n# occurrence of a parameter in the derivative formula\\nIDENT_REGEX = r\\\"(^|\\\\W){}($|\\\\W)\\\"\\n\\n\\n# TODO: Use a real parser here; this will get bamboozled\\ndef split_name_params(schema: str) -> tuple[str, list[str]]:\\n m = re.match(r\\\"(\\\\w+)(\\\\.\\\\w+)?\\\\((.*)\\\\)\\\", schema)\\n if m is None:\\n raise RuntimeError(f\\\"Unsupported function schema: {schema}\\\")\\n name, _, params = m.groups()\\n return name, params.split(\\\", \\\")\\n\\n\\nT = TypeVar(\\\"T\\\")\\nS = TypeVar(\\\"S\\\")\\n\\n# These two functions purposely return generators in analogy to map()\\n# so that you don't mix up when you need to list() them\\n\\n\\n# Map over function that may return None; omit Nones from output sequence\\ndef mapMaybe(func: Callable[[T], S | None], xs: Iterable[T]) -> Iterator[S]:\\n for x in xs:\\n r = func(x)\\n if r is not None:\\n yield r\\n\\n\\n# Map over function that returns sequences and cat them all together\\ndef concatMap(func: Callable[[T], Sequence[S]], xs: Iterable[T]) -> Iterator[S]:\\n for x in xs:\\n yield from func(x)\\n\\n\\n# Conveniently add error context to exceptions raised. Lets us\\n# easily say that an error occurred while processing a specific\\n# context.\\n@contextlib.contextmanager\\ndef context(msg_fn: Callable[[], str]) -> Iterator[None]:\\n try:\\n yield\\n except Exception as e:\\n # TODO: this does the wrong thing with KeyError\\n msg = msg_fn()\\n msg = textwrap.indent(msg, \\\" \\\")\\n msg = f\\\"{e.args[0]}\\\\n{msg}\\\" if e.args else msg\\n e.args = (msg,) + e.args[1:]\\n raise\\n\\n\\n# A little trick from https://github.com/python/mypy/issues/6366\\n# for getting mypy to do exhaustiveness checking\\n# TODO: put this somewhere else, maybe\\ndef assert_never(x: NoReturn) -> NoReturn:\\n raise AssertionError(f\\\"Unhandled type: {type(x).__name__}\\\")\\n\\n\\n@functools.lru_cache(maxsize=None)\\ndef _read_template(template_fn: str) -> CodeTemplate:\\n return CodeTemplate.from_file(template_fn)\\n\\n\\n# String hash that's stable across different executions, unlike builtin hash\\ndef string_stable_hash(s: str) -> int:\\n sha1 = hashlib.sha1(s.encode(\\\"latin1\\\")).digest()\\n return int.from_bytes(sha1, byteorder=\\\"little\\\")\\n\\n\\n# A small abstraction for writing out generated files and keeping track\\n# of what files have been written (so you can write out a list of output\\n# files)\\nclass FileManager:\\n install_dir: str\\n template_dir: str\\n dry_run: bool\\n filenames: set[str]\\n\\n def __init__(self, install_dir: str, template_dir: str, dry_run: bool) -> None:\\n self.install_dir = install_dir\\n self.template_dir = template_dir\\n self.filenames = set()\\n self.dry_run = dry_run\\n\\n def _write_if_changed(self, filename: str, contents: str) -> None:\\n old_contents: str | None\\n try:\\n with open(filename) as f:\\n old_contents = f.read()\\n except OSError:\\n old_contents = None\\n if contents != old_contents:\\n # Create output directory if it doesn't exist\\n os.makedirs(os.path.dirname(filename), exist_ok=True)\\n with open(filename, \\\"w\\\") as f:\\n f.write(contents)\\n\\n # Read from template file and replace pattern with callable (type could be dict or str).\\n def substitute_with_template(\\n self, template_fn: str, env_callable: Callable[[], str | dict[str, Any]]\\n ) -> str:\\n template_path = os.path.join(self.template_dir, template_fn)\\n env = env_callable()\\n if isinstance(env, dict):\\n if \\\"generated_comment\\\" not in env:\\n generator_default = REPO_ROOT / \\\"torchgen\\\" / \\\"gen.py\\\"\\n try:\\n generator = Path(\\n sys.modules[\\\"__main__\\\"].__file__ or generator_default\\n ).absolute()\\n except (KeyError, AttributeError):\\n generator = generator_default.absolute()\\n\\n try:\\n generator_path = generator.relative_to(REPO_ROOT).as_posix()\\n except ValueError:\\n generator_path = generator.name\\n\\n env = {\\n **env, # copy the original dict instead of mutating it\\n \\\"generated_comment\\\": (\\n \\\"@\\\" + f\\\"generated by {generator_path} from {template_fn}\\\"\\n ),\\n }\\n template = _read_template(template_path)\\n return template.substitute(env)\\n elif isinstance(env, str):\\n return env\\n else:\\n assert_never(env)\\n\\n def write_with_template(\\n self,\\n filename: str,\\n template_fn: str,\\n env_callable: Callable[[], str | dict[str, Any]],\\n ) -> None:\\n filename = f\\\"{self.install_dir}/{filename}\\\"\\n assert filename not in self.filenames, \\\"duplicate file write {filename}\\\"\\n self.filenames.add(filename)\\n if not self.dry_run:\\n substitute_out = self.substitute_with_template(\\n template_fn=template_fn,\\n env_callable=env_callable,\\n )\\n self._write_if_changed(filename=filename, contents=substitute_out)\\n\\n def write(\\n self,\\n filename: str,\\n env_callable: Callable[[], str | dict[str, Any]],\\n ) -> None:\\n self.write_with_template(filename, filename, env_callable)\\n\\n def write_sharded(\\n self,\\n filename: str,\\n items: Iterable[T],\\n *,\\n key_fn: Callable[[T], str],\\n env_callable: Callable[[T], dict[str, list[str]]],\\n num_shards: int,\\n base_env: dict[str, Any] | None = None,\\n sharded_keys: set[str],\\n ) -> None:\\n everything: dict[str, Any] = {\\\"shard_id\\\": \\\"Everything\\\"}\\n shards: list[dict[str, Any]] = [\\n {\\\"shard_id\\\": f\\\"_{i}\\\"} for i in range(num_shards)\\n ]\\n all_shards = [everything] + shards\\n\\n if base_env is not None:\\n for shard in all_shards:\\n shard.update(base_env)\\n\\n for key in sharded_keys:\\n for shard in all_shards:\\n if key in shard:\\n assert isinstance(\\n shard[key], list\\n ), \\\"sharded keys in base_env must be a list\\\"\\n shard[key] = shard[key].copy()\\n else:\\n shard[key] = []\\n\\n def merge_env(into: dict[str, list[str]], from_: dict[str, list[str]]) -> None:\\n for k, v in from_.items():\\n assert k in sharded_keys, f\\\"undeclared sharded key {k}\\\"\\n into[k] += v\\n\\n if self.dry_run:\\n # Dry runs don't write any templates, so incomplete environments are fine\\n items = ()\\n\\n for item in items:\\n key = key_fn(item)\\n sid = string_stable_hash(key) % num_shards\\n env = env_callable(item)\\n\\n merge_env(shards[sid], env)\\n merge_env(everything, env)\\n\\n dot_pos = filename.rfind(\\\".\\\")\\n if dot_pos == -1:\\n dot_pos = len(filename)\\n base_filename = filename[:dot_pos]\\n extension = filename[dot_pos:]\\n\\n for shard in all_shards:\\n shard_id = shard[\\\"shard_id\\\"]\\n self.write_with_template(\\n f\\\"{base_filename}{shard_id}{extension}\\\", filename, lambda: shard\\n )\\n\\n # filenames is used to track compiled files, but FooEverything.cpp isn't meant to be compiled\\n self.filenames.discard(\\n f\\\"{self.install_dir}/{base_filename}Everything{extension}\\\"\\n )\\n\\n def write_outputs(self, variable_name: str, filename: str) -> None:\\n \\\"\\\"\\\"Write a file containing the list of all outputs which are\\n generated by this script.\\\"\\\"\\\"\\n content = \\\"set({}\\\\n {})\\\".format(\\n variable_name,\\n \\\"\\\\n \\\".join('\\\"' + name + '\\\"' for name in sorted(self.filenames)),\\n )\\n self._write_if_changed(filename, content)\\n\\n def template_dir_for_comments(self) -> str:\\n \\\"\\\"\\\"\\n This needs to be deterministic. The template dir is an absolute path\\n that varies across builds. So, just use the path relative to this file,\\n which will point to the codegen source but will be stable.\\n \\\"\\\"\\\"\\n return os.path.relpath(self.template_dir, os.path.dirname(__file__))\\n\\n\\n# Helper function to generate file manager\\ndef make_file_manager(\\n options: Namespace, install_dir: str | None = None\\n) -> FileManager:\\n template_dir = os.path.join(options.source_path, \\\"templates\\\")\\n install_dir = install_dir if install_dir else options.install_dir\\n return FileManager(\\n install_dir=install_dir, template_dir=template_dir, dry_run=options.dry_run\\n )\\n\\n\\n# Helper function to create a pretty representation for dataclasses\\ndef dataclass_repr(\\n obj: Any,\\n indent: int = 0,\\n width: int = 80,\\n) -> str:\\n # built-in pprint module support dataclasses from python 3.10\\n if sys.version_info >= (3, 10):\\n from pprint import pformat\\n\\n return pformat(obj, indent, width)\\n\\n return _pformat(obj, indent=indent, width=width)\\n\\n\\ndef _pformat(\\n obj: Any,\\n indent: int,\\n width: int,\\n curr_indent: int = 0,\\n) -> str:\\n assert is_dataclass(obj), f\\\"obj should be a dataclass, received: {type(obj)}\\\"\\n\\n class_name = obj.__class__.__name__\\n # update current indentation level with class name\\n curr_indent += len(class_name) + 1\\n\\n fields_list = [(f.name, getattr(obj, f.name)) for f in fields(obj) if f.repr]\\n\\n fields_str = []\\n for name, attr in fields_list:\\n # update the current indent level with the field name\\n # dict, list, set and tuple also add indent as done in pprint\\n _curr_indent = curr_indent + len(name) + 1\\n if is_dataclass(attr):\\n str_repr = _pformat(attr, indent, width, _curr_indent)\\n elif isinstance(attr, dict):\\n str_repr = _format_dict(attr, indent, width, _curr_indent)\\n elif isinstance(attr, (list, set, tuple)):\\n str_repr = _format_list(attr, indent, width, _curr_indent)\\n else:\\n str_repr = repr(attr)\\n\\n fields_str.append(f\\\"{name}={str_repr}\\\")\\n\\n indent_str = curr_indent * \\\" \\\"\\n body = f\\\",\\\\n{indent_str}\\\".join(fields_str)\\n return f\\\"{class_name}({body})\\\"\\n\\n\\ndef _format_dict(\\n attr: dict[Any, Any],\\n indent: int,\\n width: int,\\n curr_indent: int,\\n) -> str:\\n curr_indent += indent + 3\\n dict_repr = []\\n for k, v in attr.items():\\n k_repr = repr(k)\\n v_str = (\\n _pformat(v, indent, width, curr_indent + len(k_repr))\\n if is_dataclass(v)\\n else repr(v)\\n )\\n dict_repr.append(f\\\"{k_repr}: {v_str}\\\")\\n\\n return _format(dict_repr, indent, width, curr_indent, \\\"{\\\", \\\"}\\\")\\n\\n\\ndef _format_list(\\n attr: list[Any] | set[Any] | tuple[Any, ...],\\n indent: int,\\n width: int,\\n curr_indent: int,\\n) -> str:\\n curr_indent += indent + 1\\n list_repr = [\\n _pformat(l, indent, width, curr_indent) if is_dataclass(l) else repr(l)\\n for l in attr\\n ]\\n start, end = (\\\"[\\\", \\\"]\\\") if isinstance(attr, list) else (\\\"(\\\", \\\")\\\")\\n return _format(list_repr, indent, width, curr_indent, start, end)\\n\\n\\ndef _format(\\n fields_str: list[str],\\n indent: int,\\n width: int,\\n curr_indent: int,\\n start: str,\\n end: str,\\n) -> str:\\n delimiter, curr_indent_str = \\\"\\\", \\\"\\\"\\n # if it exceed the max width then we place one element per line\\n if len(repr(fields_str)) >= width:\\n delimiter = \\\"\\\\n\\\"\\n curr_indent_str = \\\" \\\" * curr_indent\\n\\n indent_str = \\\" \\\" * indent\\n body = f\\\", {delimiter}{curr_indent_str}\\\".join(fields_str)\\n return f\\\"{start}{indent_str}{body}{end}\\\"\\n\\n\\nclass NamespaceHelper:\\n \\\"\\\"\\\"A helper for constructing the namespace open and close strings for a nested set of namespaces.\\n\\n e.g. for namespace_str torch::lazy,\\n\\n prologue:\\n namespace torch {\\n namespace lazy {\\n\\n epilogue:\\n } // namespace lazy\\n } // namespace torch\\n \\\"\\\"\\\"\\n\\n def __init__(\\n self, namespace_str: str, entity_name: str = \\\"\\\", max_level: int = 2\\n ) -> None:\\n # cpp_namespace can be a colon joined string such as torch::lazy\\n cpp_namespaces = namespace_str.split(\\\"::\\\")\\n assert (\\n len(cpp_namespaces) <= max_level\\n ), f\\\"Codegen doesn't support more than {max_level} level(s) of custom namespace. Got {namespace_str}.\\\"\\n self.cpp_namespace_ = namespace_str\\n self.prologue_ = \\\"\\\\n\\\".join([f\\\"namespace {n} {{\\\" for n in cpp_namespaces])\\n self.epilogue_ = \\\"\\\\n\\\".join(\\n [f\\\"}} // namespace {n}\\\" for n in reversed(cpp_namespaces)]\\n )\\n self.namespaces_ = cpp_namespaces\\n self.entity_name_ = entity_name\\n\\n @staticmethod\\n def from_namespaced_entity(\\n namespaced_entity: str, max_level: int = 2\\n ) -> NamespaceHelper:\\n \\\"\\\"\\\"\\n Generate helper from nested namespaces as long as class/function name. E.g.: \\\"torch::lazy::add\\\"\\n \\\"\\\"\\\"\\n names = namespaced_entity.split(\\\"::\\\")\\n entity_name = names[-1]\\n namespace_str = \\\"::\\\".join(names[:-1])\\n return NamespaceHelper(\\n namespace_str=namespace_str, entity_name=entity_name, max_level=max_level\\n )\\n\\n @property\\n def prologue(self) -> str:\\n return self.prologue_\\n\\n @property\\n def epilogue(self) -> str:\\n return self.epilogue_\\n\\n @property\\n def entity_name(self) -> str:\\n return self.entity_name_\\n\\n # Only allow certain level of namespaces\\n def get_cpp_namespace(self, default: str = \\\"\\\") -> str:\\n \\\"\\\"\\\"\\n Return the namespace string from joining all the namespaces by \\\"::\\\" (hence no leading \\\"::\\\").\\n Return default if namespace string is empty.\\n \\\"\\\"\\\"\\n return self.cpp_namespace_ if self.cpp_namespace_ else default\\n\\n\\nclass OrderedSet(Generic[T]):\\n storage: dict[T, Literal[None]]\\n\\n def __init__(self, iterable: Iterable[T] | None = None) -> None:\\n if iterable is None:\\n self.storage = {}\\n else:\\n self.storage = dict.fromkeys(iterable)\\n\\n def __contains__(self, item: T) -> bool:\\n return item in self.storage\\n\\n def __iter__(self) -> Iterator[T]:\\n return iter(self.storage.keys())\\n\\n def update(self, items: OrderedSet[T]) -> None:\\n self.storage.update(items.storage)\\n\\n def add(self, item: T) -> None:\\n self.storage[item] = None\\n\\n def copy(self) -> OrderedSet[T]:\\n ret: OrderedSet[T] = OrderedSet()\\n ret.storage = self.storage.copy()\\n return ret\\n\\n @staticmethod\\n def union(*args: OrderedSet[T]) -> OrderedSet[T]:\\n ret = args[0].copy()\\n for s in args[1:]:\\n ret.update(s)\\n return ret\\n\\n def __or__(self, other: OrderedSet[T]) -> OrderedSet[T]:\\n return OrderedSet.union(self, other)\\n\\n def __ior__(self, other: OrderedSet[T]) -> Self:\\n self.update(other)\\n return self\\n\\n def __eq__(self, other: object) -> bool:\\n if isinstance(other, OrderedSet):\\n return self.storage == other.storage\\n else:\\n return set(self.storage.keys()) == other\\n\\n\\nfrom __future__ import annotations\\n\\nimport threading\\nfrom contextlib import contextmanager\\nfrom typing import Iterator\\n\\n\\n# Simple dynamic scoping implementation. The name \\\"parametrize\\\" comes\\n# from Racket.\\n#\\n# WARNING WARNING: LOOKING TO EDIT THIS FILE? Think carefully about\\n# why you need to add a toggle to the global behavior of code\\n# generation. The parameters here should really only be used\\n# for \\\"temporary\\\" situations, where we need to temporarily change\\n# the codegen in some cases because we cannot conveniently update\\n# all call sites, and are slated to be eliminated once all call\\n# sites are eliminated. If you don't have a plan for how to get there,\\n# DON'T add a new entry here.\\n\\n\\nclass Locals(threading.local):\\n use_const_ref_for_mutable_tensors: bool | None = None\\n use_ilistref_for_tensor_lists: bool | None = None\\n\\n\\n_locals = Locals()\\n\\n\\ndef use_const_ref_for_mutable_tensors() -> bool:\\n assert _locals.use_const_ref_for_mutable_tensors is not None, (\\n \\\"need to initialize local.use_const_ref_for_mutable_tensors with \\\"\\n \\\"local.parametrize\\\"\\n )\\n return _locals.use_const_ref_for_mutable_tensors\\n\\n\\ndef use_ilistref_for_tensor_lists() -> bool:\\n assert _locals.use_ilistref_for_tensor_lists is not None, (\\n \\\"need to initialize local.use_ilistref_for_tensor_lists with \\\"\\n \\\"local.parametrize\\\"\\n )\\n return _locals.use_ilistref_for_tensor_lists\\n\\n\\n@contextmanager\\ndef parametrize(\\n *, use_const_ref_for_mutable_tensors: bool, use_ilistref_for_tensor_lists: bool\\n) -> Iterator[None]:\\n old_use_const_ref_for_mutable_tensors = _locals.use_const_ref_for_mutable_tensors\\n old_use_ilistref_for_tensor_lists = _locals.use_ilistref_for_tensor_lists\\n try:\\n _locals.use_const_ref_for_mutable_tensors = use_const_ref_for_mutable_tensors\\n _locals.use_ilistref_for_tensor_lists = use_ilistref_for_tensor_lists\\n yield\\n finally:\\n _locals.use_const_ref_for_mutable_tensors = (\\n old_use_const_ref_for_mutable_tensors\\n )\\n _locals.use_ilistref_for_tensor_lists = old_use_ilistref_for_tensor_lists\\n\\n\\nfrom __future__ import annotations\\n\\nimport argparse\\nimport functools\\nimport json\\nimport os\\nfrom collections import defaultdict, namedtuple, OrderedDict\\nfrom dataclasses import dataclass, field\\nfrom pathlib import Path\\nfrom typing import Any, Callable, Literal, Sequence, TypeVar\\n\\nimport yaml\\n\\nimport torchgen.api.dispatcher as dispatcher\\nimport torchgen.api.meta as meta\\nimport torchgen.api.native as native\\nimport torchgen.api.structured as structured\\nimport torchgen.dest as dest\\nfrom torchgen.aoti.fallback_ops import inductor_fallback_ops\\nfrom torchgen.api import cpp\\nfrom torchgen.api.translate import translate\\nfrom torchgen.api.types import (\\n Binding,\\n CppSignature,\\n CppSignatureGroup,\\n DispatcherSignature,\\n NamedCType,\\n NativeSignature,\\n SpecialArgName,\\n)\\nfrom torchgen.context import (\\n method_with_native_function,\\n native_function_manager,\\n with_native_function,\\n with_native_function_and_indices,\\n)\\nfrom torchgen.gen_aoti_c_shim import (\\n gen_aoti_c_shim,\\n gen_static_dispatch_backend_call_signature,\\n get_fallback_op_name,\\n get_header_for_aoti,\\n)\\nfrom torchgen.gen_functionalization_type import (\\n gen_functionalization_definition,\\n gen_functionalization_registration,\\n gen_functionalization_view_inverse_declaration,\\n GenCompositeViewCopyKernel,\\n)\\nfrom torchgen.gen_vmap_plumbing import gen_all_vmap_plumbing\\nfrom torchgen.model import (\\n Argument,\\n BackendIndex,\\n BackendMetadata,\\n BaseOperatorName,\\n DEFAULT_KERNEL_NAMESPACE,\\n DispatchKey,\\n FRAGMENT_NAMESPACES,\\n FunctionSchema,\\n is_cuda_dispatch_key,\\n is_generic_dispatch_key,\\n is_ufunc_dispatch_key,\\n is_xpu_dispatch_key,\\n Location,\\n NativeFunction,\\n NativeFunctionsGroup,\\n NativeFunctionsViewGroup,\\n OperatorName,\\n OptionalType,\\n SchemaKind,\\n SelfArgument,\\n STRUCTURED_DISPATCH_KEYS,\\n TensorOptionsArguments,\\n Type,\\n Variant,\\n ViewSchemaKind,\\n)\\nfrom torchgen.native_function_generation import (\\n add_generated_native_functions,\\n gen_composite_functional_kernel,\\n gen_composite_out_kernel,\\n pre_group_native_functions,\\n)\\nfrom torchgen.selective_build.selector import SelectiveBuilder\\nfrom torchgen.utils import (\\n assert_never,\\n concatMap,\\n context,\\n FileManager,\\n make_file_manager,\\n mapMaybe,\\n NamespaceHelper,\\n Target,\\n)\\nfrom torchgen.yaml_utils import YamlDumper, YamlLoader\\n\\n\\nT = TypeVar(\\\"T\\\")\\n\\n# Welcome to the ATen code generator v2! The ATen code generator is\\n# responsible for parsing native_functions.yaml and then generating\\n# various generated files (e.g., TypeDefault.cpp) based on the operators\\n# defined in this file. This means that the code generator knows how to\\n# parse function schema, and then translate this into various C++ types\\n# and boilerplate code.\\n#\\n# Some things to know about this file when you modify it:\\n#\\n# - This file has STRICT mypy typechecking. Typecheck it with\\n# `mypy --config mypy-strict.ini` in the root source directory\\n#\\n# - Most of the heavy lifting lives in external modules:\\n# - 'model' has the data model for native_functions.yaml. The classes\\n# in those file represent what you see when you look at\\n# a native_functions.yaml\\n# - 'api' has conversions for how to translate JIT schema into\\n# the various C++ APIs that the codegen interacts with. There\\n# are in fact THREE different C++ APIs: the public C++ API,\\n# the dispatcher API, and the legacy dispatcher API. See each\\n# of these respective files for more information\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# HELPER FUNCTIONS\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\n# A custom loader for YAML to let us also keep track of line numbers\\n# of each entry in the YAML file\\nclass LineLoader(YamlLoader):\\n def construct_mapping(self, node, deep=False): # type: ignore[no-untyped-def]\\n mapping = super().construct_mapping(node, deep=deep) # type: ignore[no-untyped-call]\\n # Add 1 so line numbering starts at 1\\n mapping[\\\"__line__\\\"] = node.start_mark.line + 1\\n return mapping\\n\\n\\n# Parse native_functions.yaml into a sequence of NativeFunctions and Backend Indices.\\nParsedYaml = namedtuple(\\\"ParsedYaml\\\", [\\\"native_functions\\\", \\\"backend_indices\\\"])\\n\\n\\n_GLOBAL_PARSE_NATIVE_YAML_CACHE: dict[str, ParsedYaml] = {}\\n_GLOBAL_PARSE_TAGS_YAML_CACHE: dict[str, set[str]] = {}\\n\\n\\ndef parse_native_yaml_struct(\\n es: object,\\n valid_tags: set[str],\\n ignore_keys: set[DispatchKey] | None = None,\\n path: str = \\\"\\\",\\n skip_native_fns_gen: bool = False,\\n) -> ParsedYaml:\\n assert isinstance(es, list)\\n rs: list[NativeFunction] = []\\n bs: dict[DispatchKey, dict[OperatorName, BackendMetadata]] = defaultdict(dict)\\n for e in es:\\n assert isinstance(e, dict), f\\\"expected to be dict: {e}\\\"\\n assert isinstance(e.get(\\\"__line__\\\"), int), e\\n loc = Location(path, e[\\\"__line__\\\"])\\n funcs = e.get(\\\"func\\\")\\n assert funcs is not None, f\\\"missed 'func' in {e}\\\"\\n with context(lambda: f\\\"in {loc}:\\\\n {funcs}\\\"):\\n func, m = NativeFunction.from_yaml(e, loc, valid_tags, ignore_keys)\\n rs.append(func)\\n BackendIndex.grow_index(bs, m)\\n error_check_native_functions(rs)\\n # Default dict is to prevent the codegen from barfing when we have a dispatch key that has no kernels yet.\\n indices: dict[DispatchKey, BackendIndex] = defaultdict(\\n lambda: BackendIndex(\\n dispatch_key=DispatchKey.Undefined,\\n use_out_as_primary=True,\\n external=False,\\n device_guard=False,\\n # I'm actually not sure about this; undefined could be hit on\\n # empty TensorList, hypothetically that could have sizes in it\\n index={},\\n )\\n )\\n if not skip_native_fns_gen:\\n add_generated_native_functions(rs, bs)\\n for k, v in bs.items():\\n # All structured in-tree operators are implemented in terms of their out operator.\\n indices[k] = BackendIndex(\\n dispatch_key=k,\\n use_out_as_primary=True,\\n external=False,\\n # Only cuda-like devices in tree require device guards\\n device_guard=is_cuda_dispatch_key(k) or is_xpu_dispatch_key(k),\\n index=v,\\n )\\n return ParsedYaml(rs, indices)\\n\\n\\ndef parse_tags_yaml_struct(es: object, path: str = \\\"\\\") -> set[str]:\\n assert isinstance(es, list)\\n rs: set[str] = set()\\n for e in es:\\n assert isinstance(e.get(\\\"__line__\\\"), int), e\\n loc = Location(path, e[\\\"__line__\\\"])\\n tags = e.get(\\\"tag\\\")\\n with context(lambda: f\\\"in {loc}:\\\\n {tags}\\\"):\\n e_i = e.copy()\\n name = e_i.pop(\\\"tag\\\")\\n desc = e_i.pop(\\\"desc\\\", \\\"\\\")\\n # ensure that each tag has a non-empty description\\n assert desc != \\\"\\\"\\n rs.add(name)\\n return rs\\n\\n\\n@functools.lru_cache(maxsize=None)\\ndef parse_tags_yaml(path: str) -> set[str]:\\n global _GLOBAL_PARSE_TAGS_YAML_CACHE\\n if path not in _GLOBAL_PARSE_TAGS_YAML_CACHE:\\n with open(path) as f:\\n es = yaml.load(f, Loader=LineLoader)\\n _GLOBAL_PARSE_TAGS_YAML_CACHE[path] = parse_tags_yaml_struct(es, path=path)\\n\\n return _GLOBAL_PARSE_TAGS_YAML_CACHE[path]\\n\\n\\ndef parse_native_yaml(\\n path: str,\\n tags_yaml_path: str,\\n ignore_keys: set[DispatchKey] | None = None,\\n *,\\n skip_native_fns_gen: bool = False,\\n loaded_yaml: object | None = None,\\n) -> ParsedYaml:\\n global _GLOBAL_PARSE_NATIVE_YAML_CACHE\\n if path not in _GLOBAL_PARSE_NATIVE_YAML_CACHE:\\n valid_tags = parse_tags_yaml(tags_yaml_path)\\n\\n # if a loaded yaml is provided, use that instead of reading from path\\n if loaded_yaml is None:\\n with open(path) as f:\\n es = yaml.load(f, Loader=LineLoader)\\n else:\\n es = loaded_yaml\\n\\n _GLOBAL_PARSE_NATIVE_YAML_CACHE[path] = parse_native_yaml_struct(\\n es,\\n valid_tags,\\n ignore_keys,\\n path=path,\\n skip_native_fns_gen=skip_native_fns_gen,\\n )\\n\\n return _GLOBAL_PARSE_NATIVE_YAML_CACHE[path]\\n\\n\\n# Some assertions are already performed during parsing, but those are only within a single NativeFunction.\\n# Assertions here are meant to be performed across NativeFunctions.\\ndef error_check_native_functions(funcs: Sequence[NativeFunction]) -> None:\\n func_map: dict[OperatorName, NativeFunction] = {}\\n base_func_map: dict[BaseOperatorName, list[NativeFunction]] = defaultdict(list)\\n for f in funcs:\\n func_map[f.func.name] = f\\n base_func_map[f.func.name.name].append(f)\\n for f in funcs:\\n if f.structured_delegate is not None:\\n delegate_func = func_map.get(f.structured_delegate)\\n assert delegate_func is not None, (\\n f\\\"{f.func.name} is marked as a structured_delegate pointing to \\\"\\n f\\\"{f.structured_delegate}, but {f.structured_delegate} is missing.\\\"\\n )\\n assert delegate_func.structured, (\\n f\\\"{f.func.name} is marked as a structured_delegate pointing to \\\"\\n f\\\"{f.structured_delegate}, but {f.structured_delegate} is not marked as structured. \\\"\\n f\\\"Consider adding 'structured=True' to the delegated operator\\\"\\n )\\n # See Note [resize_ in Functionalization]\\n # resize_() is technically an inplace view op (and therefore needs the tag),\\n # but it would be overkill to add a true \\\"view\\\" variant of resize.\\n # Instead, resize_() gets special treatment in functionalization,\\n # and we have a resize() op that is non-aliasing + functional.\\n if (\\n \\\"inplace_view\\\" in f.tags\\n and str(f.func.name) != \\\"resize_\\\"\\n and str(f.func.name) != \\\"resize_as_\\\"\\n and str(f.func.name.name) != \\\"set_\\\"\\n ):\\n base_name = f.func.name.name\\n assert base_name.inplace, (\\n f\\\"{f.func.name} is marked with tag: inplace_view, but it doesn't follow the naming \\\"\\n \\\"convention for inplace ops - the codegen expects the base name to have a trailing underscore. \\\"\\n )\\n out_of_place_base_name = BaseOperatorName(\\n base_name.base, False, base_name.dunder_method\\n )\\n assert len(base_func_map[out_of_place_base_name]) > 0, (\\n f\\\"{f.func.name} is marked with tag: inplace_view. The codegen expects there to be a corresponding \\\"\\n f\\\"out-of-place view op with the name '{base_name}' and matching schema, but it didn't find one. \\\"\\n )\\n\\n\\ndef cpp_string(s: str) -> str:\\n \\\"\\\"\\\"Convert a python string into a c++ string literal\\\"\\\"\\\"\\n s = s.replace(\\\"\\\\\\\\\\\", \\\"\\\\\\\\\\\\\\\\\\\")\\n s = s.replace('\\\"', '\\\\\\\\\\\"')\\n s = s.replace(\\\"\\\\a\\\", \\\"\\\\\\\\a\\\")\\n s = s.replace(\\\"\\\\b\\\", \\\"\\\\\\\\b\\\")\\n s = s.replace(\\\"\\\\f\\\", \\\"\\\\\\\\f\\\")\\n s = s.replace(\\\"\\\\n\\\", \\\"\\\\\\\\n\\\")\\n s = s.replace(\\\"\\\\v\\\", \\\"\\\\\\\\v\\\")\\n s = s.replace(\\\"\\\\t\\\", \\\"\\\\\\\\t\\\")\\n return f'\\\"{s}\\\"'\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# C++ CODE GENERATION\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n# Most functions in this section are curried: they consist of a function\\n# that takes some parameters (e.g., what is to be generated) which itself\\n# returns a function that actually maps NativeFunction to the code\\n# to be generated. This pattern makes it convenient to use map, concatMap\\n# and similar functional combinators.\\n\\n\\ndef static_dispatch_keys(backends: list[BackendIndex]) -> list[DispatchKey]:\\n if len(backends) == 0:\\n return []\\n else:\\n return [backend.dispatch_key for backend in backends] + [\\n DispatchKey.CompositeImplicitAutograd,\\n DispatchKey.CompositeImplicitAutogradNestedTensor,\\n DispatchKey.CompositeExplicitAutograd,\\n DispatchKey.CompositeExplicitAutogradNonFunctional,\\n ]\\n\\n\\ndef get_static_dispatch_backend(\\n f: NativeFunction, backend_index: BackendIndex\\n) -> DispatchKey | None:\\n if f.structured_delegate is not None or backend_index.has_kernel(f):\\n # TODO: for ops with structured_delegate it should check the dispatch table of\\n # the out variant instead. For now, these structured ops all have CPU/CUDA kernels\\n # so we always dispatch to the `backend`, but this could be wrong when we\\n # migrate math/default_backend ops to use structured delegate.\\n return backend_index.dispatch_key\\n elif f.has_composite_explicit_autograd_kernel:\\n return DispatchKey.CompositeExplicitAutograd\\n elif f.has_composite_explicit_autograd_non_functional_kernel:\\n return DispatchKey.CompositeExplicitAutogradNonFunctional\\n elif f.has_composite_implicit_autograd_kernel:\\n return DispatchKey.CompositeImplicitAutograd\\n elif f.has_composite_implicit_autograd_nested_tensor_kernel:\\n return DispatchKey.CompositeImplicitAutogradNestedTensor\\n return None\\n\\n\\ndef static_dispatch_ops_header(\\n f: NativeFunction, backend_index: list[BackendIndex]\\n) -> str | None:\\n if backend_index is None or f.manual_kernel_registration:\\n return None\\n\\n output = []\\n for index in backend_index:\\n dispatch_key = get_static_dispatch_backend(f, index)\\n if dispatch_key is not None:\\n output.append(\\n f\\\"#include \\\"\\n )\\n return \\\"\\\\n\\\".join(output)\\n\\n\\ndef static_dispatch_extra_headers(backends: list[BackendIndex]) -> list[str]:\\n return [\\n f\\\"#include \\\"\\n for dispatch_key in static_dispatch_keys(backends)\\n ]\\n\\n\\n# Translates arguments of `sig` to CppSignature bindings.\\n# Note that we have a special case for `memory_format` argument and this case is not covered by\\n# tools.codegen.api.translate() yet as its application is limited to static dispatch.\\ndef translate_args(\\n sig: CppSignature | DispatcherSignature,\\n cpp_sig: CppSignature,\\n) -> str:\\n # Adds SpecialArgName.possibly_redundant_memory_format NamedCType for memory_format bindings\\n def add_spl_memory_format_binding(input_bindings: list[Binding]) -> list[Binding]:\\n output_bindings: list[Binding] = []\\n for binding in input_bindings:\\n if binding.name == \\\"memory_format\\\":\\n spl_mem_format_binding = Binding(\\n nctype=NamedCType(\\n SpecialArgName.possibly_redundant_memory_format,\\n binding.nctype.type,\\n ),\\n name=binding.name,\\n default=binding.default,\\n argument=binding.argument,\\n )\\n output_bindings.append(spl_mem_format_binding)\\n else:\\n output_bindings.append(binding)\\n return output_bindings\\n\\n src_bindings = list(sig.arguments())\\n goal_bindings = list(cpp_sig.arguments())\\n # When last argument of CPP signature has SpecialArgName.possibly_redundant_memory_format NCType,\\n # get memory_format bindings of dispatcher signature to have the same NCType as well\\n for arg in goal_bindings:\\n if arg.nctype.name == SpecialArgName.possibly_redundant_memory_format:\\n src_bindings = add_spl_memory_format_binding(src_bindings)\\n break\\n exprs = translate(src_bindings, goal_bindings)\\n return \\\", \\\".join(a.expr for a in exprs)\\n\\n\\ndef generate_static_dispatch_backend_call(\\n sig: CppSignature | DispatcherSignature,\\n f: NativeFunction,\\n backend_index: BackendIndex,\\n) -> str:\\n cpp_sig = gen_static_dispatch_backend_call_signature(sig, f)\\n name = cpp_sig.name()\\n exprs = translate_args(sig, cpp_sig)\\n backend_metadata = backend_index.get_kernel(f)\\n kernel_ns = (\\n backend_metadata.cpp_namespace\\n if backend_metadata and backend_metadata.cpp_namespace\\n else DEFAULT_KERNEL_NAMESPACE\\n )\\n ns = kernel_ns.replace(\\\"::native\\\", \\\"\\\")\\n return f\\\"return {ns}::{backend_index.dispatch_key.lower()}::{name}({exprs});\\\"\\n\\n\\ndef generate_static_dispatch_fallback_call(\\n sig: CppSignature | DispatcherSignature,\\n f: NativeFunction,\\n backend_indices: list[BackendIndex],\\n) -> str:\\n cpp_sigs = CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=False\\n )\\n if sig.symint and f.func.has_symint():\\n cpp_sig = cpp_sigs.symint_signature\\n else:\\n cpp_sig = cpp_sigs.signature\\n assert cpp_sig is not None\\n name = cpp_sig.name()\\n exprs = translate_args(sig, cpp_sig)\\n ns = DEFAULT_KERNEL_NAMESPACE.replace(\\\"::native\\\", \\\"\\\")\\n if f.has_composite_explicit_autograd_kernel:\\n return f\\\"return {ns}::{DispatchKey.CompositeExplicitAutograd.lower()}::{name}({exprs});\\\"\\n elif f.has_composite_explicit_autograd_non_functional_kernel:\\n return f\\\"return {ns}::{DispatchKey.CompositeExplicitAutogradNonFunctional.lower()}::{name}({exprs});\\\"\\n elif f.has_composite_implicit_autograd_kernel:\\n return f\\\"return {ns}::{DispatchKey.CompositeImplicitAutograd.lower()}::{name}({exprs});\\\"\\n elif f.has_composite_implicit_autograd_nested_tensor_kernel:\\n return f\\\"return {ns}::{DispatchKey.CompositeImplicitAutogradNestedTensor.lower()}::{name}({exprs});\\\"\\n else:\\n return f\\\"\\\"\\\"TORCH_CHECK(false, \\\"Static dispatch does not support {name} for\\\\\\n{', '.join([str(index.dispatch_key)for index in backend_indices])} \\\");\\\"\\\"\\\"\\n\\n\\ndef static_dispatch(\\n sig: CppSignature | DispatcherSignature,\\n f: NativeFunction,\\n backend_indices: list[BackendIndex],\\n) -> str:\\n \\\"\\\"\\\"\\n For a given `NativeFunction`, find out the corresponding backend and dispatch to it. If more than one\\n backends exsit, fallback to static dispatch by determining dispatch key from inputs.\\n Arguments:\\n sig: A CppSignature or DispatcherSignature for this native function we want to use.\\n f: NativeFunction to generate static dispatch.\\n backend_indices: All available backends.\\n Return:\\n C++ code to call backend-specific functions, e.g., \\\"return at::cpu::add(self, other, scale);\\\"\\n \\\"\\\"\\\"\\n if len(backend_indices) == 0 or f.manual_kernel_registration:\\n return \\\"\\\"\\n\\n keys = [\\n b\\n for b in backend_indices\\n if b.has_kernel(f)\\n or (\\n f.structured_delegate is not None\\n and b.dispatch_key in STRUCTURED_DISPATCH_KEYS\\n )\\n ]\\n if len(keys) == 1:\\n return generate_static_dispatch_backend_call(sig, f, keys[0])\\n elif len(keys) == 0:\\n return generate_static_dispatch_fallback_call(sig, f, backend_indices)\\n\\n native_tensor_args = [\\n a.name\\n for a in sig.arguments()\\n if isinstance(a.argument, SelfArgument)\\n or isinstance(a.argument, Argument)\\n and a.argument.type.is_tensor_like()\\n ]\\n tensor_args = \\\", \\\".join(native_tensor_args)\\n tensor_opts = f.func.arguments.tensor_options\\n\\n stmts = []\\n subexprs: list[str] = []\\n if tensor_opts is not None:\\n subexprs.append(\\n \\\"DispatchKeySet(c10::computeDispatchKey(dtype, layout, device))\\\"\\n )\\n if tensor_args != \\\"\\\":\\n subexprs.append(f\\\"c10::detail::multi_dispatch_key_set({tensor_args})\\\")\\n stmts.append(f\\\"\\\"\\\"DispatchKeySet _dk_set = {' | '.join(subexprs)};\\\"\\\"\\\")\\n stmts.append(\\\"DispatchKey _dk = c10::highestPriorityBackendTypeId(_dk_set);\\\")\\n\\n dispatch_code = []\\n for index in keys:\\n dispatch_code.append(f\\\"\\\"\\\"case DispatchKey::{index.dispatch_key}:\\\"\\\"\\\")\\n dispatch_code.append(\\n f\\\"\\\"\\\"\\\\t{generate_static_dispatch_backend_call(sig, f, index)};\\\"\\\"\\\"\\n )\\n\\n fallback = generate_static_dispatch_fallback_call(sig, f, backend_indices)\\n connector = \\\"\\\\n\\\\t\\\\t\\\"\\n\\n return f\\\"\\\"\\\"\\n {connector.join(stmts)}\\n switch (_dk) {{\\n {connector.join(dispatch_code)}\\n default:\\n {fallback}\\n }}\\n \\\"\\\"\\\"\\n\\n\\n# Generates RegisterSchema.cpp. Depending on the selector, either\\n# all schemas are registered, or only some are (in the case of\\n# selective build)\\n@dataclass(frozen=True)\\nclass RegisterSchema:\\n selector: SelectiveBuilder\\n known_tags: dict[str, int] = field(default_factory=dict)\\n\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str | None:\\n if not self.selector.is_native_function_selected(f):\\n return None\\n tags = \\\"{\\\" + \\\", \\\".join(f\\\"at::Tag::{tag}\\\" for tag in sorted(f.tags)) + \\\"}\\\"\\n if tags == \\\"{}\\\":\\n return f\\\"m.def({cpp_string(str(f.func))}, {{}});\\\\n\\\"\\n maybe_tags = \\\"\\\"\\n if tags not in self.known_tags:\\n idx = len(self.known_tags)\\n self.known_tags[tags] = idx\\n maybe_tags = f\\\"const std::vector tags_{idx} = {tags};\\\\n\\\"\\n return f\\\"{maybe_tags}m.def({cpp_string(str(f.func))}, tags_{self.known_tags[tags]});\\\\n\\\"\\n\\n\\n# Generates Operators.h and Operators.cpp.\\n# These provide macros that, given an operator and overload name, allow users\\n# to access an \\\"un-overloaded\\\" function version of the operator. This\\n# is useful for extension writers who want to (1) want to decltype the operator\\n# and (2) don't want to worry about method-only operators.\\n@dataclass(frozen=True)\\nclass ComputeOperators:\\n target: Literal[Target.DECLARATION, Target.DEFINITION]\\n static_dispatch_backend_indices: list[BackendIndex]\\n\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str:\\n sig = DispatcherSignature.from_schema(f.func)\\n name = f.func.name.unambiguous_name()\\n\\n if self.target is Target.DECLARATION:\\n # Note [The ATen Operators API]\\n # The ATen Operators API lives in the at::_ops namespace, and contains compile-time\\n # metadata about each operator + entry points into the Dispatcher.\\n # The C++ function, method, and redispatch API's are all implemented as wrappers\\n # into various bits of the structs defined here.\\n #\\n # Important characteristics about the Operators API:\\n # (1) It follows the Dispatcher API.\\n # This is kind of necessary to avoid overhead.\\n # For example: if it followed the C++ API, then all of the faithful C++ factory functions\\n # would need to wrap their arguments into TensorOptions only to unwrap them again.\\n # (2) Overload names are disambiguated.\\n # This is helpful for pytorch extenders who would like to decltype() an aten operator,\\n # that has overloads, e.g. decltype(at::_ops::mul_Tensor::call)\\n # (3) No argument defaulting is allowed.\\n # This is more of an implementation detail to avoid #include cycles,\\n # since TensorBody.h (which defines the Tensor class) needs to include this file.\\n # (4) manual_cpp_bindings and faithful names are not included in the API.\\n # This applies to stuff like __dispatch__is_complex(), and add_outf().\\n # These aren't \\\"real aten ops\\\", they're just additional functions provided by the C++ API.\\n # They're implemented as wrappers in Functions.h that call into the actual operators\\n # defined here, i.e. at::_ops::is_complex::call() and at::_ops::add_out::call().\\n # This means that ATEN_OP(is_complex) will not fastpath, and will go through the dispatcher.\\n return f\\\"\\\"\\\"\\nstruct TORCH_API {name} {{\\n using schema = {sig.type()};\\n using ptr_schema = schema*;\\n // See Note [static constexpr char* members for windows NVCC]\\n STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(name, \\\"aten::{f.func.name.name}\\\")\\n STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(overload_name, \\\"{f.func.name.overload_name}\\\")\\n STATIC_CONSTEXPR_STR_INL_EXCEPT_WIN_CUDA(schema_str, {cpp_string(str(f.func))})\\n static {sig.defn(name=\\\"call\\\", is_redispatching_fn=False)};\\n static {sig.defn(name=\\\"redispatch\\\", is_redispatching_fn=True)};\\n}};\\\"\\\"\\\"\\n\\n elif self.target is Target.DEFINITION:\\n defns = f\\\"\\\"\\\"\\nSTATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, name, \\\"aten::{f.func.name.name}\\\")\\nSTATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, overload_name, \\\"{f.func.name.overload_name}\\\")\\nSTATIC_CONST_STR_OUT_OF_LINE_FOR_WIN_CUDA({name}, schema_str, {cpp_string(str(f.func))})\\n\\n// aten::{f.func}\\nstatic C10_NOINLINE c10::TypedOperatorHandle<{name}::schema> create_{name}_typed_handle() {{\\n return c10::Dispatcher::singleton()\\n .findSchemaOrThrow({name}::name, {name}::overload_name)\\n .typed<{name}::schema>();\\n}}\\n\\\"\\\"\\\"\\n for is_redispatching_fn in [False, True]:\\n if is_redispatching_fn:\\n dispatcher_exprs_str = \\\", \\\".join(\\n [\\\"dispatchKeySet\\\"] + [a.name for a in sig.arguments()]\\n )\\n method_base = \\\"redispatch\\\"\\n else:\\n dispatcher_exprs_str = \\\", \\\".join([a.name for a in sig.arguments()])\\n method_base = \\\"call\\\"\\n\\n dispatcher_call = method_base\\n method_name = f\\\"{name}::{method_base}\\\"\\n\\n fn_body = f\\\"\\\"\\\"\\n static auto op = create_{name}_typed_handle();\\n return op.{dispatcher_call}({dispatcher_exprs_str});\\\"\\\"\\\"\\n\\n if (\\n not is_redispatching_fn\\n and len(self.static_dispatch_backend_indices) > 0\\n ):\\n # call() should go through static dispatch\\n fn_body = static_dispatch(\\n sig, f, backend_indices=self.static_dispatch_backend_indices\\n )\\n defns += f\\\"\\\"\\\"\\n// aten::{f.func}\\n{sig.defn(name=method_name, is_redispatching_fn=is_redispatching_fn)} {{\\n {fn_body}\\n}}\\n\\\"\\\"\\\"\\n return defns\\n else:\\n assert_never(self.target)\\n\\n\\n# Generates Functions.h, which provides the functional public C++ API,\\n# and the scaffolding to call into the dispatcher from these functions.\\n@dataclass(frozen=True)\\nclass ComputeFunction:\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str | None:\\n sig_group = CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=f.manual_cpp_binding\\n )\\n has_symint = f.func.has_symint()\\n\\n result = \\\"\\\"\\n for sig in sig_group.signatures():\\n # See Note [The ATen Operators API]\\n target_sig = DispatcherSignature.from_schema(f.func)\\n exprs = translate(sig.arguments(), target_sig.arguments())\\n exprs_str = \\\", \\\".join([e.expr for e in exprs])\\n\\n if sig.symint:\\n intlike_t = \\\"c10::SymInt\\\"\\n else:\\n intlike_t = \\\"int64_t\\\"\\n\\n if Variant.function in f.variants:\\n result += f\\\"\\\"\\\"\\n// aten::{f.func}\\ninline {sig.decl()} {{\\n return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str});\\n}}\\\"\\\"\\\"\\n\\n # The template function can be used from template situations\\n # where you want to switch between the symint or not version\\n # depending on a template argument\\n #\\n # NB: we ALWAYS generate this even for methods. But we put it in\\n # this header so it can take advantage of per-op headers\\n if has_symint:\\n result += f\\\"\\\"\\\"\\nnamespace symint {{\\n template ::value>>\\n {sig.decl(suppress_symint_suffix=True)} {{\\n return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str});\\n }}\\n}}\\n\\\"\\\"\\\"\\n return result\\n\\n\\n# Generates TensorBody.h. This file provides the object-oriented (method-based)\\n# public C++ API, and the scaffolding to call into the dispatcher from these functions.\\n@dataclass(frozen=True)\\nclass ComputeTensorMethod:\\n target: Literal[Target.DECLARATION, Target.DEFINITION]\\n static_dispatch_backend_indices: list[BackendIndex]\\n\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str | None:\\n if Variant.method not in f.variants:\\n return None\\n\\n assert not f.func.is_out_fn()\\n assert f.func.arguments.self_arg is not None\\n\\n sig_group = CppSignatureGroup.from_native_function(\\n f, method=True, fallback_binding=f.manual_cpp_binding\\n )\\n\\n if self.target is Target.DECLARATION:\\n result = \\\"\\\"\\n for sig in sig_group.signatures():\\n result += f\\\"{sig.decl()} const;\\\\n\\\"\\n return result\\n\\n if self.target is not Target.DEFINITION:\\n assert_never(self.target)\\n\\n result = \\\"\\\"\\n\\n for sig in sig_group.signatures():\\n target_sig = DispatcherSignature.from_schema(f.func)\\n exprs = translate(sig.arguments(), target_sig.arguments(), method=True)\\n exprs_str = \\\", \\\".join([e.expr for e in exprs])\\n\\n result += f\\\"\\\"\\\"\\n// aten::{f.func}\\ninline {sig.defn(prefix=\\\"Tensor::\\\")} const {{\\n return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str});\\n}}\\n\\\"\\\"\\\"\\n\\n return result\\n\\n\\n# Generates RedispatchFunctions.h.\\n# This is similar to the C++ API defined in Functions.h, but provides access\\n# to the dispatcher's redispatch API.\\n@dataclass(frozen=True)\\nclass ComputeRedispatchFunction:\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str | None:\\n # We unconditionally generate function variants of the redispatch API.\\n # This is mainly because we can namespace functions separately, but not methods,\\n sig_group = CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=f.manual_cpp_binding\\n )\\n\\n result = \\\"\\\"\\n for sig in sig_group.signatures():\\n target_sig = DispatcherSignature.from_schema(f.func)\\n exprs = translate(sig.arguments(), target_sig.arguments())\\n exprs_str = \\\", \\\".join([\\\"dispatchKeySet\\\"] + [a.expr for a in exprs])\\n\\n result += f\\\"\\\"\\\"\\n// aten::{f.func}\\ninline {sig.decl(is_redispatching_fn=True)} {{\\n return at::_ops::{f.func.name.unambiguous_name()}::redispatch({exprs_str});\\n}}\\n\\\"\\\"\\\"\\n\\n return result\\n\\n\\n# Generates ATenOpList.cpp, a runtime accessible list of all aten\\n# operators.\\n# TODO: This was historically used to help some JIT interop code\\n# figure out whether or not to treat aten namespace'd operators\\n# one way or another, we should reevaluate if this is actually needed.\\n@with_native_function\\ndef compute_aten_op(f: NativeFunction) -> str:\\n return f'{{\\\"aten::{f.func.name.name}\\\", \\\"{f.func.name.overload_name}\\\"}},'\\n\\n\\n# Generates MetaFunctions.h\\ndef compute_meta_function_declaration(g: NativeFunctionsGroup) -> str | None:\\n if not g.structured:\\n return None\\n with native_function_manager(g.out):\\n name = meta.name(g)\\n args = structured.meta_arguments(g)\\n args_str = \\\", \\\".join(a.decl() for a in args)\\n parent_class = g.out.structured_inherits\\n if parent_class is None:\\n parent_class = \\\"at::impl::MetaBase\\\"\\n meta_return = \\\"void\\\"\\n precomputed = g.out.precomputed if g.structured else None\\n\\n if precomputed:\\n # Generate the template declaration with one bool parameter for each\\n # precomputed element. Each parameter is true if the corresponding (in\\n # terms of position) precomputed element has been set.\\n precomputed_values = [*precomputed.replace.values(), precomputed.add]\\n precomputed_elements = [\\n elem for replace_list in precomputed_values for elem in replace_list\\n ]\\n precomputed_template_parameters = [\\n elem.name.upper() for elem in precomputed_elements\\n ]\\n precomputed_template_params_str = \\\", \\\".join(\\n f\\\"bool {param} = false\\\" for param in precomputed_template_parameters\\n )\\n precompute_template_decl = f\\\"template <{precomputed_template_params_str}>\\\"\\n\\n # Generate a string containing declarations of all precomputed elements.\\n precomputed_elements_with_cpp_types = [\\n structured.argument_type(elem, binds=elem.name)\\n for elem in precomputed_elements\\n ]\\n\\n precomputed_elements_decl = \\\";\\\\n\\\".join(\\n f\\\"{elem.cpp_type(strip_ref=True)} {elem.name}\\\"\\n for elem in precomputed_elements_with_cpp_types\\n )\\n\\n # Generate \\\"setter\\\" methods for each precomputed element. Each method will return\\n # a new instance of precompute_out with the template parameter that corresponds to\\n # the member set by the method to true (to indicate that it has been set).\\n setter_methods = []\\n for i, elem in enumerate(precomputed_elements):\\n # Generate the signature. The return type will be the same\\n # as the type of `this` but with the template parameter\\n # corresponding to the element set by this method set to true.\\n # The assert generated below will ensure that this template\\n # parameter is false on the type of `this`.\\n return_ty_templates = \\\", \\\".join(\\n precomputed_template_parameters[:i]\\n + [\\\"true\\\"]\\n + precomputed_template_parameters[i + 1 :]\\n )\\n return_ty = f\\\"precompute_out<{return_ty_templates}>\\\"\\n elem_cpp_ty = precomputed_elements_with_cpp_types[i].cpp_type(\\n strip_ref=True\\n )\\n signature = f\\\"{return_ty} set_{elem.name}({elem_cpp_ty} value)\\\"\\n\\n # Generate an assert which checks that the\\n # template parameter corresponding to the precomputed\\n # element that is set by this method is false on the\\n # class corresponding to the object that `this` points to.\\n # This ensures that each element can be set only once.\\n assert_msg = f'\\\"{elem.name} already set\\\"'\\n assert_stmt = f\\\"static_assert({precomputed_template_parameters[i]} == false, {assert_msg});\\\"\\n\\n # Generate the new object construction block. All state\\n # except the element that this method sets is copied from the\\n # object that `this` points to. The value for the element that\\n # the method sets is taken from a method parameter.\\n construction_stmts = []\\n construction_stmts.append(f\\\"{return_ty} ret;\\\")\\n\\n for j, elem in enumerate(precomputed_elements):\\n if i == j:\\n construction_stmts.append(f\\\"ret.{elem.name} = value;\\\")\\n else:\\n construction_stmts.append(\\n f\\\"ret.{elem.name} = this->{elem.name};\\\"\\n )\\n\\n construction_stmts.append(\\\"return ret;\\\")\\n construction_block = \\\"\\\\n\\\".join(construction_stmts)\\n\\n setter_methods.append(\\n f\\\"\\\"\\\"\\n {signature} {{\\n {assert_stmt}\\n {construction_block}\\n }}\\n \\\"\\\"\\\"\\n )\\n setter_methods_decl = \\\"\\\\n\\\".join(setter_methods)\\n\\n # Meta should return an instance of the struct containing the precomputed elements.\\n meta_return_template_params = \\\", \\\".join(\\n [\\\"true\\\"] * len(precomputed_template_parameters)\\n )\\n # This typedef (actually a using statement) is needed so that TORCH_META_FUNC can reuse the return\\n # type (which has a variable number of template parameters).\\n meta_return_typedef = f\\\"using meta_return_ty = precompute_out <{meta_return_template_params}>;\\\"\\n meta_return = \\\"meta_return_ty\\\"\\n precomputed_decl = f\\\"\\\"\\\"\\n {precompute_template_decl}\\n struct TORCH_API precompute_out {{\\n {setter_methods_decl}\\n {precomputed_elements_decl};\\n }};\\\"\\\"\\\"\\n else:\\n meta_return_typedef = \\\"\\\"\\n precomputed_decl = \\\"\\\"\\n\\n return f\\\"\\\"\\\"\\\\\\nstruct TORCH_API structured_{name} : public {parent_class} {{\\n {precomputed_decl}\\n {meta_return_typedef}\\n {meta_return} meta({args_str});\\n}};\\n\\\"\\\"\\\"\\n\\n\\ndef needs_backend_select(f: NativeFunction, selector: SelectiveBuilder) -> bool:\\n name = str(f.func.name.name)\\n if name.endswith(\\\"_like\\\") or name.startswith(\\\"new_\\\"):\\n return False\\n if f.func.arguments.tensor_options is None:\\n return False\\n return selector.is_native_function_selected(f)\\n\\n\\n# Generates RegisterBackendSelect.cpp, a series of kernels which provide\\n# specialized computation of dispatch key for operator signatures which cannot\\n# be easily done automatically using templating.\\n@dataclass(frozen=True)\\nclass ComputeBackendSelect:\\n target: Literal[Target.DEFINITION, Target.REGISTRATION]\\n\\n # Selector object to determine which operators to generate\\n # registration code for.\\n selector: SelectiveBuilder\\n\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str | None:\\n if not needs_backend_select(f, self.selector):\\n return None\\n\\n name = native.name(f.func)\\n # BackendSelect can go to Meta, so it must preserve symints\\n native_sig = NativeSignature(f.func, symint=True)\\n\\n native_tensor_args = [\\n a\\n for a in native_sig.arguments()\\n if isinstance(a.argument, Argument) and a.argument.type.is_tensor_like()\\n ]\\n\\n dispatcher_sig = DispatcherSignature.from_schema(f.func)\\n\\n sig: NativeSignature | DispatcherSignature\\n sig = dispatcher_sig\\n dispatcher_exprs = dispatcher_sig.exprs()\\n dispatch_key = \\\"c10::computeDispatchKey(dtype, layout, device)\\\"\\n\\n if self.target is Target.DEFINITION:\\n # I don't think there's actually a good reason to generate\\n # these two cases differently\\n # The first case could probably be improved though- it calls computeDispatchKeySet(),\\n # which looks at TLS dispatch keys- there should not be any by the time we reach backend select.\\n if native_tensor_args:\\n assert f.func.arguments.has_tensor_arg()\\n tensor_args = \\\", \\\".join(a.name for a in native_tensor_args)\\n compute_dk = f\\\"\\\"\\\"\\\\\\nDispatchKeySet _dk_set = c10::DispatchKeySet({dispatch_key}) | c10::detail::multi_dispatch_key_set({tensor_args});\\nDispatchKeySet _dk_mask = c10::DispatchKeySet(DispatchKeySet::FULL_AFTER, DispatchKey::BackendSelect);\\nDispatchKeySet _dk = c10::impl::computeDispatchKeySet(_dk_set, _dk_mask);\\\"\\\"\\\"\\n else:\\n assert not f.func.arguments.has_tensor_arg()\\n compute_dk = (\\n f\\\"DispatchKeySet _dk = c10::DispatchKeySet({dispatch_key});\\\"\\n )\\n return f\\\"\\\"\\\"\\\\\\n// aten::{f.func}\\nC10_ALWAYS_INLINE\\n{sig.defn(name)} {{\\n {compute_dk}\\n return at::_ops::{f.func.name.unambiguous_name()}::redispatch(\\n _dk, {', '.join(a.expr for a in dispatcher_exprs)});\\n}}\\n\\\"\\\"\\\"\\n elif self.target is Target.REGISTRATION:\\n return f\\\"\\\"\\\"m.impl(\\\"aten::{f.func.name}\\\", TORCH_FN({name}));\\\"\\\"\\\"\\n else:\\n assert_never(self.target)\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# YAML CODE GENERATION\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef format_yaml(data: object) -> str:\\n # Ignore alias in Dumper\\n YamlDumper.ignore_aliases = lambda self, data: True # type: ignore[assignment]\\n\\n # Support serializing OrderedDict\\n def dict_representer(dumper: Any, data: Any) -> Any:\\n return dumper.represent_dict(data.items())\\n\\n YamlDumper.add_representer(OrderedDict, dict_representer) # type: ignore[no-untyped-call]\\n # Some yaml parsers (e.g. Haskell's) don't understand line breaks.\\n # width=1e9 turns off optional line breaks and improves\\n # the portability of the outputted yaml.\\n return yaml.dump(data, default_flow_style=False, Dumper=YamlDumper, width=1e9) # type: ignore[no-any-return, call-overload]\\n\\n\\n# For some reason, some defaults we write to YAML are written as native\\n# YAML objects, rather than doing them uniformly as strings. This\\n# function detects those cases and converts them into native Python\\n# objects.\\ndef pythonify_default(s: str) -> object:\\n if s == \\\"true\\\":\\n return True\\n elif s == \\\"false\\\":\\n return False\\n\\n try:\\n return int(s)\\n except ValueError:\\n try:\\n return float(s)\\n except ValueError:\\n return s\\n\\n\\n# What is a dynamic type? Over time, the semantic meaning of\\n# dynamic type has degraded to meaninglessness (in the old days,\\n# it captured dtype-ness of types, but that has gone away with\\n# the removal of TH). These days, it's mostly the same thing as\\n# the C++ API argument type, except that Tensor and Tensor?\\n# arguments simply present as Tensor.\\n#\\n# TODO: Get rid of dynamic_type, after getting tools/autograd\\n# to use the new codegen framework\\ndef dynamic_type(t: Type) -> str:\\n if isinstance(t, OptionalType):\\n return dynamic_type(t.elem)\\n # Note we don't use t.is_tensor_like() here because it would\\n # also include Tensor[]\\n if str(t) == \\\"Tensor\\\":\\n return \\\"at::Tensor\\\"\\n # This is a legacy concept, so never report SymInt\\n return cpp.argumenttype_type(\\n t, mutable=False, binds=\\\"__placeholder__\\\", symint=False\\n ).cpp_type()\\n\\n\\ndef compute_method_of_yaml(variants: set[Variant]) -> list[str]:\\n # This is written out explicitly to ensure that Tensor and\\n # namespace are put into the list in the right order\\n method_of = [\\\"Type\\\"]\\n if Variant.method in variants:\\n method_of.append(\\\"Tensor\\\")\\n if Variant.function in variants:\\n method_of.append(\\\"namespace\\\")\\n return method_of\\n\\n\\ndef compute_returns_yaml(\\n f: NativeFunction,\\n) -> tuple[list[dict[str, str]], dict[str, str]]:\\n # Note [name and field_name]\\n # ~~~~~~~~~~~~~~~~~~~~~~~~~~\\n # To understand name_to_field_name, we must first talk about this\\n # schema:\\n #\\n # lstsq.X(Tensor self, Tensor A, *, Tensor(a!) X, Tensor(b!) qr) -> (Tensor(a!) solution, Tensor(b!) QR)\\n #\\n # There is something very odd about this schema: it is an out\\n # variant of the function (that is to say, it will convert into\\n # at::lstsq_out() in the C++ API), but the names of the output\\n # return arguments don't match the keyword argument names of\\n # the inputs. It TURNS OUT that in this situation, the historical\\n # Declarations.yaml we want to output is this (abbreviated to\\n # only show relevant fields):\\n #\\n # arguments:\\n # ...\\n # - field_name: solution\\n # name: X\\n # - field_name: QR\\n # name: qr\\n # ...\\n #\\n # returns:\\n # - field_name: solution\\n # name: X\\n # - field_name: QR\\n # name: qr\\n #\\n # The name of the return fields is stored in 'field_name', and the\\n # name of the arguments is stored in 'name'. So when we process\\n # arguments, we need a way to get at the corresponding return. At\\n # the moment, this is most conveniently done by constructing a\\n # mapping from name (the argument concept) to field_name (the\\n # return concept) while processing return arguments, since we don't\\n # directly maintain this correspondence in the modeling of function\\n # schema itself.\\n #\\n # See also https://github.com/pytorch/pytorch/issues/43114\\n name_to_field_name: dict[str, str] = {}\\n\\n # Compute the returns field of the YAML entry\\n names = cpp.return_names(f)\\n returns = []\\n for i, (r, name) in enumerate(zip(f.func.returns, names)):\\n ret = {\\n \\\"dynamic_type\\\": dynamic_type(r.type),\\n \\\"name\\\": name,\\n # legacy, report ints\\n \\\"type\\\": cpp.return_type(r, symint=False).cpp_type(),\\n }\\n\\n if r.name:\\n # See Note [name and field_name]\\n ret[\\\"field_name\\\"] = r.name\\n if f.func.is_out_fn():\\n name_to_field_name[f.func.arguments.out[i].name] = r.name\\n\\n returns.append(ret)\\n\\n return returns, name_to_field_name\\n\\n\\n# arguments in yaml roughly corresponds to the public C++ API\\ndef compute_cpp_argument_yaml(\\n cpp_a: Binding,\\n *,\\n schema_order: bool,\\n kwarg_only_set: set[str],\\n out_arg_set: set[str],\\n name_to_field_name: dict[str, str],\\n) -> object:\\n if isinstance(cpp_a.argument, TensorOptionsArguments):\\n arg: dict[str, object] = {\\n \\\"annotation\\\": None,\\n \\\"dynamic_type\\\": \\\"at::TensorOptions\\\",\\n \\\"is_nullable\\\": False,\\n \\\"name\\\": cpp_a.name,\\n \\\"type\\\": cpp_a.type,\\n \\\"kwarg_only\\\": True,\\n }\\n if cpp_a.default is not None:\\n arg[\\\"default\\\"] = cpp_a.default\\n return arg\\n elif isinstance(cpp_a.argument, SelfArgument):\\n raise AssertionError\\n elif isinstance(cpp_a.argument, Argument):\\n return compute_argument_yaml(\\n cpp_a.argument,\\n schema_order=schema_order,\\n kwarg_only_set=kwarg_only_set,\\n out_arg_set=out_arg_set,\\n name_to_field_name=name_to_field_name,\\n )\\n\\n\\ndef compute_argument_yaml(\\n a: Argument,\\n *,\\n schema_order: bool,\\n kwarg_only_set: set[str],\\n out_arg_set: set[str],\\n name_to_field_name: dict[str, str],\\n) -> object:\\n arg: dict[str, object] = {\\n \\\"annotation\\\": str(a.annotation) if a.annotation else None,\\n \\\"dynamic_type\\\": dynamic_type(a.type),\\n \\\"is_nullable\\\": a.type.is_nullable(),\\n \\\"name\\\": a.name,\\n # legacy, report ints\\n \\\"type\\\": cpp.argument_type(a, binds=\\\"__placeholder__\\\", symint=False).cpp_type(),\\n }\\n if a.default is not None:\\n arg[\\\"default\\\"] = pythonify_default(\\n cpp.default_expr(a.default, a.type, symint=False)\\n )\\n if a.name in kwarg_only_set:\\n arg[\\\"kwarg_only\\\"] = True\\n if a.name in out_arg_set:\\n arg[\\\"output\\\"] = True\\n arg[\\\"allocate\\\"] = True\\n # See Note [name and field_name]\\n if a.name in name_to_field_name:\\n arg[\\\"field_name\\\"] = name_to_field_name[a.name]\\n # Historically, booleans don't get their size recorded, because it\\n # is already built into the cpp type (e.g., std::array)\\n l = a.type.is_list_like()\\n if l is not None and l.size is not None and str(l.elem) != \\\"bool\\\":\\n arg[\\\"size\\\"] = l.size\\n return arg\\n\\n\\n@with_native_function\\ndef compute_declaration_yaml(f: NativeFunction) -> object:\\n returns, name_to_field_name = compute_returns_yaml(f)\\n\\n # These sets are used to conveniently test if an argument is a\\n # kwarg-only or out argument\\n kwarg_only_set = {a.name for a in f.func.arguments.flat_kwarg_only}\\n out_arg_set = {a.name for a in f.func.arguments.out}\\n\\n sig_group = CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=False\\n )\\n cpp_args = sig_group.signature.arguments()\\n arguments = [\\n compute_cpp_argument_yaml(\\n cpp_a,\\n schema_order=False,\\n kwarg_only_set=kwarg_only_set,\\n out_arg_set=out_arg_set,\\n name_to_field_name=name_to_field_name,\\n )\\n for cpp_a in cpp_args\\n ]\\n\\n schema_order_jit_arguments = list(f.func.schema_order_arguments())\\n\\n schema_order_arguments = [\\n compute_argument_yaml(\\n a,\\n schema_order=True,\\n kwarg_only_set=kwarg_only_set,\\n out_arg_set=out_arg_set,\\n name_to_field_name=name_to_field_name,\\n )\\n for a in schema_order_jit_arguments\\n ]\\n\\n cpp_schema_order_types = [\\n # NB: method here doesn't matter\\n r.type\\n for a in schema_order_jit_arguments\\n for r in cpp.argument(\\n a,\\n method=False,\\n cpp_no_default_args=set(),\\n faithful=False,\\n symint=False,\\n has_tensor_options=False,\\n )\\n ]\\n\\n # legacy, report ints\\n cpp_returns = cpp.returns_type(f.func.returns, symint=False).cpp_type()\\n schema_order_cpp_signature = f\\\"{cpp_returns} ({', '.join(cpp_schema_order_types)})\\\"\\n\\n is_factory_method = (\\n any(isinstance(a.argument, TensorOptionsArguments) for a in cpp_args)\\n and Variant.method not in f.variants\\n )\\n\\n return OrderedDict(\\n [\\n (\\\"name\\\", cpp.name(f.func)),\\n (\\\"operator_name\\\", str(f.func.name.name)),\\n (\\\"overload_name\\\", str(f.func.name.overload_name)),\\n (\\\"manual_kernel_registration\\\", f.manual_kernel_registration),\\n (\\n \\\"category_override\\\",\\n f.category_override if f.category_override is not None else \\\"\\\",\\n ),\\n (\\\"schema_string\\\", f\\\"aten::{f.func}\\\"),\\n (\\\"arguments\\\", arguments),\\n (\\\"schema_order_cpp_signature\\\", schema_order_cpp_signature),\\n (\\\"schema_order_arguments\\\", schema_order_arguments),\\n (\\\"method_of\\\", compute_method_of_yaml(f.variants)),\\n (\\\"mode\\\", \\\"native\\\"),\\n (\\\"python_module\\\", \\\"\\\" if f.python_module is None else f.python_module),\\n (\\\"returns\\\", returns),\\n (\\\"inplace\\\", f.func.name.name.inplace),\\n (\\\"is_factory_method\\\", is_factory_method),\\n (\\\"abstract\\\", f.is_abstract),\\n (\\\"device_guard\\\", f.device_guard),\\n (\\\"with_gil\\\", False),\\n (\\\"deprecated\\\", False),\\n (\\\"has_math_kernel\\\", f.has_composite_implicit_autograd_kernel),\\n ]\\n )\\n\\n\\n# See Note [Auto generated composite kernels]\\ndef has_autogenerated_composite_kernel(f: NativeFunction) -> bool:\\n return (f.structured or f.structured_delegate is not None) and (\\n f.func.kind() == SchemaKind.functional or f.func.kind() == SchemaKind.inplace\\n )\\n\\n\\n@with_native_function_and_indices\\ndef compute_registration_declarations(\\n f: NativeFunction, backend_indices: dict[DispatchKey, BackendIndex]\\n) -> str:\\n name = dispatcher.name(f.func)\\n returns_type = dispatcher.returns_type(\\n f.func.returns\\n ).cpp_type_registration_declarations()\\n args = dispatcher.arguments(f.func)\\n args_str = \\\", \\\".join(a.no_default().decl_registration_declarations() for a in args)\\n comment_data: dict[str, str] = {\\n \\\"schema\\\": f\\\"aten::{f.func}\\\",\\n # TODO: What exactly is the semantics of the 'dispatch' field?\\n \\\"dispatch\\\": str(\\n {k for k, v in backend_indices.items() if v.has_kernel(f)}\\n != {DispatchKey.CompositeImplicitAutograd}\\n and {k for k, v in backend_indices.items() if v.has_kernel(f)}\\n != {\\n DispatchKey.CompositeImplicitAutograd,\\n DispatchKey.CompositeImplicitAutogradNestedTensor,\\n }\\n ),\\n \\\"default\\\": str(f.has_composite_kernel or has_autogenerated_composite_kernel(f)),\\n }\\n return f\\\"\\\"\\\"{returns_type} {name}({args_str}); // {json.dumps(comment_data)}\\n\\\"\\\"\\\"\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# RUN IT ALL\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef get_custom_build_selector(\\n provided_op_registration_allowlist: list[str] | None,\\n op_selection_yaml_path: str | None,\\n) -> SelectiveBuilder:\\n assert not (\\n provided_op_registration_allowlist is not None\\n and op_selection_yaml_path is not None\\n ), (\\n \\\"Both provided_op_registration_allowlist and \\\"\\n + \\\"op_selection_yaml_path can NOT be provided at the \\\"\\n + \\\"same time.\\\"\\n )\\n\\n op_registration_allowlist: set[str] | None = None\\n if provided_op_registration_allowlist is not None:\\n op_registration_allowlist = set(provided_op_registration_allowlist)\\n\\n if op_registration_allowlist is not None:\\n selector = SelectiveBuilder.from_legacy_op_registration_allow_list(\\n op_registration_allowlist,\\n True,\\n False,\\n )\\n elif op_selection_yaml_path is not None:\\n selector = SelectiveBuilder.from_yaml_path(op_selection_yaml_path)\\n else:\\n selector = SelectiveBuilder.get_nop_selector()\\n\\n return selector\\n\\n\\ndef get_grouped_by_view_native_functions(\\n native_functions: Sequence[NativeFunction],\\n) -> Sequence[NativeFunction | NativeFunctionsViewGroup]:\\n def maybe_create_view_group(\\n d: dict[ViewSchemaKind | SchemaKind, NativeFunction]\\n ) -> list[NativeFunction | NativeFunctionsViewGroup]:\\n funcs: list[NativeFunction | NativeFunctionsViewGroup] = []\\n if ViewSchemaKind.aliasing in d:\\n view = d.pop(ViewSchemaKind.aliasing)\\n view_inplace = d.pop(ViewSchemaKind.aliasing_inplace, None)\\n view_copy = d.pop(SchemaKind.functional, None)\\n\\n funcs.append(\\n NativeFunctionsViewGroup(\\n view=view,\\n view_copy=view_copy,\\n view_inplace=view_inplace,\\n )\\n )\\n # Take the remaining functions that weren't part of the view group\\n # and emit them separately\\n funcs.extend(d.values())\\n return funcs\\n\\n grouped_by_views: dict[\\n FunctionSchema, dict[SchemaKind | ViewSchemaKind, NativeFunction]\\n ] = defaultdict(dict)\\n for f in native_functions:\\n schema = f.func.view_signature()\\n view_kind: ViewSchemaKind = f.view_schema_kind\\n # We need to group up ops relevant to the same \\\"view\\\", consisting of:\\n # view op (ViewSchemaKind.aliasing)\\n # view_inplace op (ViewSchemaKind.aliasing_inplace)\\n # view_copy op (SchemaKind.functional)\\n if view_kind == ViewSchemaKind.non_aliasing:\\n kind = f.func.kind()\\n assert kind not in grouped_by_views[schema]\\n grouped_by_views[schema][kind] = f\\n else:\\n assert (\\n view_kind not in grouped_by_views[schema]\\n ), f\\\"{view_kind} already in {grouped_by_views[schema].keys()}\\\"\\n grouped_by_views[schema][view_kind] = f\\n\\n return list(concatMap(maybe_create_view_group, grouped_by_views.values()))\\n\\n\\ndef get_grouped_native_functions(\\n native_functions: Sequence[NativeFunction],\\n) -> Sequence[NativeFunction | NativeFunctionsGroup]:\\n def flatten_pre_group(\\n d: dict[SchemaKind, NativeFunction]\\n ) -> Sequence[NativeFunction | NativeFunctionsGroup]:\\n r = NativeFunctionsGroup.from_dict(d)\\n if r is None:\\n # Invariant: any NativeFunctions that are code-generated\\n # should have been grouped into NativeFunctionsGroup objects\\n assert not any(\\\"generated\\\" in f.tags for f in d.values())\\n return list(d.values())\\n else:\\n return [r]\\n\\n # TODO: how come ValuesView isn't a Sequence lol\\n pre_grouped_native_functions = pre_group_native_functions(native_functions)\\n return list(\\n concatMap(flatten_pre_group, list(pre_grouped_native_functions.values()))\\n )\\n\\n\\ndef get_ns_grouped_kernels(\\n *,\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n backend_indices: dict[DispatchKey, BackendIndex],\\n native_function_decl_gen: Callable[\\n [NativeFunctionsGroup | NativeFunction, BackendIndex], list[str]\\n ] = dest.compute_native_function_declaration,\\n) -> dict[str, list[str]]:\\n ns_grouped_kernels: dict[str, list[str]] = defaultdict(list)\\n for f in grouped_native_functions:\\n native_function_namespaces = set()\\n dispatch_keys = set()\\n for dispatch_key, backend_idx in backend_indices.items():\\n backend_metadata = backend_idx.get_kernel(f)\\n if backend_metadata:\\n namespace = backend_metadata.cpp_namespace\\n dispatch_keys.add(dispatch_key)\\n native_function_namespaces.add(namespace)\\n else:\\n namespace = DEFAULT_KERNEL_NAMESPACE\\n assert (\\n len(native_function_namespaces) <= 1\\n ), f\\\"Codegen only supports one namespace per operator, got {native_function_namespaces} from {dispatch_keys}\\\"\\n ns_grouped_kernels[namespace].extend(\\n native_function_decl_gen(f, backend_idx)\\n )\\n return ns_grouped_kernels\\n\\n\\ndef get_native_function_declarations_from_ns_grouped_kernels(\\n *,\\n ns_grouped_kernels: dict[str, list[str]],\\n) -> list[str]:\\n declarations: list[str] = []\\n newline = \\\"\\\\n\\\"\\n for namespace, kernels in ns_grouped_kernels.items():\\n ns_helper = NamespaceHelper(\\n namespace_str=namespace,\\n entity_name=\\\"\\\",\\n max_level=4,\\n )\\n # Convert to a set first to remove duplicate kernel names. Backends are\\n # allowed to repeat kernel names; only generate the declaration once!\\n ordered_kernels = list(OrderedDict.fromkeys(kernels))\\n declarations.extend(\\n f\\\"\\\"\\\"\\n{ns_helper.prologue}\\n{newline.join(ordered_kernels)}\\n{ns_helper.epilogue}\\n \\\"\\\"\\\".split(\\n newline\\n )\\n )\\n return declarations\\n\\n\\n# Return native function declarations grouped by their namespaces.\\ndef get_native_function_declarations(\\n *,\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n backend_indices: dict[DispatchKey, BackendIndex],\\n native_function_decl_gen: Callable[\\n [NativeFunctionsGroup | NativeFunction, BackendIndex], list[str]\\n ] = dest.compute_native_function_declaration,\\n) -> list[str]:\\n \\\"\\\"\\\"\\n Generate kernel declarations, in `NativeFunction(s).h`.\\n :param grouped_native_functions: a sequence of `NativeFunction` or `NativeFunctionGroup`.\\n :param backend_indices: kernel collections grouped by dispatch key.\\n :param native_function_decl_gen: callable to generate kernel declaration for each `NativeFunction`.\\n :return: a list of string, from the string with all declarations, grouped by namespaces, split by newline.\\n \\\"\\\"\\\"\\n\\n ns_grouped_kernels = get_ns_grouped_kernels(\\n grouped_native_functions=grouped_native_functions,\\n backend_indices=backend_indices,\\n native_function_decl_gen=native_function_decl_gen,\\n )\\n return get_native_function_declarations_from_ns_grouped_kernels(\\n ns_grouped_kernels=ns_grouped_kernels\\n )\\n\\n\\ndef get_kernel_namespace(\\n *, f: NativeFunction | NativeFunctionsGroup, backend_idx: BackendIndex\\n) -> str:\\n backend_metadata = backend_idx.get_kernel(f)\\n assert not backend_metadata or \\\"::native\\\" in backend_metadata.cpp_namespace, (\\n f\\\"The kernel for function {f.func.name if isinstance(f, NativeFunction) else f.functional.func.name} \\\"\\n f\\\"with dispatch key {backend_idx.dispatch_key}\\\"\\n f\\\" has a namespace {backend_metadata.cpp_namespace} and it's not ending with '::native'.\\\"\\n )\\n return (\\n backend_metadata.cpp_namespace if backend_metadata else DEFAULT_KERNEL_NAMESPACE\\n )\\n\\n\\n# Return native function definitions grouped by dispatch key and custom namespace.\\n# Used in RegisterDispatchKey.cpp and etc.\\ndef get_native_function_definitions(\\n *,\\n fm: FileManager,\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n dispatch_key: DispatchKey,\\n backend_idx: BackendIndex,\\n selector: SelectiveBuilder,\\n rocm: bool,\\n symint: bool,\\n skip_dispatcher_op_registration: bool,\\n gen_dispatch_helpers: bool,\\n) -> list[str]:\\n definitions: list[str] = []\\n ns_definitions: dict[str, list[str]] = defaultdict(list)\\n anonymous_definitions: dict[str, list[str]] = defaultdict(list)\\n registrations: dict[str, dict[str, list[str]]] = defaultdict(dict)\\n newline = \\\"\\\\n\\\"\\n ns_gen = dest.RegisterDispatchKey(\\n backend_idx,\\n Target.NAMESPACED_DEFINITION,\\n selector,\\n rocm=rocm,\\n symint=symint,\\n class_method_name=None,\\n skip_dispatcher_op_registration=skip_dispatcher_op_registration,\\n )\\n anonymous_gen = dest.RegisterDispatchKey(\\n backend_idx,\\n Target.ANONYMOUS_DEFINITION,\\n selector,\\n rocm=rocm,\\n symint=symint,\\n class_method_name=None,\\n skip_dispatcher_op_registration=skip_dispatcher_op_registration,\\n )\\n reg_gen = dest.RegisterDispatchKey(\\n backend_idx,\\n Target.REGISTRATION,\\n selector,\\n rocm=rocm,\\n symint=symint,\\n class_method_name=None,\\n skip_dispatcher_op_registration=skip_dispatcher_op_registration,\\n )\\n for f in grouped_native_functions:\\n kernel_namespace = get_kernel_namespace(f=f, backend_idx=backend_idx).replace(\\n \\\"::native\\\", \\\"\\\"\\n )\\n\\n ns_definitions[kernel_namespace].extend(\\n ns_gen(f),\\n )\\n anonymous_definitions[kernel_namespace].extend(\\n anonymous_gen(f),\\n )\\n namespace = (\\n f.namespace if isinstance(f, NativeFunction) else f.functional.namespace\\n )\\n if namespace not in registrations[kernel_namespace]:\\n registrations[kernel_namespace] = defaultdict(list)\\n registrations[kernel_namespace][namespace].extend(\\n reg_gen(f),\\n )\\n\\n for kernel_namespace in ns_definitions:\\n if len(ns_definitions[kernel_namespace]) == 0:\\n continue\\n ns_helper = NamespaceHelper(namespace_str=kernel_namespace)\\n registration_body = \\\"\\\"\\n for namespace in registrations[kernel_namespace]:\\n if not registrations[kernel_namespace][namespace]:\\n continue\\n registration_body += f\\\"\\\"\\\"\\nTORCH_LIBRARY_IMPL({namespace}, {dispatch_key}, m) {{\\n {newline.join(registrations[kernel_namespace][namespace])}\\n}};\\\"\\\"\\\"\\n definitions.extend(\\n fm.substitute_with_template(\\n \\\"RegisterDispatchDefinitions.ini\\\",\\n lambda: {\\n \\\"ns_prologue\\\": ns_helper.prologue,\\n \\\"ns_epilogue\\\": ns_helper.epilogue,\\n \\\"dispatch_helpers\\\": dest.gen_registration_helpers(backend_idx)\\n if gen_dispatch_helpers\\n else [],\\n \\\"dispatch_anonymous_definitions\\\": anonymous_definitions[\\n kernel_namespace\\n ],\\n \\\"static_init_dispatch_registrations\\\": \\\"\\\"\\n if skip_dispatcher_op_registration\\n else registration_body,\\n \\\"deferred_dispatch_registrations\\\": \\\"\\\",\\n \\\"dispatch_namespace\\\": dispatch_key.lower(),\\n \\\"dispatch_namespaced_definitions\\\": ns_definitions[kernel_namespace],\\n },\\n ).split(newline)\\n )\\n\\n return definitions\\n\\n\\n# Return native function declarations grouped by dispatch key and custom namespace.\\n# Used in CPUFunctions_inl.h and etc.\\ndef get_namespaced_declaration(\\n *,\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n dispatch_key: DispatchKey,\\n backend_idx: BackendIndex,\\n selector: SelectiveBuilder,\\n rocm: bool,\\n symint: bool,\\n) -> list[str]:\\n declarations: list[str] = []\\n ns_grouped_kernels: dict[str, list[str]] = defaultdict(list)\\n newline = \\\"\\\\n\\\"\\n func = dest.RegisterDispatchKey(\\n backend_idx,\\n Target.NAMESPACED_DECLARATION,\\n selector,\\n rocm=rocm,\\n class_method_name=None,\\n skip_dispatcher_op_registration=False,\\n symint=symint,\\n )\\n for f in grouped_native_functions:\\n namespace = get_kernel_namespace(f=f, backend_idx=backend_idx).replace(\\n \\\"native\\\", dispatch_key.lower()\\n )\\n\\n ns_grouped_kernels[namespace].extend(\\n func(f),\\n )\\n\\n for namespace, kernels in ns_grouped_kernels.items():\\n if len(kernels) == 0:\\n continue\\n ns_helper = NamespaceHelper(\\n namespace_str=namespace, entity_name=\\\"\\\", max_level=3\\n )\\n ordered_kernels = list(OrderedDict.fromkeys(kernels))\\n declarations.extend(\\n f\\\"\\\"\\\"\\n{ns_helper.prologue}\\n{newline.join(ordered_kernels)}\\n{ns_helper.epilogue}\\n \\\"\\\"\\\".split(\\n newline\\n )\\n )\\n return declarations\\n\\n\\n# Return native function schema registration code for aten and other namespaces.\\ndef get_native_function_schema_registrations(\\n *,\\n native_functions: Sequence[NativeFunction],\\n schema_selector: SelectiveBuilder,\\n) -> tuple[list[str], str]:\\n ns_native_functions: dict[str, list[NativeFunction]] = defaultdict(list)\\n for native_function in native_functions:\\n ns_native_functions[native_function.namespace].append(native_function)\\n schema_registrations = \\\"\\\"\\n aten_schema_registrations = []\\n custom_namespace = None\\n for namespace, funcs in ns_native_functions.items():\\n schema_registrations_body = list(\\n mapMaybe(RegisterSchema(schema_selector), funcs)\\n )\\n # NB: we have to separate aten namespace registration from other namespaces,\\n # because in the template we hardcoded an operator for ATen already.\\n if namespace == \\\"aten\\\":\\n aten_schema_registrations = schema_registrations_body\\n else:\\n custom_namespace = namespace\\n tab = \\\"\\\\t\\\"\\n # if the namespace is predefined, we should use define a library fragment\\n # instead of a new library\\n torch_library_macro = (\\n \\\"TORCH_LIBRARY_FRAGMENT\\\"\\n if namespace in FRAGMENT_NAMESPACES\\n else \\\"TORCH_LIBRARY\\\"\\n )\\n schema_registrations += f\\\"\\\"\\\"\\n{torch_library_macro}({custom_namespace}, m) {{\\n {tab.join(schema_registrations_body)}\\n}};\\\"\\\"\\\"\\n return (aten_schema_registrations, schema_registrations)\\n\\n\\ndef gen_aggregated_headers(\\n *,\\n native_functions: Sequence[NativeFunction],\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n structured_native_functions: Sequence[NativeFunctionsGroup],\\n static_dispatch_idx: list[BackendIndex],\\n selector: SelectiveBuilder,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n cpu_fm: FileManager,\\n cuda_fm: FileManager,\\n functions_keys: set[DispatchKey],\\n dispatch_keys: Sequence[DispatchKey],\\n rocm: bool,\\n) -> None:\\n # Buck doesn't support dynamic output files, so we aggregate all operator\\n # headers into a single file\\n cpu_fm.write(\\n \\\"NativeMetaFunctions.h\\\",\\n lambda: {\\n \\\"NativeMetaFunctions_includes\\\": [],\\n \\\"NativeMetaFunctions_declarations\\\": list(\\n mapMaybe(compute_meta_function_declaration, structured_native_functions)\\n ),\\n },\\n )\\n method_native_functions = [\\n fn for fn in native_functions if Variant.method in fn.variants\\n ]\\n non_method_native_functions = [\\n fn for fn in native_functions if fn not in method_native_functions\\n ]\\n cpu_fm.write(\\n \\\"MethodOperators.h\\\",\\n lambda: {\\n \\\"MethodOperators_includes\\\": [],\\n \\\"MethodOperators_declarations\\\": list(\\n mapMaybe(\\n ComputeOperators(\\n Target.DECLARATION,\\n static_dispatch_backend_indices=static_dispatch_idx,\\n ),\\n method_native_functions,\\n )\\n ),\\n },\\n )\\n cpu_fm.write(\\n \\\"Operators.h\\\",\\n lambda: {\\n \\\"Operators_includes\\\": [\\\"#include \\\"],\\n \\\"Operators_declarations\\\": list(\\n mapMaybe(\\n ComputeOperators(\\n Target.DECLARATION,\\n static_dispatch_backend_indices=static_dispatch_idx,\\n ),\\n non_method_native_functions,\\n )\\n ),\\n },\\n )\\n cpu_fm.write(\\n \\\"Functions.h\\\",\\n lambda: {\\n \\\"static_dispatch_extra_headers\\\": static_dispatch_extra_headers(\\n static_dispatch_idx\\n ),\\n \\\"Functions_includes\\\": [\\\"#include \\\"],\\n \\\"Functions_declarations\\\": list(\\n mapMaybe(\\n ComputeFunction(),\\n native_functions,\\n )\\n ),\\n },\\n )\\n declarations = get_native_function_declarations(\\n grouped_native_functions=grouped_native_functions,\\n backend_indices=backend_indices,\\n )\\n cpu_fm.write(\\n \\\"NativeFunctions.h\\\",\\n lambda: {\\n \\\"NativeFunctions_includes\\\": [\\\"#include \\\"],\\n \\\"NativeFunctions_declarations\\\": declarations,\\n },\\n )\\n\\n for dispatch_key in dispatch_keys:\\n fm = cuda_fm if is_cuda_dispatch_key(dispatch_key) else cpu_fm\\n if dispatch_key in functions_keys:\\n inl_headers = f\\\"#include \\\"\\n\\n fm.write_with_template(\\n f\\\"{dispatch_key}Functions.h\\\",\\n \\\"DispatchKeyFunctions.h\\\",\\n lambda: {\\n \\\"dispatch_key\\\": str(dispatch_key),\\n \\\"inline_headers\\\": inl_headers,\\n },\\n )\\n fm.write_with_template(\\n f\\\"{dispatch_key}Functions_inl.h\\\",\\n \\\"DispatchKeyFunctions_inl.h\\\",\\n lambda: {\\n \\\"DispatchKeyFunctions_inl_includes\\\": [],\\n \\\"dispatch_namespace\\\": dispatch_key.lower(),\\n \\\"dispatch_namespaced_declarations\\\": get_namespaced_declaration(\\n grouped_native_functions=grouped_native_functions,\\n dispatch_key=dispatch_key,\\n backend_idx=backend_indices[dispatch_key],\\n selector=selector,\\n rocm=rocm,\\n symint=True,\\n ),\\n },\\n )\\n\\n del fm\\n\\n\\ndef gen_per_operator_headers(\\n *,\\n native_functions: Sequence[NativeFunction],\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n static_dispatch_idx: list[BackendIndex],\\n selector: SelectiveBuilder,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n cpu_fm: FileManager,\\n cuda_fm: FileManager,\\n ops_fm: FileManager,\\n functions_keys: set[DispatchKey],\\n dispatch_keys: Sequence[DispatchKey],\\n rocm: bool,\\n) -> None:\\n # For CMake builds, split operator declarations into separate headers in\\n # the ATen/ops folder to split up header dependencies\\n functions_by_root_name: dict[str, list[NativeFunction]] = defaultdict(list)\\n for fn in native_functions:\\n functions_by_root_name[fn.root_name].append(fn)\\n\\n grouped_functions_by_root_name: dict[\\n str, list[NativeFunction | NativeFunctionsGroup]\\n ] = defaultdict(list)\\n for group in grouped_native_functions:\\n name = group.root_name\\n grouped_functions_by_root_name[name].append(group)\\n\\n for name, functions in functions_by_root_name.items():\\n ops_fm.write_with_template(\\n f\\\"{name}_ops.h\\\",\\n \\\"Operator.h\\\",\\n lambda: {\\n \\\"declarations\\\": list(\\n mapMaybe(\\n ComputeOperators(\\n Target.DECLARATION,\\n static_dispatch_backend_indices=static_dispatch_idx,\\n ),\\n functions,\\n )\\n ),\\n },\\n )\\n\\n ops_fm.write_with_template(\\n f\\\"{name}.h\\\",\\n \\\"Function.h\\\",\\n lambda: {\\n \\\"static_dispatch_ops_headers\\\": list(\\n mapMaybe(\\n lambda fn: static_dispatch_ops_header(\\n fn, backend_index=static_dispatch_idx\\n ),\\n functions,\\n )\\n ),\\n \\\"operator_includes\\\": f\\\"#include \\\",\\n \\\"function_definitions\\\": list(\\n mapMaybe(\\n ComputeFunction(),\\n functions,\\n )\\n ),\\n },\\n )\\n\\n grouped_functions = grouped_functions_by_root_name.get(name, [])\\n structured_functions = [\\n fn\\n for fn in grouped_functions\\n if isinstance(fn, NativeFunctionsGroup) and fn.structured\\n ]\\n is_structured = len(structured_functions) > 0\\n\\n if is_structured:\\n ops_fm.write_with_template(\\n f\\\"{name}_meta.h\\\",\\n \\\"NativeMetaFunction.h\\\",\\n lambda: {\\n \\\"meta_function_declarations\\\": list(\\n mapMaybe(\\n compute_meta_function_declaration, structured_functions\\n )\\n ),\\n },\\n )\\n declarations = get_native_function_declarations(\\n grouped_native_functions=grouped_functions,\\n backend_indices=backend_indices,\\n native_function_decl_gen=dest.compute_native_function_declaration,\\n )\\n ops_fm.write_with_template(\\n f\\\"{name}_native.h\\\",\\n \\\"NativeFunction.h\\\",\\n lambda: {\\n \\\"extra_includes\\\": (\\n f\\\"#include \\\" if is_structured else []\\n ),\\n \\\"native_function_declarations\\\": declarations,\\n },\\n )\\n\\n for category, suffix in [\\n (\\\"Functions\\\", \\\"\\\"),\\n (\\\"Operators\\\", \\\"_ops\\\"),\\n (\\\"NativeMetaFunctions\\\", \\\"_meta\\\"),\\n (\\\"NativeFunctions\\\", \\\"_native\\\"),\\n ]:\\n cpu_fm.write(\\n f\\\"{category}.h\\\",\\n lambda: {\\n f\\\"{category}_includes\\\": [\\n f\\\"#include \\\"\\n for name in sorted(functions_by_root_name.keys())\\n ],\\n f\\\"{category}_declarations\\\": [],\\n },\\n )\\n\\n for dispatch_key in dispatch_keys:\\n if dispatch_key not in functions_keys:\\n continue\\n\\n dispatch_namespace = dispatch_key.lower()\\n dispatch_names = []\\n\\n for name, functions in functions_by_root_name.items():\\n grouped_functions = grouped_functions_by_root_name.get(name, [])\\n declarations = list(\\n concatMap(\\n dest.RegisterDispatchKey(\\n backend_indices[dispatch_key],\\n Target.NAMESPACED_DECLARATION,\\n selector,\\n rocm=rocm,\\n symint=True,\\n class_method_name=None,\\n skip_dispatcher_op_registration=False,\\n ),\\n grouped_functions,\\n )\\n )\\n\\n if len(declarations) == 0:\\n continue\\n\\n dispatch_names.append(name)\\n ops_fm.write_with_template(\\n f\\\"{name}_{dispatch_namespace}_dispatch.h\\\",\\n \\\"DispatchKeyFunction.h\\\",\\n lambda: {\\n \\\"dispatch_namespace\\\": dispatch_namespace,\\n \\\"dispatch_namespaced_declarations\\\": declarations,\\n },\\n )\\n\\n fm = cuda_fm if is_cuda_dispatch_key(dispatch_key) else cpu_fm\\n inl_headers = f\\\"#include \\\"\\n\\n fm.write_with_template(\\n f\\\"{dispatch_key}Functions.h\\\",\\n \\\"DispatchKeyFunctions.h\\\",\\n lambda: {\\n \\\"dispatch_key\\\": str(dispatch_key),\\n \\\"inline_headers\\\": inl_headers,\\n },\\n )\\n fm.write_with_template(\\n f\\\"{dispatch_key}Functions_inl.h\\\",\\n \\\"DispatchKeyFunctions_inl.h\\\",\\n lambda: {\\n \\\"dispatch_namespace\\\": dispatch_namespace,\\n \\\"DispatchKeyFunctions_inl_includes\\\": [\\n f\\\"#include \\\"\\n for name in sorted(dispatch_names)\\n ],\\n \\\"dispatch_namespaced_declarations\\\": [],\\n },\\n )\\n del fm\\n\\n cpu_fm.write(\\n \\\"MethodOperators.h\\\",\\n lambda: {\\n \\\"MethodOperators_includes\\\": sorted(\\n f\\\"#include \\\"\\n for name, functions in functions_by_root_name.items()\\n if any(Variant.method in fn.variants for fn in functions)\\n ),\\n \\\"MethodOperators_declarations\\\": [],\\n },\\n )\\n\\n\\ndef gen_headers(\\n *,\\n native_functions: Sequence[NativeFunction],\\n valid_tags: set[str],\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n structured_native_functions: Sequence[NativeFunctionsGroup],\\n static_dispatch_idx: list[BackendIndex],\\n selector: SelectiveBuilder,\\n backend_indices: dict[DispatchKey, BackendIndex],\\n core_fm: FileManager,\\n cpu_fm: FileManager,\\n cuda_fm: FileManager,\\n ops_fm: FileManager,\\n dispatch_keys: Sequence[DispatchKey],\\n functions_keys: set[DispatchKey],\\n rocm: bool,\\n per_operator_headers: bool,\\n) -> None:\\n if per_operator_headers:\\n gen_per_operator_headers(\\n native_functions=native_functions,\\n grouped_native_functions=grouped_native_functions,\\n static_dispatch_idx=static_dispatch_idx,\\n selector=selector,\\n backend_indices=backend_indices,\\n cpu_fm=cpu_fm,\\n cuda_fm=cuda_fm,\\n ops_fm=ops_fm,\\n dispatch_keys=dispatch_keys,\\n functions_keys=functions_keys,\\n rocm=rocm,\\n )\\n else:\\n gen_aggregated_headers(\\n native_functions=native_functions,\\n grouped_native_functions=grouped_native_functions,\\n structured_native_functions=structured_native_functions,\\n static_dispatch_idx=static_dispatch_idx,\\n selector=selector,\\n backend_indices=backend_indices,\\n cpu_fm=cpu_fm,\\n cuda_fm=cuda_fm,\\n dispatch_keys=dispatch_keys,\\n functions_keys=functions_keys,\\n rocm=rocm,\\n )\\n\\n core_fm.write(\\n \\\"TensorBody.h\\\",\\n lambda: {\\n \\\"tensor_method_declarations\\\": list(\\n mapMaybe(\\n ComputeTensorMethod(\\n target=Target.DECLARATION,\\n static_dispatch_backend_indices=static_dispatch_idx,\\n ),\\n native_functions,\\n )\\n ),\\n \\\"tensor_method_definitions\\\": list(\\n mapMaybe(\\n ComputeTensorMethod(\\n target=Target.DEFINITION,\\n static_dispatch_backend_indices=static_dispatch_idx,\\n ),\\n native_functions,\\n )\\n ),\\n },\\n )\\n\\n cpu_fm.write(\\n \\\"RedispatchFunctions.h\\\",\\n lambda: {\\n \\\"function_redispatch_definitions\\\": list(\\n mapMaybe(ComputeRedispatchFunction(), native_functions)\\n ),\\n },\\n )\\n\\n cpu_fm.write(\\n \\\"RegistrationDeclarations.h\\\",\\n lambda: {\\n \\\"registration_declarations\\\": [\\n compute_registration_declarations(f, backend_indices)\\n for f in native_functions\\n ],\\n },\\n )\\n\\n cpu_fm.write(\\n \\\"VmapGeneratedPlumbing.h\\\", lambda: gen_all_vmap_plumbing(native_functions)\\n )\\n\\n def gen_aten_interned_strings() -> dict[str, str]:\\n attrs: set[str] = set() # All function argument names\\n names = set() # All ATen function names\\n for func in native_functions:\\n names.add(str(func.func.name.name))\\n # Some operators don't have a functional variant but we still create a\\n # symbol without the underscore\\n names.add(func.func.name.name.base)\\n\\n attrs.update(arg.name for arg in func.func.schema_order_arguments())\\n\\n # These are keywords in C++, so aren't valid symbol names\\n # https://en.cppreference.com/w/cpp/language/operator_alternative\\n names -= {\\n \\\"and\\\",\\n \\\"and_eq\\\",\\n \\\"bitand\\\",\\n \\\"bitor\\\",\\n \\\"compl\\\",\\n \\\"not\\\",\\n \\\"not_eq\\\",\\n \\\"or\\\",\\n \\\"or_eq\\\",\\n \\\"xor\\\",\\n \\\"xor_eq\\\",\\n }\\n\\n return {\\n \\\"aten_symbols\\\": \\\" \\\\\\\\\\\\n\\\".join(\\n [f\\\"_(aten, {name})\\\" for name in sorted(names)]\\n ),\\n \\\"attr_symbols\\\": \\\" \\\\\\\\\\\\n\\\".join(\\n [f\\\"_(attr, {name})\\\" for name in sorted(attrs)]\\n ),\\n }\\n\\n core_fm.write(\\\"aten_interned_strings.h\\\", gen_aten_interned_strings)\\n\\n def gen_tags_enum() -> dict[str, str]:\\n return {\\\"enum_of_valid_tags\\\": (\\\",\\\\n\\\".join(sorted(valid_tags)))}\\n\\n core_fm.write(\\\"enum_tag.h\\\", gen_tags_enum)\\n\\n\\ndef gen_source_files(\\n *,\\n native_functions: Sequence[NativeFunction],\\n grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup],\\n structured_native_functions: Sequence[NativeFunctionsGroup],\\n view_groups: Sequence[NativeFunctionsViewGroup],\\n selector: SelectiveBuilder,\\n static_dispatch_idx: list[BackendIndex],\\n backend_indices: dict[DispatchKey, BackendIndex],\\n aoti_fm: FileManager,\\n core_fm: FileManager,\\n cpu_fm: FileManager,\\n cpu_vec_fm: FileManager,\\n cuda_fm: FileManager,\\n dispatch_keys: Sequence[DispatchKey],\\n functions_keys: set[DispatchKey],\\n rocm: bool,\\n force_schema_registration: bool,\\n per_operator_headers: bool,\\n skip_dispatcher_op_registration: bool,\\n update_aoti_c_shim: bool,\\n) -> None:\\n extra_cuda_headers = \\\"\\\"\\\"\\\\\\n#include \\n#include \\n#include \\n#include \\\"\\\"\\\"\\n if rocm:\\n extra_cuda_headers = \\\"\\\"\\\"\\\\\\n#include \\n#include \\n#include \\n#include \\\"\\\"\\\"\\n\\n for dispatch_key in dispatch_keys:\\n fm = cuda_fm if is_cuda_dispatch_key(dispatch_key) else cpu_fm\\n\\n if per_operator_headers:\\n\\n def operator_headers() -> list[str]:\\n headers = []\\n for g in grouped_native_functions:\\n is_registered = False\\n if backend_index.has_kernel(g):\\n is_registered = True\\n # The above has_kernel test on a group will only test for\\n # the existence of out dispatch, because that's how\\n # structured kernels work. But sometimes functions can be\\n # grouped but not be structured, and then you need to check\\n # each individual piece, as they may have manual dispatch\\n # entries.\\n elif isinstance(g, NativeFunctionsGroup) and any(\\n backend_index.has_kernel(fn) for fn in g.functions()\\n ):\\n is_registered = True\\n # TODO: this condition is a bit questionable\\n # (It has to do with the fact that structured kernels get generated kernels\\n # to the Meta + CompositeExplicitAutogradNonFunctional keys).\\n elif g.structured and dispatch_key in (\\n DispatchKey.Meta,\\n DispatchKey.CompositeExplicitAutogradNonFunctional,\\n ):\\n is_registered = True\\n if not is_registered:\\n continue\\n\\n headers.append(f\\\"#include \\\")\\n if (\\n dispatch_key\\n == DispatchKey.CompositeExplicitAutogradNonFunctional\\n ):\\n headers.append(f\\\"#include \\\")\\n if dispatch_key in functions_keys:\\n headers.append(\\n f\\\"#include \\\"\\n )\\n\\n return sorted(set(headers))\\n\\n else:\\n\\n def operator_headers() -> list[str]:\\n headers = [\\\"#include \\\"]\\n if dispatch_key == DispatchKey.CompositeExplicitAutogradNonFunctional:\\n headers.append(\\\"#include \\\")\\n if dispatch_key in functions_keys:\\n headers.append(f\\\"#include \\\")\\n return headers\\n\\n backend_index = backend_indices[dispatch_key]\\n ns_grouped_native_functions = defaultdict(list)\\n for grouped_native_function in grouped_native_functions:\\n namespace = (\\n grouped_native_function.namespace\\n if isinstance(grouped_native_function, NativeFunction)\\n else grouped_native_function.functional.namespace\\n )\\n ns_grouped_native_functions[namespace].append(grouped_native_function)\\n\\n dispatch_namespace = str(dispatch_key).lower()\\n\\n # CompositeImplicitAutogradNestdTensor does not currently user the helpers generated\\n # compilation will fail when `-Werror=unused-function` flag is set\\n gen_dispatch_helpers: bool = (\\n dispatch_key != DispatchKey.CompositeImplicitAutogradNestedTensor\\n )\\n\\n dispatch_definitions = get_native_function_definitions(\\n fm=fm,\\n grouped_native_functions=grouped_native_functions,\\n dispatch_key=dispatch_key,\\n backend_idx=backend_index,\\n selector=selector,\\n rocm=rocm,\\n symint=True,\\n skip_dispatcher_op_registration=skip_dispatcher_op_registration,\\n gen_dispatch_helpers=gen_dispatch_helpers,\\n )\\n fm.write_with_template(\\n f\\\"Register{dispatch_key}.cpp\\\",\\n \\\"RegisterDispatchKey.cpp\\\",\\n lambda: {\\n \\\"extra_cuda_headers\\\": extra_cuda_headers\\n if is_cuda_dispatch_key(dispatch_key)\\n else \\\"\\\",\\n \\\"external_backend_headers\\\": \\\"\\\",\\n \\\"dispatch_headers\\\": dest.gen_registration_headers(\\n backend_index, per_operator_headers, rocm\\n ),\\n \\\"ops_headers\\\": operator_headers(),\\n \\\"dispatch_helpers\\\": \\\"\\\",\\n \\\"dispatch_definitions\\\": dispatch_definitions,\\n },\\n )\\n\\n for g in structured_native_functions:\\n if not g.out.ufunc_inner_loop or not is_ufunc_dispatch_key(dispatch_key):\\n continue\\n name = g.functional.func.name.name\\n if dispatch_key is DispatchKey.CPU:\\n assert fm is cpu_fm\\n fm.write_with_template(\\n f\\\"UfuncCPU_{name}.cpp\\\",\\n \\\"UfuncCPU.cpp\\\",\\n lambda: {\\n \\\"meta_declaration\\\": compute_meta_function_declaration(g),\\n \\\"native_declaration\\\": dest.compute_native_function_declaration(\\n g, backend_indices[dispatch_key]\\n ),\\n \\\"native_definitions\\\": dest.compute_ufunc_cpu(g),\\n },\\n )\\n cpu_vec_fm.write_with_template(\\n f\\\"UfuncCPUKernel_{name}.cpp\\\",\\n \\\"UfuncCPUKernel.cpp\\\",\\n lambda: {\\n \\\"name\\\": name,\\n \\\"native_definitions\\\": dest.compute_ufunc_cpu_kernel(g),\\n },\\n )\\n elif dispatch_key is DispatchKey.CUDA:\\n cuda_headers = \\\"#include \\\"\\n if rocm:\\n cuda_headers = \\\"#include \\\"\\n fm.write_with_template(\\n f\\\"UfuncCUDA_{name}.cu\\\",\\n \\\"UfuncCUDA.cu\\\",\\n lambda: {\\n \\\"name\\\": name,\\n \\\"cuda_headers\\\": cuda_headers,\\n \\\"meta_declaration\\\": compute_meta_function_declaration(g),\\n \\\"native_declaration\\\": dest.compute_native_function_declaration(\\n g, backend_indices[dispatch_key]\\n ),\\n \\\"native_definitions\\\": dest.compute_ufunc_cuda(g),\\n },\\n )\\n else:\\n raise AssertionError(f\\\"unrecognized {dispatch_key} for ufunc\\\")\\n\\n structured_func_group_dict = {}\\n for func_group in structured_native_functions:\\n for func in func_group.functions():\\n if func.structured_delegate is not None:\\n structured_func_group_dict[func.structured_delegate] = func_group\\n break\\n\\n if dispatch_key in (DispatchKey.CPU, DispatchKey.CUDA):\\n fallbacks = {}\\n for func in native_functions:\\n op_name = get_fallback_op_name(func)\\n if op_name in inductor_fallback_ops:\\n fallbacks[op_name] = func\\n fallback_native_functions = tuple(\\n value for _, value in sorted(fallbacks.items())\\n )\\n\\n # header files were checked in for ABI-compatiblilty checking\\n header_file_name = f\\\"c_shim_{dispatch_key.lower()}.h\\\"\\n new_header = gen_aoti_c_shim(\\n fallback_native_functions,\\n structured_func_group_dict,\\n dispatch_key,\\n backend_indices,\\n header=True,\\n includes=\\\"\\\",\\n )\\n if update_aoti_c_shim:\\n aoti_fm.write(\\n header_file_name,\\n lambda: new_header,\\n )\\n else:\\n try:\\n with open(\\n os.path.join(aoti_fm.install_dir, header_file_name)\\n ) as old_file:\\n old_header = old_file.read()\\n assert (\\n old_header == new_header\\n ), \\\"\\\"\\\"\\n\\nWARNING: The generated AOTInductor C shim header files have unexpectedly changed. This\\nindicates an AOTInductor fallback operator ABI backward compatibility breakage!!!\\nOnly in a limited number of situations, this is allowed:\\n\\n1. You added a fallback op to the inductor_fallback_ops list in torchgen/aoti/fallback_ops.py.\\nIf that's the case, run `python torchgen/gen.py --update-aoti-c-shim` to update the existing\\nC shim header files.\\n\\n2. You added a new default argument to an existing fallback op. This is clearly a BC breaking\\nchange in the AOTInductor land. In this case, you need to keep a manual copy of that existing\\nfallback op in a file, e.g. torch/csrc/inductor/aoti_torch/c/shim.h, bump up the version\\nnumber of that fallback op in the newly generated C shim files, and update the cpp wrapper\\ncodegen to generate the correct cpp call for this op. Contact AOTInductor team for assistance.\\n\\n \\\"\\\"\\\"\\n except FileNotFoundError:\\n print(\\n f\\\"{os.path.join(aoti_fm.install_dir, header_file_name)} not found\\\"\\n )\\n\\n # cpp files are always generated on-the-fly\\n def headers_for_aoti() -> str:\\n headers = []\\n for func in fallback_native_functions:\\n header = get_header_for_aoti(\\n func, structured_func_group_dict, dispatch_key, backend_indices\\n )\\n if header is not None:\\n headers.append(header)\\n return \\\"\\\\n\\\".join(sorted(set(headers)))\\n\\n extra_headers = (\\n extra_cuda_headers if is_cuda_dispatch_key(dispatch_key) else \\\"\\\"\\n )\\n\\n aoti_fm.write(\\n f\\\"c_shim_{dispatch_key.lower()}.cpp\\\",\\n lambda: gen_aoti_c_shim(\\n fallback_native_functions,\\n structured_func_group_dict,\\n dispatch_key,\\n backend_indices,\\n header=False,\\n includes=headers_for_aoti() + \\\"\\\\n\\\" + extra_headers,\\n ),\\n )\\n\\n del fm\\n\\n # BackendSelect is generated specially\\n def gen_backend_select() -> dict[str, list[str]]:\\n relevant_fns = [\\n fn for fn in native_functions if needs_backend_select(fn, selector)\\n ]\\n return {\\n \\\"ops_headers\\\": [\\n f\\\"#include \\\" for fn in relevant_fns\\n ],\\n \\\"backend_select_method_definitions\\\": list(\\n mapMaybe(\\n ComputeBackendSelect(Target.DEFINITION, selector), relevant_fns\\n )\\n ),\\n \\\"backend_select_function_registrations\\\": list(\\n mapMaybe(\\n ComputeBackendSelect(Target.REGISTRATION, selector), relevant_fns\\n )\\n ),\\n }\\n\\n cpu_fm.write(\\\"RegisterBackendSelect.cpp\\\", gen_backend_select)\\n\\n schema_selector = selector\\n if force_schema_registration:\\n schema_selector = SelectiveBuilder.get_nop_selector()\\n\\n (\\n aten_schema_registrations,\\n schema_registrations,\\n ) = get_native_function_schema_registrations(\\n native_functions=native_functions, schema_selector=schema_selector\\n )\\n cpu_fm.write(\\n \\\"RegisterSchema.cpp\\\",\\n lambda: {\\n \\\"aten_schema_registrations\\\": []\\n if skip_dispatcher_op_registration\\n else aten_schema_registrations,\\n \\\"schema_registrations\\\": []\\n if skip_dispatcher_op_registration\\n else schema_registrations,\\n },\\n )\\n\\n def key_func(\\n fn: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup,\\n ) -> str:\\n return fn.root_name\\n\\n cpu_fm.write_sharded(\\n \\\"Operators.cpp\\\",\\n native_functions,\\n key_fn=key_func,\\n env_callable=lambda fn: {\\n \\\"operator_headers\\\": [f\\\"#include \\\"],\\n \\\"definitions\\\": [\\n ComputeOperators(\\n Target.DEFINITION,\\n static_dispatch_backend_indices=static_dispatch_idx,\\n )(fn)\\n ],\\n },\\n base_env={\\n \\\"static_dispatch_extra_headers\\\": static_dispatch_extra_headers(\\n static_dispatch_idx\\n ),\\n },\\n num_shards=5,\\n sharded_keys={\\n \\\"operator_headers\\\",\\n \\\"definitions\\\",\\n \\\"static_dispatch_extra_headers\\\",\\n },\\n )\\n\\n cpu_fm.write(\\\"Functions.cpp\\\", dict)\\n\\n core_fm.write(\\\"TensorMethods.cpp\\\", dict)\\n\\n core_fm.write(\\n \\\"ATenOpList.cpp\\\",\\n lambda: {\\n \\\"aten_ops\\\": list(mapMaybe(compute_aten_op, native_functions)),\\n },\\n )\\n\\n def functionalization_env_callable(\\n g: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup,\\n ) -> dict[str, list[str]]:\\n def gen_op_headers(\\n g: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup,\\n ) -> list[str]:\\n if isinstance(g, NativeFunctionsViewGroup):\\n # view ops always get a functionalization kernel\\n headers = [\\n f\\\"#include \\\",\\n f\\\"#include \\\",\\n ]\\n if g.view_copy is not None:\\n headers += [\\n f\\\"#include \\\",\\n f\\\"#include \\\",\\n ]\\n return headers\\n elif isinstance(g, NativeFunctionsGroup):\\n headers = [\\n f\\\"#include \\\",\\n f\\\"#include \\\",\\n f\\\"#include \\\",\\n f\\\"#include \\\",\\n ]\\n if g.inplace is not None:\\n headers += [\\n f\\\"#include \\\",\\n f\\\"#include \\\",\\n ]\\n if g.mutable is not None:\\n headers += [\\n f\\\"#include \\\",\\n f\\\"#include \\\",\\n ]\\n return headers\\n else:\\n return [\\n f\\\"#include \\\",\\n f\\\"#include \\\",\\n ]\\n\\n return {\\n \\\"ops_headers\\\": gen_op_headers(g),\\n \\\"func_definitions\\\": gen_functionalization_definition(\\n selector,\\n g,\\n ),\\n \\\"func_registrations\\\": gen_functionalization_registration(\\n selector,\\n g,\\n backend_indices[DispatchKey.CompositeImplicitAutograd],\\n ),\\n }\\n\\n all_groups: list[\\n NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup\\n ] = list(structured_native_functions) + list(\\n view_groups # type: ignore[assignment, arg-type, operator]\\n )\\n # Note: all operators that functionalization needs to handle (mutable and aliasing ops) should be grouped properly.\\n # The only reason we really need to deal with direct NativeFunctions here (instead of the groups) is because:\\n # (1) We can provide better error checking (error out if someone introduces a mutable op that doesn't obey the grouping logic)\\n # (2) functionalization needs to manually register CompositeImplicitAutograd kernels, which might not be grouped.\\n # Although this could go away long-term if we add a dedicated dispatch key for decompositions.\\n structured_map: dict[OperatorName, NativeFunction] = {\\n f.func.name: f\\n for f in concatMap(lambda g: list(g.functions()), structured_native_functions)\\n }\\n view_map: dict[OperatorName, NativeFunction] = {\\n f.func.name: f for f in concatMap(lambda g: list(g.functions()), view_groups)\\n }\\n for f in native_functions:\\n if f.func.name not in structured_map and f.func.name not in view_map:\\n all_groups.append(f)\\n\\n cpu_fm.write_sharded(\\n \\\"RegisterFunctionalization.cpp\\\",\\n all_groups,\\n key_fn=key_func,\\n env_callable=functionalization_env_callable,\\n num_shards=4,\\n sharded_keys={\\n \\\"ops_headers\\\",\\n \\\"func_definitions\\\",\\n \\\"func_registrations\\\",\\n \\\"func_add_back_views_definitions\\\",\\n \\\"func_add_back_views_registrations\\\",\\n },\\n )\\n\\n cpu_fm.write(\\n \\\"FunctionalInverses.h\\\",\\n lambda: {\\n \\\"view_inverse_declarations\\\": list(\\n mapMaybe(\\n lambda g: gen_functionalization_view_inverse_declaration(\\n selector, g\\n ),\\n view_groups,\\n )\\n )\\n },\\n )\\n\\n # Note [view_copy NativeFunctions]\\n # Every view operator in native_functions.yaml that is not CompositeImplicitAutograd\\n # needs to have a corresponding non-aliasing {view}_copy variant.\\n # Backends that use functionalization and don't know how to handle aliasing ops\\n # are expected to implement kernels for these {view}_copy kernels instead.\\n # The code for {view}_copy operators in core is pretty boilerplate-heavy however,\\n # so we codegen the following:\\n # (1) A CompositeExplicitAutogradNonFunctional kernel for every {view}_copy operator.\\n # These are never explicitly invoked by the functionalization pass,\\n # but they could theoretically be called from user code (I added these kernels for completeness,\\n # since the ops are part of the public API).\\n # (2) A derivative formula for every {view}_copy operator\\n # {view}_copy operators can re-use the same derivative formulas as their {view} op counterparts,\\n # so rather than stamping all of the entries out in derivatives.yaml,\\n # we codegen them in.\\n # This is similar to how autograd codegen doesn't require inplace ops to have a derivatives.yaml entry.\\n cpu_fm.write(\\n \\\"CompositeViewCopyKernels.cpp\\\",\\n lambda: {\\n \\\"ops_headers\\\": [\\n \\\"\\\\n\\\".join(\\n f\\\"#include \\\\n\\\"\\n # NB: this include is important as it ensures we\\n # set the visibility on generated view_copy kernels\\n # correctly\\n f\\\"#include \\\"\\n for f in (\\n [g.view] if g.view_copy is None else [g.view, g.view_copy]\\n )\\n )\\n for g in view_groups\\n ]\\n + [\\n \\\"\\\\n\\\".join(\\n f\\\"#include \\\\n\\\"\\n # NB: this include is also important for correct visibility\\n f\\\"#include \\\"\\n for f in [g.inplace, g.mutable, g.functional]\\n if f is not None and \\\"generated\\\" not in f.tags\\n )\\n for g in structured_native_functions\\n ],\\n \\\"CompositeViewCopyKernel_Definitions\\\": list(\\n mapMaybe(\\n GenCompositeViewCopyKernel(\\n backend_indices[\\n DispatchKey.CompositeExplicitAutogradNonFunctional\\n ]\\n ),\\n view_groups,\\n )\\n ),\\n \\\"GeneratedCompositeFunctional_Definitions\\\": list(\\n mapMaybe(\\n gen_composite_functional_kernel,\\n structured_native_functions,\\n )\\n ),\\n \\\"GeneratedCompositeOut_Definitions\\\": list(\\n mapMaybe(\\n gen_composite_out_kernel,\\n structured_native_functions,\\n )\\n ),\\n },\\n )\\n\\n\\ndef gen_declarations_yaml(\\n cpu_fm: FileManager, native_functions: Sequence[NativeFunction]\\n) -> None:\\n cpu_fm.write(\\n \\\"Declarations.yaml\\\",\\n lambda: format_yaml([compute_declaration_yaml(f) for f in native_functions]),\\n )\\n\\n\\ndef get_torchgen_root() -> Path:\\n \\\"\\\"\\\"\\n If you're depending on torchgen out-of-tree, you can use the root to figure\\n out the path to native_functions.yaml\\n \\\"\\\"\\\"\\n return Path(__file__).parent.resolve()\\n\\n\\ndef main() -> None:\\n parser = argparse.ArgumentParser(description=\\\"Generate ATen source files\\\")\\n parser.add_argument(\\n \\\"-s\\\",\\n \\\"--source-path\\\",\\n help=\\\"path to source directory for ATen\\\",\\n default=\\\"aten/src/ATen\\\",\\n )\\n parser.add_argument(\\n \\\"-o\\\",\\n \\\"--output-dependencies\\\",\\n help=\\\"output a list of dependencies into the given file and exit\\\",\\n )\\n parser.add_argument(\\n \\\"--dry-run\\\",\\n action=\\\"store_true\\\",\\n help=\\\"run without writing any files (still updates outputs)\\\",\\n )\\n parser.add_argument(\\n \\\"--per-operator-headers\\\",\\n action=\\\"store_true\\\",\\n help=\\\"generate separate headers per operator in ATen/ops\\\",\\n )\\n parser.add_argument(\\n \\\"-d\\\",\\n \\\"--install-dir\\\",\\n \\\"--install_dir\\\",\\n help=\\\"output directory\\\",\\n default=\\\"build/aten/src/ATen\\\",\\n )\\n parser.add_argument(\\n \\\"--aoti-install-dir\\\",\\n \\\"--aoti_install_dir\\\",\\n help=\\\"output directory for AOTInductor shim\\\",\\n default=\\\"torch/csrc/inductor/aoti_torch/generated\\\",\\n )\\n parser.add_argument(\\n \\\"--rocm\\\",\\n action=\\\"store_true\\\",\\n help=\\\"reinterpret CUDA as ROCm/HIP and adjust filepaths accordingly\\\",\\n )\\n parser.add_argument(\\n \\\"--mps\\\",\\n action=\\\"store_true\\\",\\n help=\\\"Generate MPS registration code when set\\\",\\n )\\n # TODO: --op-registration-whitelist will be removed when all call-sites\\n # for gen.py are moved over to using the operator YAML file for mobile\\n # custom build.\\n parser.add_argument(\\n \\\"--op-registration-whitelist\\\",\\n \\\"--op_registration_whitelist\\\",\\n nargs=\\\"*\\\",\\n help=\\\"filter op registrations by the whitelist (if set); \\\"\\n \\\"each item is `namespace`::`operator name` without overload name; \\\"\\n \\\"e.g.: aten::empty aten::conv2d ...\\\",\\n )\\n parser.add_argument(\\n \\\"--op-selection-yaml-path\\\",\\n \\\"--op_selection_yaml_path\\\",\\n help=\\\"Provide a path to the operator selection (for custom build) YAML \\\"\\n \\\"that contains the information about the set of selected operators \\\"\\n \\\"and their categories (training, ...). Each operator is either a \\\"\\n \\\"full operator name with overload or just a bare operator name. \\\"\\n \\\"The operator names also contain the namespace prefix (e.g. aten::)\\\",\\n )\\n parser.add_argument(\\n \\\"--backend-whitelist\\\",\\n \\\"--backend_whitelist\\\",\\n nargs=\\\"*\\\",\\n help=\\\"filter dispatch backend by the whitelist (if set), \\\"\\n \\\"e.g.: CPU CUDA QuantizedCPU ...\\\",\\n )\\n parser.add_argument(\\n \\\"--static-dispatch-backend\\\",\\n \\\"--static_dispatch_backend\\\",\\n nargs=\\\"*\\\",\\n help=\\\"generate static dispatch code for the specific backend (if set)\\\",\\n )\\n parser.add_argument(\\n \\\"--skip-dispatcher-op-registration\\\",\\n \\\"--skip_dispatcher_op_registration\\\",\\n action=\\\"store_true\\\",\\n help=\\\"Avoid registering operators into the dispatcher.\\\",\\n )\\n parser.add_argument(\\n \\\"--force-schema-registration\\\",\\n \\\"--force_schema_registration\\\",\\n action=\\\"store_true\\\",\\n help=\\\"force it to generate schema-only registrations for all ops, including\\\"\\n \\\"those that are not listed on --op-registration-whitelist\\\",\\n )\\n parser.add_argument(\\n \\\"--generate\\\",\\n type=str,\\n nargs=\\\"*\\\",\\n choices=[\\\"headers\\\", \\\"sources\\\", \\\"declarations_yaml\\\"],\\n default=[\\\"headers\\\", \\\"sources\\\", \\\"declarations_yaml\\\"],\\n help=\\\"Generate only a subset of files\\\",\\n )\\n parser.add_argument(\\n \\\"--update-aoti-c-shim\\\",\\n action=\\\"store_true\\\",\\n help=\\\"Update AOTInductor C shim after adding an entry to inductor_fallback_ops in torchgen/aoti/fallback_ops.py. \\\"\\n \\\"WARNING: Do not use this unless you are sure what you are doing!!!\\\",\\n )\\n\\n options = parser.parse_args()\\n\\n selector = get_custom_build_selector(\\n options.op_registration_whitelist,\\n options.op_selection_yaml_path,\\n )\\n\\n native_yaml_path = os.path.join(options.source_path, \\\"native/native_functions.yaml\\\")\\n tags_yaml_path = os.path.join(options.source_path, \\\"native/tags.yaml\\\")\\n\\n from torchgen.model import dispatch_keys\\n\\n # TODO: stop generating CUDA kernels for non-CUDA builds\\n ignore_keys = set()\\n if not options.mps:\\n ignore_keys.add(DispatchKey.MPS)\\n\\n if DispatchKey.MPS in dispatch_keys:\\n del dispatch_keys[dispatch_keys.index(DispatchKey.MPS)]\\n\\n parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path, ignore_keys)\\n valid_tags = _GLOBAL_PARSE_TAGS_YAML_CACHE[tags_yaml_path]\\n native_functions, backend_indices = (\\n parsed_yaml.native_functions,\\n parsed_yaml.backend_indices,\\n )\\n\\n grouped_native_functions = get_grouped_native_functions(native_functions)\\n\\n structured_native_functions = [\\n g for g in grouped_native_functions if isinstance(g, NativeFunctionsGroup)\\n ]\\n native_functions_with_view_groups = get_grouped_by_view_native_functions(\\n native_functions\\n )\\n view_groups = [\\n g\\n for g in native_functions_with_view_groups\\n if isinstance(g, NativeFunctionsViewGroup)\\n ]\\n\\n # NB: It is mandatory to NOT use os.path.join here, as the install directory\\n # will eventually be ingested by cmake, which does not respect Windows style\\n # path slashes. If you switch this to use os.path.join, you'll get an error\\n # like:\\n #\\n # Syntax error in cmake code when parsing string\\n #\\n # C:/Jenkins/workspace/pytorch-builds/pytorch-win-ws2016-cuda9-cudnn7-py3-build/build/aten/src/ATen\\\\core/TensorMethods.h\\n #\\n # Invalid character escape '\\\\c'.\\n core_install_dir = f\\\"{options.install_dir}/core\\\"\\n Path(core_install_dir).mkdir(parents=True, exist_ok=True)\\n ops_install_dir = f\\\"{options.install_dir}/ops\\\"\\n Path(ops_install_dir).mkdir(parents=True, exist_ok=True)\\n aoti_install_dir = f\\\"{options.aoti_install_dir}\\\"\\n Path(aoti_install_dir).mkdir(parents=True, exist_ok=True)\\n\\n core_fm = make_file_manager(options=options, install_dir=core_install_dir)\\n cpu_fm = make_file_manager(options=options)\\n cpu_vec_fm = make_file_manager(options=options)\\n cuda_fm = make_file_manager(options=options)\\n ops_fm = make_file_manager(options=options, install_dir=ops_install_dir)\\n aoti_fm = make_file_manager(options=options, install_dir=aoti_install_dir)\\n\\n # Only a limited set of dispatch keys get CPUFunctions.h headers generated\\n # for them; this is the set\\n functions_keys = {\\n DispatchKey.CPU,\\n DispatchKey.CUDA,\\n DispatchKey.CompositeImplicitAutograd,\\n DispatchKey.CompositeImplicitAutogradNestedTensor,\\n DispatchKey.CompositeExplicitAutograd,\\n DispatchKey.CompositeExplicitAutogradNonFunctional,\\n DispatchKey.Meta,\\n }\\n if options.mps:\\n functions_keys.add(DispatchKey.MPS)\\n\\n if options.backend_whitelist:\\n dispatch_keys = [\\n k\\n for k in dispatch_keys\\n if is_generic_dispatch_key(k) or str(k) in options.backend_whitelist\\n ]\\n\\n static_dispatch_idx: list[BackendIndex] = []\\n if options.static_dispatch_backend:\\n static_dispatch_idx = [\\n backend_indices[DispatchKey.parse(key)]\\n for key in options.static_dispatch_backend\\n ]\\n for key in options.static_dispatch_backend:\\n dp_key = DispatchKey.parse(key)\\n if dp_key not in functions_keys:\\n functions_keys.add(dp_key)\\n\\n if \\\"sources\\\" in options.generate:\\n gen_source_files(\\n native_functions=native_functions,\\n grouped_native_functions=grouped_native_functions,\\n structured_native_functions=structured_native_functions,\\n view_groups=view_groups,\\n selector=selector,\\n static_dispatch_idx=static_dispatch_idx,\\n backend_indices=backend_indices,\\n aoti_fm=aoti_fm,\\n core_fm=core_fm,\\n cpu_fm=cpu_fm,\\n cpu_vec_fm=cpu_vec_fm,\\n cuda_fm=cuda_fm,\\n dispatch_keys=dispatch_keys,\\n functions_keys=functions_keys,\\n rocm=options.rocm,\\n force_schema_registration=options.force_schema_registration,\\n per_operator_headers=options.per_operator_headers,\\n skip_dispatcher_op_registration=options.skip_dispatcher_op_registration,\\n update_aoti_c_shim=options.update_aoti_c_shim,\\n )\\n\\n if \\\"headers\\\" in options.generate:\\n gen_headers(\\n native_functions=native_functions,\\n valid_tags=valid_tags,\\n grouped_native_functions=grouped_native_functions,\\n structured_native_functions=structured_native_functions,\\n static_dispatch_idx=static_dispatch_idx,\\n selector=selector,\\n backend_indices=backend_indices,\\n core_fm=core_fm,\\n cpu_fm=cpu_fm,\\n cuda_fm=cuda_fm,\\n ops_fm=ops_fm,\\n dispatch_keys=dispatch_keys,\\n functions_keys=functions_keys,\\n rocm=options.rocm,\\n per_operator_headers=options.per_operator_headers,\\n )\\n\\n if \\\"declarations_yaml\\\" in options.generate:\\n gen_declarations_yaml(native_functions=native_functions, cpu_fm=cpu_fm)\\n\\n if options.output_dependencies:\\n depfile_path = Path(options.output_dependencies).resolve()\\n depfile_name = depfile_path.name\\n depfile_stem = depfile_path.stem\\n\\n for fm, prefix in [\\n (cpu_fm, \\\"\\\"),\\n (cpu_vec_fm, \\\"cpu_vec_\\\"),\\n (core_fm, \\\"core_\\\"),\\n (cuda_fm, \\\"cuda_\\\"),\\n (ops_fm, \\\"ops_\\\"),\\n ]:\\n varname = prefix + depfile_stem\\n path = depfile_path.parent / (prefix + depfile_name)\\n fm.write_outputs(varname, str(path))\\n\\n\\nif __name__ == \\\"__main__\\\":\\n main()\\n\\n\\nfrom __future__ import annotations\\n\\nimport argparse\\nimport os\\nfrom collections import defaultdict\\nfrom dataclasses import dataclass\\nfrom pathlib import Path\\nfrom typing import Any, Callable, Sequence, TextIO, TYPE_CHECKING\\n\\nimport yaml\\n\\n# Parse native_functions.yaml into a sequence of NativeFunctions and Backend Indices.\\nfrom torchgen import dest\\nfrom torchgen.api import cpp as aten_cpp\\nfrom torchgen.api.types import CppSignature, CppSignatureGroup, CType, NamedCType\\nfrom torchgen.context import (\\n method_with_native_function,\\n method_with_nested_native_function,\\n with_native_function_and_index,\\n)\\nfrom torchgen.executorch.api import et_cpp\\nfrom torchgen.executorch.api.custom_ops import (\\n ComputeNativeFunctionStub,\\n gen_custom_ops_registration,\\n)\\nfrom torchgen.executorch.api.types import contextArg, ExecutorchCppSignature\\nfrom torchgen.executorch.api.unboxing import Unboxing\\nfrom torchgen.executorch.model import ETKernelIndex, ETKernelKey, ETParsedYaml\\nfrom torchgen.executorch.parse import ET_FIELDS, parse_et_yaml, parse_et_yaml_struct\\nfrom torchgen.gen import (\\n get_custom_build_selector,\\n get_native_function_declarations,\\n get_native_function_declarations_from_ns_grouped_kernels,\\n get_native_function_schema_registrations,\\n LineLoader,\\n parse_native_yaml,\\n)\\nfrom torchgen.model import (\\n BackendIndex,\\n BackendMetadata,\\n DEFAULT_KERNEL_NAMESPACE,\\n DispatchKey,\\n FunctionSchema,\\n Location,\\n NativeFunction,\\n NativeFunctionsGroup,\\n OperatorName,\\n Variant,\\n)\\nfrom torchgen.utils import (\\n context,\\n FileManager,\\n make_file_manager,\\n mapMaybe,\\n NamespaceHelper,\\n)\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.selective_build.selector import SelectiveBuilder\\n\\n\\ndef _sig_decl_wrapper(sig: CppSignature | ExecutorchCppSignature) -> str:\\n \\\"\\\"\\\"\\n A wrapper function to basically get `sig.decl(include_context=True)`.\\n For ATen kernel, the codegen has no idea about ET contextArg, so we\\n use this wrapper to add it.\\n \\\"\\\"\\\"\\n if isinstance(sig, ExecutorchCppSignature):\\n return sig.decl()\\n\\n returns_type = aten_cpp.returns_type(sig.func.returns).cpp_type()\\n cpp_args = [a.decl() for a in sig.arguments()]\\n cpp_args_str = \\\", \\\".join([contextArg.decl()] + cpp_args)\\n sig_decl = f\\\"{returns_type} {sig.name()}({cpp_args_str})\\\"\\n return sig_decl\\n\\n\\ndef static_dispatch(\\n sig: CppSignature | ExecutorchCppSignature,\\n f: NativeFunction,\\n backend_indices: list[BackendIndex],\\n) -> str:\\n \\\"\\\"\\\"\\n For a given `NativeFunction`, find out the corresponding native function and dispatch to it. If zero or more than one\\n native function exists, error out. A simplified version of register_dispatch_key.py\\n Arguments:\\n sig: A CppSignature for this native function we want to use.\\n f: NativeFunction to generate static dispatch.\\n backend_indices: All available backends.\\n Return:\\n C++ code to call backend-specific functions, e.g., \\\"return at::native::add(self, other, scale);\\\"\\n \\\"\\\"\\\"\\n if len(backend_indices) == 0 or f.manual_kernel_registration:\\n return \\\"\\\"\\n\\n backends = [b for b in backend_indices if b.has_kernel(f)]\\n static_block = None\\n if len(backends) == 1:\\n backend_metadata = backends[0].get_kernel(f)\\n if backend_metadata:\\n args = \\\", \\\".join(a.name for a in sig.arguments())\\n # Here we are assuming there's no difference between CppSignature and NativeSignature for Executorch.\\n static_block = f\\\"return ::{backend_metadata.cpp_namespace}::{backend_metadata.kernel}({args});\\\"\\n else:\\n static_block = f\\\"\\\"\\\"\\nET_ASSERT_UNREACHABLE_MSG(\\\"The number of native function(s) binding to {f.func.name} is {len(backends)}.\\\");\\n \\\"\\\"\\\"\\n return f\\\"\\\"\\\"\\n// {f.namespace}::{f.func}\\nTORCH_API inline {_sig_decl_wrapper(sig)} {{\\n {static_block}\\n}}\\n\\\"\\\"\\\"\\n\\n\\n# Generates Functions.h, which provides the functional public C++ API,\\n# and the scaffolding to call into the dispatcher from these functions.\\n@dataclass(frozen=True)\\nclass ComputeFunction:\\n static_dispatch_backend_indices: list[BackendIndex]\\n\\n selector: SelectiveBuilder\\n\\n use_aten_lib: bool\\n\\n is_custom_op: Callable[[NativeFunction], bool]\\n\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str | None:\\n is_method_variant = False\\n if not self.selector.is_root_operator(f\\\"{f.namespace}::{f.func.name}\\\"):\\n return None\\n\\n if Variant.function not in f.variants and Variant.method in f.variants:\\n is_method_variant = True\\n\\n # only valid remaining case is only function is in f.variants\\n elif not (Variant.function in f.variants and Variant.method not in f.variants):\\n raise Exception( # noqa: TRY002\\n f\\\"Can't handle native function {f.func} with the following variant specification {f.variants}.\\\"\\n )\\n\\n sig: CppSignature | ExecutorchCppSignature = (\\n CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=f.manual_cpp_binding\\n ).most_faithful_signature()\\n if self.use_aten_lib\\n else ExecutorchCppSignature.from_native_function(f)\\n )\\n if self.use_aten_lib and not self.is_custom_op(f):\\n comma = \\\", \\\"\\n\\n if is_method_variant:\\n return f\\\"\\\"\\\"\\n// {f.namespace}::{f.func}\\nTORCH_API inline {_sig_decl_wrapper(sig)} {{\\n return {sig.arguments()[0].name}.{sig.name()}({comma.join(e.name for e in sig.arguments()[1:])});\\n}}\\n\\\"\\\"\\\"\\n else:\\n return f\\\"\\\"\\\"\\n// {f.namespace}::{f.func}\\nTORCH_API inline {_sig_decl_wrapper(sig)} {{\\n return at::{sig.name()}({comma.join(e.name for e in sig.arguments())});\\n}}\\n\\\"\\\"\\\"\\n\\n else:\\n return static_dispatch(\\n sig,\\n f,\\n backend_indices=self.static_dispatch_backend_indices,\\n )\\n\\n\\n# Generates RegisterCodegenUnboxedKernels.cpp.\\n@dataclass(frozen=True)\\nclass ComputeCodegenUnboxedKernels:\\n selector: SelectiveBuilder\\n\\n use_aten_lib: bool\\n\\n @method_with_nested_native_function\\n def __call__(\\n self,\\n unbox_kernel_entry: tuple[NativeFunction, tuple[ETKernelKey, BackendMetadata]],\\n ) -> str:\\n f: NativeFunction = unbox_kernel_entry[0]\\n kernel_key: ETKernelKey | list[ETKernelKey] = unbox_kernel_entry[1][0]\\n kernel_meta: BackendMetadata = unbox_kernel_entry[1][1]\\n\\n op_name = f\\\"{f.namespace}::{f.func.name}\\\"\\n if not self.selector.is_root_operator(op_name):\\n return \\\"\\\"\\n\\n if not isinstance(kernel_key, list):\\n kernel_key = [kernel_key]\\n used_kernel_keys = self.selector.et_get_selected_kernels(\\n op_name, [k.to_native_string() for k in kernel_key]\\n )\\n if not used_kernel_keys:\\n return \\\"\\\"\\n sig: CppSignature | ExecutorchCppSignature\\n argument_type_gen: Callable[..., NamedCType]\\n return_type_gen: Callable[..., CType]\\n if self.use_aten_lib:\\n sig = CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=f.manual_cpp_binding\\n ).most_faithful_signature()\\n argument_type_gen = aten_cpp.argumenttype_type\\n return_type_gen = aten_cpp.returns_type\\n arguments = sig.arguments()\\n kernel_call = f\\\"torch::executor::{f.namespace}::{sig.name()}\\\"\\n else:\\n sig = ExecutorchCppSignature.from_native_function(f)\\n argument_type_gen = et_cpp.argumenttype_type\\n return_type_gen = et_cpp.returns_type\\n arguments = sig.arguments(include_context=False)\\n kernel_call = f\\\"{kernel_meta.cpp_namespace}::{kernel_meta.kernel}\\\"\\n # parse arguments into C++ code\\n binding_list, code_list = Unboxing(\\n argument_type_gen=argument_type_gen\\n ).convert_arguments(arguments)\\n\\n # for each C++ argument, generate the conversion code\\n code_connector = \\\"\\\\n\\\\t\\\"\\n arg_connector = \\\", \\\"\\n\\n args_str = f\\\"{arg_connector.join(e.name for e in binding_list)}\\\"\\n event_tracer_output_logging = \\\"\\\"\\n output_ids = []\\n\\n if len(f.func.returns) == 0:\\n if len(f.func.arguments.out) == 0:\\n raise Exception( # noqa: TRY002\\n f\\\"Can't handle native function {f.func} with no returns and no out yet.\\\"\\n )\\n out = f.func.arguments.out[0]\\n return_assignment = f\\\"\\\"\\\"stack[{len(binding_list)}] = &{out.name};\\\"\\\"\\\"\\n ret_prefix = \\\"\\\"\\n output_ids = [len(binding_list)]\\n else:\\n if len(f.func.arguments.out) == 0:\\n return_assignment = (\\n f\\\"\\\"\\\"*stack[{len(binding_list)}] = EValue(result_);\\\"\\\"\\\"\\n )\\n ret_prefix = return_type_gen(f.func.returns).cpp_type() + \\\" result_ = \\\"\\n output_ids = [len(binding_list)]\\n else:\\n return_assignment = \\\"\\\"\\n ret_prefix = \\\"\\\"\\n output_ids = [\\n len(binding_list) - (i + 1)\\n for i in reversed(range(len(f.func.arguments.out)))\\n ]\\n\\n for output_id in output_ids:\\n event_tracer_output_logging += (\\n f\\\"internal::event_tracer_log_evalue(\\\"\\n f\\\"context.internal_event_tracer(), \\\"\\n f\\\"*stack[{output_id}]);\\\\n\\\"\\n )\\n\\n newline = \\\"\\\\n \\\"\\n return \\\"\\\\n\\\".join(\\n [\\n f\\\"\\\"\\\"\\nKernel(\\n \\\"{f.namespace}::{f.func.name}\\\",{newline + '\\\"' + (k + '\\\",') if k != 'default' else ''}\\n []({contextArg.defn()}, EValue** stack) {{\\n {code_connector.join(code_list)}\\n\\n internal::EventTracerProfileScope event_tracer_scope(context.internal_event_tracer(), \\\"native_call_{f.func.name}\\\");\\n EXECUTORCH_SCOPE_PROF(\\\"native_call_{f.func.name}\\\");\\n {ret_prefix}{kernel_call}(context, {args_str});\\n {event_tracer_output_logging}\\n {return_assignment}\\n }}\\n),\\n\\\"\\\"\\\"\\n for k in used_kernel_keys\\n ]\\n )\\n\\n\\ndef gen_unboxing(\\n *,\\n native_functions: Sequence[NativeFunction],\\n cpu_fm: FileManager,\\n selector: SelectiveBuilder,\\n use_aten_lib: bool,\\n kernel_index: ETKernelIndex,\\n manual_registration: bool,\\n) -> None:\\n # Iterable type for write_sharded is a Tuple of (native_function, (kernel_key, metadata))\\n def key_func(\\n item: tuple[NativeFunction, tuple[ETKernelKey, BackendMetadata]]\\n ) -> str:\\n return item[0].root_name + \\\":\\\" + item[1][0].to_native_string()\\n\\n items: list[tuple[NativeFunction, tuple[ETKernelKey, BackendMetadata]]] = [\\n (native_function, (kernel_key, metadata))\\n for native_function in native_functions\\n for kernel_key, metadata in kernel_index.get_kernels(native_function).items()\\n ]\\n\\n header = [\\\"Functions.h\\\" if use_aten_lib else \\\"NativeFunctions.h\\\"]\\n filename = (\\n \\\"RegisterKernels.cpp\\\"\\n if manual_registration\\n else \\\"RegisterCodegenUnboxedKernels.cpp\\\"\\n )\\n cpu_fm.write_sharded(\\n filename,\\n items,\\n key_fn=key_func,\\n env_callable=lambda unbox_kernel_entry: {\\n \\\"unboxed_kernels\\\": [\\n ComputeCodegenUnboxedKernels(selector, use_aten_lib)(unbox_kernel_entry)\\n ],\\n \\\"fn_header\\\": header\\n if unbox_kernel_entry == items[0]\\n else [], # Only write header once\\n },\\n num_shards=1,\\n sharded_keys={\\\"unboxed_kernels\\\", \\\"fn_header\\\"},\\n )\\n\\n\\n@with_native_function_and_index # type: ignore[arg-type]\\ndef compute_native_function_declaration(\\n g: NativeFunctionsGroup | NativeFunction, kernel_index: ETKernelIndex\\n) -> list[str]:\\n assert isinstance(g, NativeFunction)\\n sig = ExecutorchCppSignature.from_native_function(f=g)\\n metadata_list = kernel_index.get_kernels(g).values()\\n if metadata_list is None:\\n return []\\n\\n # for kernels in lean mode, we declare two versions, one with context and one without.\\n # In the end we will cleanup the unused one.\\n def gen_decl(metadata: BackendMetadata, include_context: bool) -> str:\\n return f\\\"{sig.decl(name=metadata.kernel, include_context=include_context)};\\\"\\n\\n return [\\n gen_decl(metadata, include_context)\\n for include_context in [False, True]\\n for metadata in metadata_list\\n ]\\n\\n\\ndef gen_functions_declarations(\\n *,\\n native_functions: Sequence[NativeFunction],\\n kernel_index: ETKernelIndex,\\n selector: SelectiveBuilder,\\n use_aten_lib: bool,\\n custom_ops_native_functions: Sequence[NativeFunction] | None = None,\\n) -> str:\\n \\\"\\\"\\\"\\n Generates namespace separated C++ function API inline declaration/definitions.\\n Native functions are grouped by namespaces and the generated code is wrapped inside\\n namespace blocks.\\n\\n E.g., for `custom_1::foo.out` in yaml file we will generate a C++ API as a symbol\\n in `torch::executor::custom_1::foo_out`. This way we avoid symbol conflict when\\n the other `custom_2::foo.out` is available.\\n \\\"\\\"\\\"\\n\\n # convert kernel index to BackendIndex. This is because we can't handle ETKernelIndex yet.\\n # TODO larryliu: evaluate if this code is still needed. If yes let it handle ETKernelIndex.\\n\\n backend_index = kernel_index._to_backend_index()\\n\\n ns_grouped_functions = defaultdict(list)\\n for native_function in native_functions:\\n ns_grouped_functions[native_function.namespace].append(native_function)\\n functions_declarations = \\\"\\\"\\n newline = \\\"\\\\n\\\"\\n for namespace in ns_grouped_functions:\\n ns_helper = NamespaceHelper(\\n namespace_str=namespace,\\n entity_name=\\\"\\\",\\n max_level=3,\\n )\\n declarations = list(\\n mapMaybe(\\n ComputeFunction(\\n static_dispatch_backend_indices=[backend_index],\\n selector=selector,\\n use_aten_lib=use_aten_lib,\\n is_custom_op=lambda f: custom_ops_native_functions is not None\\n and f in custom_ops_native_functions,\\n ),\\n ns_grouped_functions[namespace],\\n )\\n )\\n functions_declarations += f\\\"\\\"\\\"\\n{ns_helper.prologue}\\n{newline.join(declarations)}\\n{ns_helper.epilogue}\\n \\\"\\\"\\\"\\n return functions_declarations\\n\\n\\ndef get_ns_grouped_kernels(\\n *,\\n native_functions: Sequence[NativeFunction],\\n kernel_index: ETKernelIndex,\\n native_function_decl_gen: Callable[\\n [\\n NativeFunctionsGroup | NativeFunction,\\n ETKernelIndex,\\n ],\\n list[str],\\n ],\\n) -> dict[str, list[str]]:\\n ns_grouped_kernels: dict[str, list[str]] = defaultdict(list)\\n for f in native_functions:\\n native_function_namespaces = set()\\n op_kernels = kernel_index.get_kernels(f)\\n for backend_metadata in op_kernels.values():\\n if backend_metadata:\\n namespace = backend_metadata.cpp_namespace\\n native_function_namespaces.add(namespace)\\n else:\\n namespace = DEFAULT_KERNEL_NAMESPACE\\n assert (\\n len(native_function_namespaces) <= 1\\n ), f\\\"Codegen only supports one namespace per operator, got {native_function_namespaces}\\\"\\n ns_grouped_kernels[namespace].extend(\\n native_function_decl_gen(f, kernel_index)\\n )\\n return ns_grouped_kernels\\n\\n\\ndef gen_headers(\\n *,\\n native_functions: Sequence[NativeFunction],\\n gen_custom_ops_header: bool,\\n custom_ops_native_functions: Sequence[NativeFunction],\\n selector: SelectiveBuilder,\\n kernel_index: ETKernelIndex,\\n cpu_fm: FileManager,\\n use_aten_lib: bool,\\n) -> None:\\n \\\"\\\"\\\"Generate headers.\\n\\n Args:\\n native_functions (Sequence[NativeFunction]): a collection of NativeFunction for ATen ops.\\n gen_custom_ops_header (bool): whether we should generate CustomOpsNativeFunctions.h\\n custom_ops_native_functions (Sequence[NativeFunction]): a collection of NativeFunction for custom ops.\\n kernel_index (ETKernelIndex): kernel collection\\n cpu_fm (FileManager): file manager manages output stream\\n use_aten_lib (bool): whether we are generating for PyTorch types or Executorch types.\\n \\\"\\\"\\\"\\n aten_headers = [\\\"#include \\\"]\\n backend_indices = {DispatchKey.CPU: kernel_index._to_backend_index()}\\n if gen_custom_ops_header:\\n cpu_fm.write_with_template(\\n \\\"CustomOpsNativeFunctions.h\\\",\\n \\\"NativeFunctions.h\\\",\\n lambda: {\\n \\\"nativeFunctions_declarations\\\": get_native_function_declarations(\\n grouped_native_functions=custom_ops_native_functions,\\n backend_indices=backend_indices,\\n native_function_decl_gen=dest.compute_native_function_declaration,\\n ),\\n \\\"headers\\\": [\\n \\\"#include \\\",\\n \\\"#include \\\",\\n ],\\n },\\n )\\n aten_headers.append('#include \\\"CustomOpsNativeFunctions.h\\\"')\\n cpu_fm.write(\\n \\\"Functions.h\\\",\\n lambda: {\\n \\\"static_dispatch_extra_headers\\\": aten_headers\\n if use_aten_lib\\n else ['#include \\\"NativeFunctions.h\\\"'],\\n \\\"Functions_declarations\\\": gen_functions_declarations(\\n native_functions=native_functions,\\n kernel_index=kernel_index,\\n selector=selector,\\n use_aten_lib=use_aten_lib,\\n custom_ops_native_functions=custom_ops_native_functions,\\n ),\\n },\\n )\\n cpu_fm.write(\\n \\\"RegisterKernels.h\\\",\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\" + \\\"generated by torchgen/gen_executorch.py\\\",\\n },\\n )\\n headers = {\\n \\\"headers\\\": [\\n \\\"#include // at::Tensor etc.\\\",\\n \\\"#include \\\",\\n ],\\n }\\n if use_aten_lib:\\n headers[\\\"headers\\\"].append(\\\"#include // TORCH_API\\\")\\n cpu_fm.write(\\n \\\"NativeFunctions.h\\\",\\n lambda: dict(\\n {\\n \\\"nativeFunctions_declarations\\\": get_native_function_declarations(\\n grouped_native_functions=native_functions,\\n backend_indices=backend_indices,\\n native_function_decl_gen=dest.compute_native_function_declaration,\\n ),\\n },\\n **headers,\\n ),\\n )\\n else:\\n ns_grouped_kernels = get_ns_grouped_kernels(\\n native_functions=native_functions,\\n kernel_index=kernel_index,\\n native_function_decl_gen=compute_native_function_declaration, # type: ignore[arg-type]\\n )\\n cpu_fm.write(\\n \\\"NativeFunctions.h\\\",\\n lambda: dict(\\n {\\n \\\"nativeFunctions_declarations\\\": get_native_function_declarations_from_ns_grouped_kernels(\\n ns_grouped_kernels=ns_grouped_kernels,\\n ),\\n },\\n **headers,\\n ),\\n )\\n\\n\\ndef gen_custom_ops(\\n *,\\n native_functions: Sequence[NativeFunction],\\n selector: SelectiveBuilder,\\n kernel_index: ETKernelIndex,\\n cpu_fm: FileManager,\\n rocm: bool,\\n) -> None:\\n dispatch_key = DispatchKey.CPU\\n (\\n anonymous_definition,\\n static_init_dispatch_registrations,\\n ) = gen_custom_ops_registration(\\n native_functions=native_functions,\\n selector=selector,\\n kernel_index=kernel_index,\\n rocm=rocm,\\n )\\n cpu_fm.write_with_template(\\n f\\\"Register{dispatch_key}CustomOps.cpp\\\",\\n \\\"RegisterDispatchKeyCustomOps.cpp\\\",\\n lambda: {\\n \\\"ops_headers\\\": '#include \\\"CustomOpsNativeFunctions.h\\\"',\\n \\\"DispatchKey\\\": dispatch_key,\\n \\\"dispatch_namespace\\\": dispatch_key.lower(),\\n \\\"dispatch_namespaced_definitions\\\": \\\"\\\",\\n \\\"dispatch_anonymous_definitions\\\": anonymous_definition,\\n \\\"static_init_dispatch_registrations\\\": static_init_dispatch_registrations,\\n },\\n )\\n cpu_fm.write_with_template(\\n f\\\"Register{dispatch_key}Stub.cpp\\\",\\n \\\"RegisterDispatchKeyCustomOps.cpp\\\",\\n lambda: {\\n \\\"ops_headers\\\": \\\"\\\",\\n \\\"DispatchKey\\\": dispatch_key,\\n \\\"dispatch_namespace\\\": dispatch_key.lower(),\\n \\\"dispatch_namespaced_definitions\\\": \\\"\\\",\\n \\\"dispatch_anonymous_definitions\\\": list(\\n mapMaybe(ComputeNativeFunctionStub(), native_functions)\\n ),\\n \\\"static_init_dispatch_registrations\\\": static_init_dispatch_registrations,\\n },\\n )\\n\\n (\\n aten_schema_registrations,\\n schema_registrations,\\n ) = get_native_function_schema_registrations(\\n native_functions=native_functions,\\n schema_selector=selector,\\n )\\n cpu_fm.write(\\n \\\"RegisterSchema.cpp\\\",\\n lambda: {\\n \\\"schema_registrations\\\": schema_registrations,\\n \\\"aten_schema_registrations\\\": aten_schema_registrations,\\n },\\n )\\n\\n\\ndef translate_native_yaml(\\n tags_yaml_path: str,\\n aten_yaml_path: str,\\n native_yaml_path: str | None,\\n use_aten_lib: bool,\\n out_file: TextIO,\\n) -> None:\\n \\\"\\\"\\\"Translates Executorch DSL dialect to use the same syntax as\\n native_functions.yaml. The major difference is that Executorch DSL dialect\\n supports \\\"op\\\" key, where it refers to the operator name in native_functions.yaml.\\n\\n For example, a functions.yaml may have the following entry:\\n\\n - op: add.out\\n ...\\n\\n It needs to be translated to the following:\\n\\n - func: add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)\\n ...\\n\\n We go in aten_yaml_path and find the operator schema for \\\"add.out\\\" and add it\\n to the original functions.yaml. We also add required field \\\"variants\\\", where for\\n Executorch it will always be \\\"function\\\".\\n\\n For ATen mode we don't have to do the translation because native_yaml_path is\\n the same as native_functions.yaml.\\n\\n Args:\\n tags_yaml_path: Path to a tags.yaml file to satisfy codegen parsing.\\n It is not optional.\\n aten_yaml_path: Path to ATen operator yaml file native_functions.yaml.\\n native_yaml_path: Path to a functions.yaml file to parse.\\n If the path does not exist in the filesystem, it is treated as an\\n empty file. If `custom_ops_yaml_path` exists, the contents of that\\n file are appended to the yaml input to be parsed.\\n use_aten_lib: We use this flag to determine if we want to generate native\\n functions. In ATen mode we should generate out= variants.\\n out_file: The IO object that we are writing into.\\n Returns:\\n None\\n \\\"\\\"\\\"\\n if use_aten_lib:\\n with open(aten_yaml_path) as aten_yaml:\\n out_file.writelines(aten_yaml.readlines())\\n return\\n\\n native_functions, persisted_fields = parse_et_yaml(\\n aten_yaml_path,\\n tags_yaml_path,\\n None,\\n skip_native_fns_gen=False,\\n )\\n\\n func_to_scoped_name: dict[FunctionSchema, str] = {\\n f.func: f\\\"{f.namespace}::{f.func.name}\\\" for f in native_functions\\n }\\n op_to_scoped_name: dict[OperatorName, str] = {\\n func.name: name for func, name in func_to_scoped_name.items()\\n }\\n\\n schema_dict = {name: str(func) for func, name in func_to_scoped_name.items()}\\n kernel_persist_dict: dict[str, dict[str, Any]] = {\\n op_to_scoped_name[op]: v for op, v in persisted_fields.items()\\n }\\n\\n if (\\n not native_yaml_path\\n or not os.path.exists(native_yaml_path)\\n or os.stat(native_yaml_path).st_size == 0\\n ):\\n return\\n with open(native_yaml_path) as native_yaml:\\n native_es = yaml.load(native_yaml, Loader=LineLoader)\\n if not native_es:\\n return\\n for e in native_es:\\n assert isinstance(e.get(\\\"__line__\\\"), int), e\\n loc = Location(native_yaml_path, e.pop(\\\"__line__\\\"))\\n with context(lambda: f\\\"in {loc}:\\\\n \\\"):\\n if \\\"variants\\\" not in e:\\n e[\\\"variants\\\"] = \\\"function\\\"\\n if \\\"func\\\" in e:\\n continue\\n assert isinstance(e.get(\\\"op\\\"), str), e\\n opname = e.pop(\\\"op\\\")\\n if \\\"::\\\" not in opname:\\n opname = \\\"aten::\\\" + opname\\n assert opname in schema_dict\\n e[\\\"func\\\"] = schema_dict.get(opname)\\n\\n # Write out persisted kernel information\\n if opname in kernel_persist_dict:\\n for k, v in kernel_persist_dict[opname].items():\\n e[k] = v\\n\\n yaml.dump(native_es, out_file, width=1000)\\n\\n\\ndef parse_yaml(\\n path: str | None,\\n tags_yaml_path: str,\\n function_filter: Callable[[NativeFunction], bool],\\n skip_native_fns_gen: bool = False,\\n) -> tuple[\\n list[NativeFunction],\\n dict[DispatchKey, dict[OperatorName, BackendMetadata]] | ETKernelIndex,\\n]:\\n if path and os.path.exists(path) and os.stat(path).st_size > 0:\\n with open(path) as f:\\n es = yaml.load(f, Loader=LineLoader)\\n\\n # Check for kernel index structure\\n kernel_index = (\\n parse_et_yaml_struct(es) if any(\\\"kernels\\\" in e for e in es) else None\\n )\\n\\n # Remove ET specific fields from entries for BC compatibility\\n for entry in es:\\n for field in ET_FIELDS:\\n entry.pop(field, None)\\n\\n parsed_yaml = parse_native_yaml(\\n path,\\n tags_yaml_path,\\n None,\\n skip_native_fns_gen=skip_native_fns_gen,\\n loaded_yaml=es,\\n )\\n native_functions = list(filter(function_filter, parsed_yaml.native_functions))\\n op_names = [f.func.name for f in native_functions]\\n\\n # (1) Return ETKernelIndex if kernel index is present\\n if kernel_index is not None:\\n filtered_index = {\\n op_name: kernel_mapping\\n for op_name, kernel_mapping in kernel_index.index.items()\\n if op_name in op_names\\n }\\n return native_functions, ETKernelIndex(index=filtered_index)\\n\\n # (2) Return BackendIndices if kernel index is absent\\n def map_index(\\n m: dict[OperatorName, BackendMetadata]\\n ) -> dict[OperatorName, BackendMetadata]:\\n return {op: m[op] for op in m if op in op_names}\\n\\n backend_indices = {\\n k: map_index(b.index) for (k, b) in parsed_yaml.backend_indices.items()\\n }\\n\\n return native_functions, backend_indices\\n else:\\n return [], {}\\n\\n\\ndef parse_yaml_files(\\n tags_yaml_path: str,\\n aten_yaml_path: str,\\n native_yaml_path: str | None,\\n custom_ops_yaml_path: str | None,\\n selector: SelectiveBuilder,\\n use_aten_lib: bool,\\n) -> tuple[ETParsedYaml, ETParsedYaml | None]:\\n \\\"\\\"\\\"Parses functions.yaml and custom_ops.yaml files.\\n\\n Args:\\n tags_yaml_path: Path to a tags.yaml file to satisfy codegen parsing.\\n It is not optional.\\n aten_yaml_path: Path to ATen operator yaml file native_functions.yaml.\\n native_yaml_path: Path to a functions.yaml file to parse.\\n If the path does not exist in the filesystem, it is treated as an\\n empty file. If `custom_ops_yaml_path` exists, the contents of that\\n file are appended to the yaml input to be parsed.\\n custom_ops_yaml_path: Path to a custom_ops.yaml file to parse. If\\n the path does not exist in the filesystem, it is ignored.\\n selector: For selective build.\\n use_aten_lib: We use this flag to determine if we want to generate native\\n functions. In ATen mode we should generate out= variants.\\n Returns:\\n A tuple with two elements:\\n [0]: The parsed results of concatenating the contents of\\n `native_yaml_path` and `custom_ops_yaml_path`.\\n [1]: The parsed results of the contents of `custom_ops_yaml_path`, if\\n present. If not present, None.\\n \\\"\\\"\\\"\\n import tempfile\\n\\n # only include selected ops, this is because we want to avoid\\n def function_filter(f: NativeFunction) -> bool:\\n return selector.is_native_function_selected(f)\\n\\n with tempfile.TemporaryDirectory() as tmpdirname:\\n translated_yaml_path = os.path.join(tmpdirname, \\\"translated.yaml\\\")\\n with open(translated_yaml_path, \\\"w\\\") as translated:\\n translate_native_yaml(\\n tags_yaml_path,\\n aten_yaml_path,\\n native_yaml_path,\\n use_aten_lib,\\n translated,\\n )\\n\\n translated_functions, translated_indices = parse_yaml(\\n translated_yaml_path, tags_yaml_path, function_filter, not use_aten_lib\\n )\\n custom_ops_functions, custom_ops_indices = parse_yaml(\\n custom_ops_yaml_path, tags_yaml_path, function_filter, True\\n )\\n\\n # Convert BackendIndices to ETKernelIndex\\n if not isinstance(translated_indices, ETKernelIndex):\\n translated_indices = ETKernelIndex.from_backend_indices(translated_indices)\\n if not isinstance(custom_ops_indices, ETKernelIndex):\\n custom_ops_indices = ETKernelIndex.from_backend_indices(custom_ops_indices)\\n\\n combined_functions = translated_functions + custom_ops_functions\\n combined_kernel_index = ETKernelIndex.merge_indices(\\n translated_indices, custom_ops_indices\\n )\\n combined_yaml = ETParsedYaml(combined_functions, combined_kernel_index)\\n custom_ops_parsed_yaml = ETParsedYaml(custom_ops_functions, custom_ops_indices)\\n\\n return combined_yaml, custom_ops_parsed_yaml\\n\\n\\ndef main() -> None:\\n parser = argparse.ArgumentParser(description=\\\"Generate operator source files\\\")\\n # Although we don't refer to --source-path directly, make_file_manager()\\n # expects it to point to a directory that contains a templates/ subdirectory\\n # containing the file templates.\\n parser.add_argument(\\n \\\"-s\\\",\\n \\\"--source-path\\\",\\n help=\\\"path to source directory for kernel templates\\\",\\n )\\n parser.add_argument(\\n \\\"--functions-yaml-path\\\",\\n \\\"--functions_yaml_path\\\",\\n help=\\\"path to the functions.yaml file to use. Optional, but at least \\\"\\n \\\"one of --functions-yaml-path and --custom-ops-yaml-path must be \\\"\\n \\\"specified.\\\",\\n )\\n parser.add_argument(\\n \\\"--custom-ops-yaml-path\\\",\\n \\\"--custom_ops_yaml_path\\\",\\n help=\\\"path to the custom_ops.yaml file to use. Optional, but at least \\\"\\n \\\"one of --functions-yaml-path and --custom-ops-yaml-path must be \\\"\\n \\\"specified.\\\",\\n )\\n parser.add_argument(\\n \\\"--aten-yaml-path\\\",\\n \\\"--aten_yaml_path\\\",\\n help=\\\"path to native_functions.yaml file.\\\",\\n )\\n # Note that make_file_manager() also looks at --install-dir.\\n parser.add_argument(\\n \\\"-d\\\",\\n \\\"--install-dir\\\",\\n \\\"--install_dir\\\",\\n help=\\\"output directory\\\",\\n default=\\\"build/generated\\\",\\n )\\n parser.add_argument(\\n \\\"-o\\\",\\n \\\"--output-dependencies\\\",\\n help=\\\"output a list of dependencies into the given file and exit\\\",\\n )\\n # Although we don't refer to --dry-run directly, make_file_manager() looks\\n # for it.\\n parser.add_argument(\\n \\\"--dry-run\\\",\\n action=\\\"store_true\\\",\\n help=\\\"run without writing any files (still updates outputs)\\\",\\n )\\n parser.add_argument(\\n \\\"--static-dispatch-backend\\\",\\n \\\"--static_dispatch_backend\\\",\\n nargs=\\\"*\\\",\\n help=\\\"generate static dispatch code for the specific backend (if set)\\\",\\n )\\n parser.add_argument(\\n \\\"--op-registration-whitelist\\\",\\n \\\"--op_registration_whitelist\\\",\\n nargs=\\\"*\\\",\\n help=\\\"filter op registrations by the whitelist (if set); \\\"\\n \\\"each item is `namespace`::`operator name` without overload name; \\\"\\n \\\"e.g.: aten::empty aten::conv2d ...\\\",\\n )\\n parser.add_argument(\\n \\\"--op-selection-yaml-path\\\",\\n \\\"--op_selection_yaml_path\\\",\\n help=\\\"Provide a path to the operator selection (for custom build) YAML \\\"\\n \\\"that contains the information about the set of selected operators \\\"\\n \\\"and their categories (training, ...). Each operator is either a \\\"\\n \\\"full operator name with overload or just a bare operator name. \\\"\\n \\\"The operator names also contain the namespace prefix (e.g. aten::)\\\",\\n )\\n parser.add_argument(\\n \\\"--tags-path\\\",\\n help=\\\"Path to tags.yaml. Required by yaml parsing in codegen system.\\\",\\n )\\n parser.add_argument(\\n \\\"--rocm\\\",\\n action=\\\"store_true\\\",\\n help=\\\"reinterpret CUDA as ROCm/HIP and adjust filepaths accordingly\\\",\\n )\\n parser.add_argument(\\n \\\"--use-aten-lib\\\",\\n \\\"--use_aten_lib\\\",\\n action=\\\"store_true\\\",\\n help=\\\"a boolean flag to indicate whether we use ATen kernels or not, in the future this flag will be per \\\"\\n \\\"operator\\\",\\n )\\n parser.add_argument(\\n \\\"--manual_registration\\\",\\n \\\"--manual-registration\\\",\\n action=\\\"store_true\\\",\\n help=\\\"a boolean flag to indicate whether we want to manually call\\\"\\n \\\"register_kernels() or rely on static init. \\\",\\n )\\n parser.add_argument(\\n \\\"--generate\\\",\\n type=str,\\n nargs=\\\"*\\\",\\n choices=[\\\"headers\\\", \\\"sources\\\"],\\n default=[\\\"headers\\\", \\\"sources\\\"],\\n help=\\\"Generate only a subset of files\\\",\\n )\\n options = parser.parse_args()\\n assert options.tags_path, \\\"tags.yaml is required by codegen yaml parsing.\\\"\\n\\n selector = get_custom_build_selector(\\n options.op_registration_whitelist,\\n options.op_selection_yaml_path,\\n )\\n\\n parsed_yaml, custom_ops_parsed_yaml = parse_yaml_files(\\n aten_yaml_path=options.aten_yaml_path,\\n tags_yaml_path=options.tags_path,\\n native_yaml_path=options.functions_yaml_path,\\n custom_ops_yaml_path=options.custom_ops_yaml_path,\\n selector=selector,\\n use_aten_lib=options.use_aten_lib,\\n )\\n native_functions, kernel_index = (\\n parsed_yaml.native_functions,\\n parsed_yaml.kernel_index,\\n )\\n custom_ops_native_functions = (\\n custom_ops_parsed_yaml.native_functions if custom_ops_parsed_yaml else []\\n )\\n\\n cpu_fm = make_file_manager(options=options)\\n\\n if \\\"headers\\\" in options.generate:\\n # generate CustomOpsNativeFunctions.h when custom_ops.yaml is present, to match the build system.\\n gen_headers(\\n native_functions=native_functions,\\n gen_custom_ops_header=options.custom_ops_yaml_path,\\n custom_ops_native_functions=custom_ops_native_functions,\\n selector=selector,\\n kernel_index=kernel_index,\\n cpu_fm=cpu_fm,\\n use_aten_lib=options.use_aten_lib,\\n )\\n\\n if \\\"sources\\\" in options.generate:\\n gen_unboxing(\\n native_functions=native_functions,\\n cpu_fm=cpu_fm,\\n selector=selector,\\n use_aten_lib=options.use_aten_lib,\\n kernel_index=kernel_index,\\n manual_registration=options.manual_registration,\\n )\\n if custom_ops_native_functions:\\n gen_custom_ops(\\n native_functions=custom_ops_native_functions,\\n selector=selector,\\n kernel_index=kernel_index,\\n cpu_fm=cpu_fm,\\n rocm=options.rocm,\\n )\\n\\n if options.output_dependencies:\\n depfile_path = Path(options.output_dependencies).resolve()\\n depfile_name = depfile_path.name\\n depfile_stem = depfile_path.stem\\n\\n for fm, prefix in [\\n (cpu_fm, \\\"\\\"),\\n ]:\\n varname = prefix + depfile_stem\\n path = depfile_path.parent / (prefix + depfile_name)\\n fm.write_outputs(varname, str(path))\\n\\n\\nif __name__ == \\\"__main__\\\":\\n main()\\n\\n\\nfrom __future__ import annotations\\n\\nimport contextlib\\nimport functools\\nfrom typing import Any, Callable, Iterator, List, Optional, Tuple, TypeVar, Union\\n\\nimport torchgen.local as local\\nfrom torchgen.model import (\\n BackendIndex,\\n DispatchKey,\\n NativeFunction,\\n NativeFunctionsGroup,\\n NativeFunctionsViewGroup,\\n)\\nfrom torchgen.utils import context, S, T\\n\\n\\n# Helper functions for defining generators on things in the model\\n\\nF = TypeVar(\\n \\\"F\\\",\\n NativeFunction,\\n NativeFunctionsGroup,\\n NativeFunctionsViewGroup,\\n Union[NativeFunction, NativeFunctionsGroup],\\n Union[NativeFunction, NativeFunctionsViewGroup],\\n)\\n\\nF2 = TypeVar(\\n \\\"F2\\\",\\n NativeFunction,\\n NativeFunctionsGroup,\\n Optional[NativeFunction],\\n bool,\\n str,\\n)\\n\\nF3 = TypeVar(\\\"F3\\\", Tuple[NativeFunction, Any], List[NativeFunction])\\n\\n\\n@contextlib.contextmanager\\ndef native_function_manager(\\n g: NativeFunctionsGroup | NativeFunctionsViewGroup | NativeFunction,\\n) -> Iterator[None]:\\n if isinstance(g, NativeFunctionsGroup):\\n # By default, we associate all errors with structured native functions\\n # with the out variant. In some cases, it might be better to have\\n # a more specific place to hang things; if so, use\\n # native_function_manager again on the inside\\n f = g.out\\n elif isinstance(g, NativeFunctionsViewGroup):\\n # We associate errors with the view operator\\n f = g.view\\n else:\\n f = g\\n with context(lambda: f\\\"in native_functions.yaml line {f.loc}:\\\\n {f.func}\\\"):\\n with local.parametrize(\\n use_const_ref_for_mutable_tensors=f.use_const_ref_for_mutable_tensors,\\n use_ilistref_for_tensor_lists=f.part_of_structured_group,\\n ):\\n yield\\n\\n\\n# Given a function that operates on NativeFunction, wrap it into a new function\\n# that sets some appropriate context managers for that native function.\\n# YOU MUST WRAP FUNCTIONS IN THIS for calls to api modules to be sound\\n# (you will get an error if we try to access the local variables without having\\n# set them).\\ndef with_native_function(func: Callable[[F], T]) -> Callable[[F], T]:\\n @functools.wraps(func)\\n def wrapper(f: F) -> T:\\n with native_function_manager(f):\\n return func(f)\\n\\n return wrapper\\n\\n\\ndef with_native_function_and(func: Callable[[F, F2], T]) -> Callable[[F, F2], T]:\\n @functools.wraps(func)\\n def wrapper(f: F, f2: F2) -> T:\\n # The first native_function is assumed to be the one with the appropriate context.\\n with native_function_manager(f):\\n return func(f, f2)\\n\\n return wrapper\\n\\n\\ndef method_with_native_function(func: Callable[[S, F], T]) -> Callable[[S, F], T]:\\n @functools.wraps(func)\\n def wrapper(slf: S, f: F) -> T:\\n with native_function_manager(f):\\n return func(slf, f)\\n\\n return wrapper\\n\\n\\ndef method_with_nested_native_function(\\n func: Callable[[S, F3], T]\\n) -> Callable[[S, F3], T]:\\n @functools.wraps(func)\\n def wrapper(slf: S, f: F3) -> T:\\n with native_function_manager(f[0]):\\n return func(slf, f)\\n\\n return wrapper\\n\\n\\n# Convenience decorator for functions that explicitly take in a BackendIndex,\\n# instead of indirectly taking one in as a closure\\ndef with_native_function_and_index(\\n func: Callable[[F, BackendIndex], T]\\n) -> Callable[[F, BackendIndex], T]:\\n @functools.wraps(func)\\n def wrapper(f: F, backend_index: BackendIndex) -> T:\\n with native_function_manager(f):\\n return func(f, backend_index)\\n\\n return wrapper\\n\\n\\n# Convenience decorator for functions that explicitly take in a Dict of BackendIndices\\ndef with_native_function_and_indices(\\n func: Callable[[F, dict[DispatchKey, BackendIndex]], T]\\n) -> Callable[[F, dict[DispatchKey, BackendIndex]], T]:\\n @functools.wraps(func)\\n def wrapper(f: F, backend_indices: dict[DispatchKey, BackendIndex]) -> T:\\n with native_function_manager(f):\\n return func(f, backend_indices)\\n\\n return wrapper\\n\\n\\nfrom __future__ import annotations\\n\\nfrom collections import defaultdict\\nfrom typing import Sequence\\n\\nimport torchgen.api.dispatcher as dispatcher\\nfrom torchgen.api.translate import translate\\nfrom torchgen.api.types import Binding, DispatcherSignature, Expr\\nfrom torchgen.context import with_native_function\\nfrom torchgen.model import (\\n Annotation,\\n Argument,\\n BackendIndex,\\n BackendMetadata,\\n BaseOperatorName,\\n BaseTy,\\n BaseType,\\n DEFAULT_KERNEL_NAMESPACE,\\n DeviceCheckType,\\n DispatchKey,\\n FunctionSchema,\\n NativeFunction,\\n NativeFunctionsGroup,\\n OperatorName,\\n Return,\\n SchemaKind,\\n Variant,\\n)\\nfrom torchgen.utils import concatMap\\n\\n\\n# See Note: [Out ops with functional variants that don't get grouped properly]\\nOUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY = [\\n # This has a functional variant, but it's currently marked private.\\n # This function should be marked private as well (*_backward ops aren't exposed to python anyway).\\n \\\"adaptive_avg_pool3d_backward.grad_input\\\",\\n # There's a functional variant, _slow_conv2d_backward.output_mask, that isn't grouped properly.\\n # Maybe we can kill this operator in favor of convolution_backward?\\n \\\"_slow_conv2d_backward.grad_input\\\",\\n]\\n\\n\\n# See Note: [Mutable ops that cannot get an out variant]\\nMUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT = [\\n # should be out=?\\n \\\"_cummax_helper\\\",\\n # should be out=?\\n \\\"_cummin_helper\\\",\\n]\\n\\n# All of these operators don't have any tensor like returns\\nFUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT = [\\n \\\"_assert_async\\\", # no return\\n \\\"_assert_async.msg\\\", # no return\\n \\\"_cslt_sparse_mm_search\\\", # returns an int\\n \\\"_assert_scalar\\\", # no return\\n \\\"_dimI\\\", # returns an int\\n \\\"_dimV\\\", # returns an int\\n \\\"_has_same_storage_numel\\\", # returns a boolean\\n \\\"_linalg_check_errors\\\", # no return\\n \\\"_local_scalar_dense\\\", # returns a Scalar\\n \\\"_nested_tensor_from_mask_left_aligned\\\", # returns a boolean\\n \\\"_nnz\\\", # returns an int\\n \\\"_use_cudnn_ctc_loss\\\", # returns a boolean\\n \\\"_use_cudnn_ctc_loss.Tensor\\\", # returns a boolean\\n \\\"_validate_compressed_sparse_indices\\\", # no return\\n \\\"allclose\\\", # returns a boolean\\n \\\"dense_dim\\\", # returns an int\\n \\\"equal\\\", # returns a boolean\\n \\\"is_coalesced\\\", # returns an boolean\\n \\\"is_pinned\\\", # returns a boolean\\n \\\"is_same_size\\\", # returns a boolean\\n \\\"is_set_to\\\", # returns a boolean\\n \\\"q_per_channel_axis\\\", # returns an int\\n \\\"q_scale\\\", # returns a float\\n \\\"q_zero_point\\\", # returns an int\\n \\\"qscheme\\\", # returns a QScheme\\n \\\"record_stream\\\", # no return\\n \\\"sparse_dim\\\", # returns an int\\n \\\"sym_constrain_range\\\", # no return\\n \\\"sym_constrain_range_for_size\\\", # no return\\n \\\"_nested_tensor_storage_offsets\\\", # returns a vector of ints\\n \\\"_chunk_grad_outputs_efficient_attention\\\", # returns a bool\\n \\\"_fused_sdp_choice\\\", # returns an int\\n \\\"_print\\\", # no return\\n \\\"_sink_tokens\\\", # no return\\n \\\"_nested_get_ragged_idx\\\", # returns an int\\n]\\n\\nINPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY = [\\n # polygamma and polygamma.out both exist, but have a\\n # pre-self arg (while polygamma_ does not)\\n # We should either fix this schema so it can be grouped properly,\\n # or allow the codegen to generate new functional/out= NativeFunctions for this op\\n # (which would require changing its overload name to prevent overload ambiguity).\\n \\\"polygamma_\\\"\\n]\\n\\n\\n# Groups \\\"similar\\\" NativeFunctions together\\n# example add.Tensor, add_.Tensor, add.out\\n# \\\"similar\\\" NativeFunctions are all expected to have an identical `signature()`,\\n# But have differing SchemaKinds.\\ndef pre_group_native_functions(\\n native_functions: Sequence[NativeFunction],\\n) -> dict[FunctionSchema, dict[SchemaKind, NativeFunction]]:\\n pre_grouped_native_functions: dict[\\n FunctionSchema, dict[SchemaKind, NativeFunction]\\n ] = defaultdict(dict)\\n for f in native_functions:\\n d = pre_grouped_native_functions[f.func.signature()]\\n assert f.func.kind() not in d\\n d[f.func.kind()] = f\\n return pre_grouped_native_functions\\n\\n\\n# Returns the out variant overload name given a base function overload name\\ndef get_expected_out_variant_overload_name(overload_name: str | None) -> str:\\n return \\\"out\\\" if not overload_name else f\\\"{overload_name}_out\\\"\\n\\n\\n# Helper function: given an inplace FunctionSchema, generate its corresponding out= variant\\n# Example before:\\n# _add_relu_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)\\n# Example after:\\n# _add_relu.Scalar_out(Tensor self, Scalar other, Scalar alpha=1, *, Tensor(a!) out)\\ndef self_to_out_signature(func: FunctionSchema) -> FunctionSchema:\\n # Generating an out= schema from an inplace schema.\\n assert func.kind() == SchemaKind.inplace\\n assert func.arguments.self_arg is not None\\n # The new out= schema has:\\n # - a new out argument with the same type as \\\"func\\\" (but with a mutable annotation)\\n # - The returns (if any) now alias the out= argument instead of \\\"func\\\"\\n # - an \\\"out\\\" overload name\\n return FunctionSchema(\\n name=func.name.remove_inplace().with_overload(\\n get_expected_out_variant_overload_name(func.name.overload_name)\\n ),\\n arguments=func.arguments.remove_self_annotation().with_out_args(\\n [\\n Argument(\\n name=\\\"out\\\",\\n type=func.arguments.self_arg.argument.type,\\n default=None,\\n annotation=func.arguments.self_arg.argument.annotation,\\n )\\n ]\\n ),\\n returns=func.returns,\\n )\\n\\n\\n# Helper function: given a functional FunctionSchema, generate its corresponding out= variant\\n# Example before:\\n# _to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None,\\n# bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor\\n# Example after:\\n# _to_copy._out(Tensor self, *, bool non_blocking=False, MemoryFormat? memory_format=None,\\n# Tensor(a!) out) -> Tensor(a!)\\ndef functional_to_out_signature(func: FunctionSchema) -> FunctionSchema:\\n # Generating an out= schema from a functional schema.\\n assert func.kind() == SchemaKind.functional\\n\\n new_returns, new_out_args = generate_out_args_from_schema(func)\\n # The new out= schema has:\\n # - one or more new out argument(s) with the same type as returns (but with a mutable annotation)\\n # - The returns now alias the out= arguments\\n # - an \\\"_out\\\" overload name\\n return FunctionSchema(\\n name=func.name.with_overload(\\n get_expected_out_variant_overload_name(func.name.overload_name)\\n ),\\n arguments=func.arguments.signature().with_out_args(\\n new_out_args,\\n ),\\n returns=tuple(new_returns),\\n )\\n\\n\\n# Helper function: given a function schema, generate corresponding out arguments, also the updated return annotations.\\ndef generate_out_args_from_schema(\\n func: FunctionSchema,\\n) -> tuple[list[Return], list[Argument]]:\\n # More of a sanity check - our existing restrictions on schemas should enforce that\\n # mutable schema kinds never return their mutable arguments.\\n assert not any(\\n r.annotation is not None and r.annotation.is_write for r in func.returns\\n )\\n\\n tensorlike_rets = [r for r in func.returns if r.type.is_tensor_like()]\\n assert len(tensorlike_rets) > 0\\n\\n used_annotations = concatMap(\\n lambda a: [] if a.annotation is None else a.annotation.alias_set,\\n func.arguments.flat_all,\\n )\\n valid_annotations = [\\n x for x in \\\"abcdefghijklmnopqrstuvwxyz\\\" if x not in used_annotations\\n ]\\n\\n all_rets_are_tensors = all(r.type == BaseType(BaseTy.Tensor) for r in func.returns)\\n\\n new_out_args: list[Argument] = []\\n # The end result of new_returns is that:\\n # - If every return is a plain tensor, then the new returns == the old returns, but with the out= alias annotations added.\\n # - Otherwise, none of the out arguments show up in the returns (and we're only left with non-tensor-like returns, if any).\\n new_returns: list[Return] = []\\n for i, r in enumerate(func.returns):\\n if r.type.is_tensor_like():\\n new_out = Argument(\\n name=\\\"out\\\" if len(func.returns) == 1 else f\\\"out{i}\\\",\\n type=r.type,\\n default=None,\\n annotation=Annotation.parse(f\\\"{valid_annotations[i]}!\\\"),\\n )\\n new_out_args.append(new_out)\\n if all_rets_are_tensors:\\n # The convention for out= schemas is that they only return their out arguments\\n # if the return is a plain Tensor (or if it's a tuple of plain Tensors)\\n new_ret = Return(\\n name=None, type=new_out.type, annotation=new_out.annotation\\n )\\n new_returns.append(new_ret)\\n else:\\n new_returns.append(r)\\n return new_returns, new_out_args\\n\\n\\n# Helper function: given a mutable FunctionSchema, generate its corresponding out= variant\\n# Example before:\\n# _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask) # noqa: B950\\n# Example after:\\n# _fused_moving_avg_obs_fq_helper._out(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False, *, Tensor(e!) out0, Tensor(f!) out1) -> (Tensor(e!), Tensor(f!)) # noqa: B950\\ndef mutable_to_out_signature(func: FunctionSchema) -> FunctionSchema:\\n # Generating an out= schema from a mutable schema.\\n assert func.kind() == SchemaKind.mutable\\n # The new out= schema has:\\n # - Any non-aliased tensor-like returns are converted to mutable, aliased out= arguments\\n # (if the argument is a tensor then we also return it for method chaining,\\n # otherwise we return nothing)\\n # - an \\\"out\\\" overload name\\n #\\n # Note that:\\n # (1) This also means that we can *only* generate an out= variant from a mutable schema\\n # if the mutable schema has at least one tensor-like non-aliasing return.\\n # (2) The generated out= variant still has mutable positional arguments,\\n # but if necessary we could probably add another out= variant that also\\n # functionalizes the mutable arguments (a functional_out variant)\\n\\n new_returns, new_out_args = generate_out_args_from_schema(func)\\n\\n return FunctionSchema(\\n name=func.name.remove_inplace().with_overload(\\n get_expected_out_variant_overload_name(func.name.overload_name)\\n ),\\n arguments=func.arguments.with_out_args(new_out_args),\\n returns=tuple(new_returns),\\n )\\n\\n\\n# This function, given function of one SchemaKind, as well as a target SchemaKind,\\n# generates a new NativeFunction with the same properties, but using the target SchemaKind.\\n# We only actually generate functions for either functional or out= SchemaKinds.\\n# This function returns a tuple, with:\\n# - The generated NativeFunction\\n# - a dictionary of `BackendIndex` objects, describing which dispatch keys\\n# we will generate kernels for, for the new NativeFunction.\\n# Details are in the function, but we only generate composite kernels (in some cases) today.\\ndef generate_function(\\n f: NativeFunction, k: SchemaKind\\n) -> tuple[NativeFunction, dict[DispatchKey, dict[OperatorName, BackendMetadata]]]:\\n from torchgen.api import cpp\\n\\n if k == SchemaKind.functional:\\n assert f.func.kind() != SchemaKind.functional\\n # The new \\\"functional\\\" NativeFunction has:\\n # - any mutable arguments have been converted into (immutable) returns.\\n # (if a mutable argument was not also a return, it gets converted to one)\\n # - \\\"_functional\\\" appended to the base name, ONLY IF this op has a mutable variant.\\n # See Note [Overload Ambiguity With Functional Variants]\\n # The default grouping logic in signature() actually already does this,\\n # so we can piggy-back off it (but we still want return names)\\n func = f.func.signature(keep_return_names=True).with_name(\\n OperatorName(\\n name=BaseOperatorName(\\n base=f.func.name.name.base,\\n inplace=False,\\n dunder_method=f.func.name.name.dunder_method,\\n # See Note [Overload Ambiguity With Functional Variants]\\n functional_overload=f.func.kind() == SchemaKind.mutable,\\n ),\\n overload_name=f.func.name.overload_name,\\n )\\n )\\n elif k == SchemaKind.out:\\n # We generate out= ops mostly just so that we can pair up NativeFunctions into groups easily,\\n # but at least today, there is no good reason to actually use them.\\n # we'll generate a dispatcher entry for them, but won't actually register any kernels for them.\\n if f.func.kind() == SchemaKind.inplace:\\n func = self_to_out_signature(f.func)\\n elif f.func.kind() == SchemaKind.mutable:\\n func = mutable_to_out_signature(f.func)\\n elif f.func.kind() == SchemaKind.functional:\\n func = functional_to_out_signature(f.func)\\n else:\\n raise AssertionError(\\n \\\"We only bother generating out= functions from either inplace or mutable or functional variants\\\"\\n )\\n else:\\n raise AssertionError(\\n \\\"We currently only generate either functional or out= NativeFunctions\\\"\\n )\\n\\n # Generated kernel naming convention for out: _. The reason for this is to\\n # disambiguate operator with the same name but different overload name, e.g., `randn.names_out` and\\n # `randn.generator_with_names_out`.\\n kernel_name = (\\n func.name.unambiguous_name()\\n if func.kind() == SchemaKind.out\\n else cpp.name(func)\\n )\\n if f.func.has_symint():\\n kernel_name += \\\"_symint\\\"\\n backend_metadata = {\\n DispatchKey.CompositeExplicitAutograd: {\\n func.name: BackendMetadata(\\n kernel=kernel_name,\\n structured=False,\\n cpp_namespace=DEFAULT_KERNEL_NAMESPACE,\\n )\\n }\\n }\\n tags = {\\\"generated\\\"} | set(\\n f.tags & {\\\"nondeterministic_seeded\\\", \\\"view_copy\\\", \\\"pt2_compliant_tag\\\"}\\n )\\n\\n return (\\n NativeFunction(\\n func=func,\\n use_const_ref_for_mutable_tensors=f.use_const_ref_for_mutable_tensors,\\n # These generated fn's aren't meant to be user friendly- don't generate methods.\\n variants={Variant.function},\\n structured=False,\\n structured_delegate=None,\\n structured_inherits=None,\\n precomputed=None,\\n autogen=[],\\n ufunc_inner_loop={},\\n manual_kernel_registration=False,\\n manual_cpp_binding=False,\\n python_module=None,\\n category_override=None,\\n device_guard=False,\\n device_check=DeviceCheckType.NoCheck,\\n loc=f.loc,\\n cpp_no_default_args=set(),\\n is_abstract=f.is_abstract,\\n has_composite_implicit_autograd_kernel=False,\\n has_composite_implicit_autograd_nested_tensor_kernel=False,\\n has_composite_explicit_autograd_kernel=True,\\n has_composite_explicit_autograd_non_functional_kernel=False,\\n # Every generated NativeFunction gets a \\\"generated\\\" tag, so it's easy to tell\\n # which NativeFunction objects did not come directly from native_functions.yaml.\\n tags=tags,\\n namespace=f.namespace,\\n ),\\n backend_metadata,\\n )\\n\\n\\n# This function is responsible for adding generated NativeFunctions which don't appear\\n# explicitly in the codegen.\\n# You can inspect the full list of NativeFunctions yourself with the torchgen package, by running\\n# torchgen.parse_native_yaml(\\\"aten/src/ATen/native/native_functions.yaml\\\", \\\"aten/src/ATen/native/tags.yaml\\\")\\n# (Maybe we should make a friendly API for this)\\n#\\n# Note: this function *mutates* its two inputs,\\n# adding the new NativeFunctions / BackendMetadata to them\\ndef add_generated_native_functions(\\n rs: list[NativeFunction],\\n indices: dict[DispatchKey, dict[OperatorName, BackendMetadata]],\\n) -> None:\\n # The main code for generating new NativeFunctions\\n # First we group of NativeFunctions by schema kind,\\n # then we detect which ones are missing and generate them.\\n pre_grouped_native_functions = pre_group_native_functions(rs)\\n for d in pre_grouped_native_functions.values():\\n has_functional = SchemaKind.functional in d\\n has_inplace = SchemaKind.inplace in d\\n has_mutable = SchemaKind.mutable in d\\n has_out = SchemaKind.out in d\\n\\n # We automatically generate a few native functions that don't exist in the yaml, for a few reasons:\\n # (1) If an operator has an inplace/out= variant but no functional variant, we can generate\\n # a simple functional variant that the functionalization pass can consume.\\n # (2) If an operator has an inplace or functional but no out= variant, we generate an out=\\n # variant, mostly so we can easily pair up functions into NativeFunctionsGroup,\\n # while maintaining the constraint that the out= variant is \\\"required\\\".\\n if has_mutable or has_inplace or has_out or has_functional:\\n # Don't bother generating functions trio's for native functions that bypass the dispatcher.\\n are_manual = all(f.manual_cpp_binding for f in d.values())\\n # Don't bother generating functional + out= variants for view operators\\n # set_ is technically an inplace_view, but for now it is treated\\n # as a normal inplace op in the codegen\\n has_view_ops = any(\\n f.is_view_op and str(f.func.name.name) != \\\"set_\\\" for f in d.values()\\n )\\n # Don't generate the other variants for CompositeImplicitAutograd operators.\\n # We could probably do this, but the main benefit of generating the function triplets\\n # is for transforms that need them, and transforms don't need to act directly\\n # on CompositeImplicitAutograd operators (since we let them decompose).\\n are_composite_implicit = all(\\n f.has_composite_implicit_autograd_kernel for f in d.values()\\n )\\n if are_manual or has_view_ops or are_composite_implicit:\\n continue\\n if has_out and len(d.values()) == 1:\\n # Note: [Out ops with functional variants that don't get grouped properly]\\n # In theory we could validly have an out= operator in native_functions.yaml\\n # that has no other variants.\\n # But today, all of the operators where that's the case actually do have\\n # functional variants, that we are just unable to pair up properly.\\n # I think banning this all together is probably safer\\n # (you can always add a functional variant yourself if you want to add a new out= operator).\\n #\\n # We should probably fix the existing cases; this check is to prevent us from adding more over time.\\n if (\\n str(d[SchemaKind.out].func.name)\\n not in OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY\\n ):\\n raise AssertionError(\\n f\\\"Found an out= operator that we could not find any other variants of: {str(d[SchemaKind.out].func)}\\\"\\n )\\n continue\\n\\n # Some inplace ops that have problematic schemas (that we should fix), which prevent us\\n # from generating out= and functional variants\\n if (\\n has_inplace\\n and str(d[SchemaKind.inplace].func.name)\\n in INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY\\n ):\\n continue\\n\\n base_fn = (\\n d[SchemaKind.inplace]\\n if has_inplace\\n else d[SchemaKind.mutable]\\n if has_mutable\\n else d[SchemaKind.out]\\n if has_out\\n else d[SchemaKind.functional]\\n )\\n\\n # Note: [Mutable ops that cannot get an out variant]\\n # We can only generate an out= variant if either:\\n # - the original function has tensor-like returns (since we can convert them to out kwargs)\\n # - or it's inplace (since we can convert `self` to an out kwarg)\\n # There are only two functions that don't fit this criteria today though,\\n # and they both look like they should be fixed to be out= variants,\\n # so if feels safer to ban this schema all-together\\n base_fn_valid = base_fn.func.kind() == SchemaKind.inplace or any(\\n r.type.is_tensor_like() for r in base_fn.func.returns\\n )\\n # Note: [Loosen the assertion that all functional should have out variant]\\n # By design all functional operators should have our variants. The needs_out check\\n # is loosening this requirement, changing it to only generate out variant if there's\\n # an `autogen` block in the native function, in the long run it should be removed.\\n # FIXME: Remove this after figuring out CI job failures related to min, max, mean\\n needs_out = any(\\\"out\\\" in str(op_name) for op_name in base_fn.autogen)\\n gets_out_variant = not has_out and base_fn_valid and needs_out\\n if not has_out and not base_fn_valid:\\n if (\\n str(base_fn.func.name)\\n not in MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT\\n and str(base_fn.func.name)\\n not in FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT\\n ):\\n raise AssertionError(\\n f\\\"\\\"\\\"Found an operator that we could not generate an out= variant for: {str(base_fn.func)}.\\nThis type of operators don't have tensor-like return, making it difficult to generate a proper out= variant. If\\nout= variant is not needed, please add the function name into FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT list.\\\"\\\"\\\"\\n )\\n\\n # Generate an out= variant\\n if gets_out_variant:\\n fn, metadata = generate_function(base_fn, SchemaKind.out)\\n d[SchemaKind.out] = fn\\n BackendIndex.grow_index(indices, metadata)\\n rs.append(fn)\\n\\n # Generate a functional variant, but only do it if the operator got an out= variant\\n # (Functional variants are only useful if we can group up the variants,\\n # which we can only do if they have an out= variant)\\n if not has_functional and (has_out or gets_out_variant):\\n fn, metadata = generate_function(base_fn, SchemaKind.functional)\\n d[SchemaKind.functional] = fn\\n BackendIndex.grow_index(indices, metadata)\\n rs.append(fn)\\n\\n\\ndef return_str(rets: tuple[Return, ...], names: list[str]) -> str:\\n assert len(rets) == len(names)\\n if len(rets) == 0:\\n return \\\"\\\"\\n elif len(rets) == 1:\\n return f\\\"return {names[0]};\\\"\\n else:\\n return f\\\"return {dispatcher.returns_type(rets).cpp_type()}({', '.join(names)});\\\"\\n\\n\\n# Given a function, and the name of a variable corresponding to the output of that function,\\n# gather up all of the individual returns that are not aliased\\ndef gather_nonaliased_inner_rets(func: FunctionSchema, out_var: str) -> list[str]:\\n aliased_rets = func.aliased_return_names()\\n non_aliased_names = []\\n is_out_var_a_tuple = len(func.returns) > 1\\n for i, r in enumerate(aliased_rets):\\n if r is None:\\n non_aliased_names.append(\\n f\\\"std::get<{i}>({out_var})\\\" if is_out_var_a_tuple else out_var\\n )\\n return non_aliased_names\\n\\n\\n# Generates functional kernels in terms of their inplace.mutable counterparts.\\n# We only do this for \\\"generated\\\" NativeFunctions\\n@with_native_function\\ndef gen_composite_functional_kernel(g: NativeFunctionsGroup) -> str | None:\\n # We should only be generating these for code-generated NativeFunctions\\n if \\\"generated\\\" not in g.functional.tags:\\n return None\\n # And we always write the kernel for a generated op in terms of a non-generated op.\\n if g.inplace is not None and \\\"generated\\\" not in g.inplace.tags:\\n target_f = g.inplace\\n elif g.mutable is not None and \\\"generated\\\" not in g.mutable.tags:\\n target_f = g.mutable\\n else:\\n # We should be guaranteed to have a valid inplace/mutable variant to call into.\\n # See Note: [Mutable Ops Not Using Functionalization]\\n raise AssertionError(str(g.functional.func))\\n\\n sig = DispatcherSignature(g.functional.func)\\n target_sig = DispatcherSignature(target_f.func)\\n\\n context: list[Binding | Expr] = []\\n clone_mutable_inputs = []\\n cloned_return_names = []\\n # We can't just directly pass all of the arguments from the functional op into the mutating op.\\n # We need to check for which inputs to the mutating operator are mutable,\\n # and clone those inputs first.\\n for a_curr, a_tgt in zip(\\n dispatcher.jit_arguments(g.functional.func),\\n dispatcher.jit_arguments(target_f.func),\\n ):\\n if a_tgt.annotation is not None and a_tgt.annotation.is_write:\\n clone_mutable_inputs.append(\\n f\\\"auto {a_curr.name}_clone = clone_arg({a_curr.name});\\\"\\n )\\n context.append(\\n Expr(\\n expr=f\\\"{a_curr.name}_clone\\\",\\n type=dispatcher.argument_type(a_curr, binds=a_curr.name),\\n )\\n )\\n # Invariant: mutable arguments on the inner mutable op are always returns on the functional op.\\n cloned_return_names.append(f\\\"{a_curr.name}_clone\\\")\\n else:\\n context.append(dispatcher.argument(a_curr))\\n exprs = \\\", \\\".join([e.expr for e in translate(context, target_sig.arguments())])\\n\\n out_name = \\\"output\\\"\\n maybe_assign = f\\\"auto {out_name} = \\\" if len(target_f.func.returns) > 0 else \\\"\\\"\\n inner_return_names = gather_nonaliased_inner_rets(target_f.func, out_name)\\n ret_str = return_str(\\n g.functional.func.returns, inner_return_names + cloned_return_names\\n )\\n\\n clone_mutable_inputs_str = \\\"\\\\n\\\".join(clone_mutable_inputs)\\n return f\\\"\\\"\\\"\\n{sig.defn(name=sig.name() + (\\\"_symint\\\" if g.out.func.has_symint() else \\\"\\\"))} {{\\n {clone_mutable_inputs_str}\\n {maybe_assign}at::_ops::{target_f.func.name.unambiguous_name()}::call({exprs});\\n {ret_str}\\n}}\\n\\\"\\\"\\\"\\n\\n\\n# Generates out= kernels in terms of their functional counterparts.\\n# We only do this for \\\"generated\\\" NativeFunctions\\n@with_native_function\\ndef gen_composite_out_kernel(g: NativeFunctionsGroup) -> str | None:\\n # We should only be generating these for code-generated NativeFunctions\\n if \\\"generated\\\" not in g.out.tags:\\n return None\\n # And we always write the kernel for the out= op in terms of the functional.\\n # Note that the functional op might have also been generated, but we don't have to\\n # worry about cycles, because the generated functional kernels are always implemented\\n # in terms of non-generated kernels (see gen_composite_functional_kernel).\\n\\n sig = DispatcherSignature(g.out.func)\\n target_sig = DispatcherSignature(g.functional.func)\\n\\n exprs = \\\", \\\".join(\\n [e.expr for e in translate(sig.arguments(), target_sig.arguments())]\\n )\\n\\n copy_outs = []\\n out_name = \\\"tmp_output\\\"\\n for i, out_arg in enumerate(g.out.func.arguments.out):\\n functional_return_name = (\\n out_name\\n if len(g.functional.func.returns) == 1\\n else f\\\"std::get<{i}>({out_name})\\\"\\n )\\n copy_outs.append(\\n f\\\"\\\"\\\"\\\\\\n resize_out_helper({out_arg.name}, {functional_return_name});\\n copy_arg({out_arg.name}, {functional_return_name});\\\"\\\"\\\"\\n )\\n\\n rets = []\\n # For each return arg in the calling (out=) operator,\\n # If it corresponds to an aliased input, return the input.\\n # Otherwise, return the corresponding output from calling the functional operator.\\n for i, ret_name in enumerate(g.out.func.aliased_return_names()):\\n if ret_name is not None:\\n rets.append(ret_name)\\n else:\\n functional_return_name = (\\n out_name\\n if len(g.functional.func.returns) == 1\\n else f\\\"std::get<{i}>({out_name})\\\"\\n )\\n rets.append(functional_return_name)\\n\\n copy_outs_str = \\\"\\\\n\\\".join(copy_outs)\\n\\n # Kernel name needs to follow the naming convention defined in `generate_function()`\\n return f\\\"\\\"\\\"\\n{sig.defn(name=g.out.func.name.unambiguous_name() + (\\\"_symint\\\" if g.out.func.has_symint() else \\\"\\\"))} {{\\n auto {out_name} = at::_ops::{g.functional.func.name.unambiguous_name()}::call({exprs});\\n {copy_outs_str}\\n {return_str(g.out.func.returns, rets)}\\n}}\\n\\\"\\\"\\\"\\n\\n\\n# Safely load fast C Yaml loader/dumper if they are available\\ntry:\\n from yaml import CSafeLoader as Loader\\nexcept ImportError:\\n from yaml import SafeLoader as Loader # type: ignore[assignment, misc]\\n\\ntry:\\n from yaml import CSafeDumper as Dumper\\nexcept ImportError:\\n from yaml import SafeDumper as Dumper # type: ignore[assignment, misc]\\nYamlDumper = Dumper\\n\\n\\n# A custom loader for YAML that errors on duplicate keys.\\n# This doesn't happen by default: see https://github.com/yaml/pyyaml/issues/165\\nclass YamlLoader(Loader):\\n def construct_mapping(self, node, deep=False): # type: ignore[no-untyped-def]\\n mapping = []\\n for key_node, value_node in node.value:\\n key = self.construct_object(key_node, deep=deep) # type: ignore[no-untyped-call]\\n assert (\\n key not in mapping\\n ), f\\\"Found a duplicate key in the yaml. key={key}, line={node.start_mark.line}\\\"\\n mapping.append(key)\\n mapping = super().construct_mapping(node, deep=deep) # type: ignore[no-untyped-call]\\n return mapping\\n\\n\\n\\\"\\\"\\\"torchgen\\n\\nThis module contains codegeneration utilities for PyTorch. It is used to\\nbuild PyTorch from source, but may also be used for out-of-tree projects\\nthat extend PyTorch.\\n\\nNote well that we provide no BC guarantees for torchgen. If you're interested\\nin using torchgen and want the PyTorch team to be aware, please reach out\\non GitHub.\\n\\\"\\\"\\\"\\n\\n\\n# Represents all kernels used by an Executorch model.\\n# It maintains a Dict[OperatorName, Dict[ETKernelKey, BackendMetadata]] structure.\\n\\nfrom __future__ import annotations\\n\\nimport itertools\\nfrom collections import defaultdict, namedtuple\\nfrom dataclasses import dataclass\\nfrom enum import IntEnum\\n\\nfrom torchgen.model import (\\n BackendIndex,\\n BackendMetadata,\\n DispatchKey,\\n NativeFunction,\\n NativeFunctionsGroup,\\n OperatorName,\\n)\\nfrom torchgen.utils import assert_never\\n\\n\\nKERNEL_KEY_VERSION = 1\\n\\n\\n# TODO: Duplicated Subset from codegen.tool.gen_oplist, remove declaration in codegen\\nclass ScalarType(IntEnum):\\n Byte = 0\\n Char = 1\\n Short = 2\\n Int = 3\\n Long = 4\\n Float = 6\\n Double = 7\\n Bool = 11\\n\\n\\nETParsedYaml = namedtuple(\\\"ETParsedYaml\\\", [\\\"native_functions\\\", \\\"kernel_index\\\"])\\n\\n\\n@dataclass(frozen=True)\\nclass ETKernelKeyOpArgMeta:\\n arg_name: str\\n dtype: str\\n # The order of the dimensions if entry is a Tensor\\n dim_order: tuple[int, ...]\\n\\n def to_native_string(self) -> str:\\n dtype_str = ScalarType[self.dtype].value\\n dim_str = str(self.dim_order)[1:-1].replace(\\\" \\\", \\\"\\\")\\n return f\\\"{dtype_str};{dim_str}\\\"\\n\\n\\n@dataclass(frozen=True)\\nclass ETKernelKey:\\n # Field undefined is default = True\\n arg_meta: tuple[ETKernelKeyOpArgMeta, ...] = ()\\n\\n # Indicator for this kernel being used as a catch all\\n default: bool = False\\n\\n version: int = KERNEL_KEY_VERSION\\n\\n @staticmethod\\n def gen_from_yaml(\\n args: dict[str, tuple[str, str]],\\n type_alias_map: dict[str, list[str]], # TODO: Support unwrapped str val\\n dim_order_alias_map: dict[str, list[int]],\\n ) -> list[ETKernelKey]:\\n \\\"\\\"\\\"Generate ETKernelKeys from arg kernel specs\\n Multiple ETKernelKeys are returned due to dtype permutations from utilizing\\n type_alias_map (actualizing each potential type permutation as a KernelKey)\\n\\n Args:\\n args: Mapping from argument name to kernel specs\\n Kernel specs are a tuple of (dtype, dim_order).\\n Currently tuple entries must be aliased via the alias map arguments\\n type_alias_map: Mapping from type alias to potential type enums\\n i.e { T0 : [Double, Int] } means T0 can be either Double or Int\\n Used for lookup by args\\n dim_order_alias_map: Mapping from alias to a list of dimension orders\\n Used for lookup by args\\n \\\"\\\"\\\"\\n # Cast to dim order to int\\n dim_order_alias_map = {\\n k: [int(alias) for alias in v] for k, v in dim_order_alias_map.items()\\n }\\n kernel_keys = []\\n\\n # Get all used Dtype Alias\\n dtype_alias_used = set()\\n for type_alias, dim_order in args.values():\\n # Enforce usage of alias initially\\n # TODO: Support inlined arguments\\n assert type_alias in type_alias_map, \\\"Undefined type alias: \\\" + str(\\n type_alias\\n )\\n assert (\\n dim_order in dim_order_alias_map\\n ), \\\"Undefined dim_order alias: \\\" + str(dim_order)\\n dtype_alias_used.add(type_alias)\\n\\n # Generate all permutations of dtype alias values\\n alias_dtypes = [\\n [(alias, dtype) for dtype in type_alias_map[alias]]\\n for alias in dtype_alias_used\\n ]\\n alias_permutations = [\\n dict(permutation) for permutation in list(itertools.product(*alias_dtypes))\\n ]\\n\\n # Using each alias value permutation, generate kernel keys\\n op_arg_cache = {}\\n for permutation in alias_permutations:\\n arg_list = []\\n for arg_name, arg_spec in args.items():\\n dtype = permutation[arg_spec[0]]\\n dim_order = dim_order_alias_map[arg_spec[1]] # type: ignore[assignment]\\n if (\\n cache_key := (arg_name, dtype, tuple(dim_order))\\n ) not in op_arg_cache:\\n op_arg_cache[cache_key] = ETKernelKeyOpArgMeta(*cache_key) # type: ignore[arg-type]\\n\\n arg_list.append(op_arg_cache[cache_key])\\n kernel_keys.append(ETKernelKey(tuple(arg_list)))\\n\\n return kernel_keys\\n\\n def to_native_string(self) -> str:\\n if self.default:\\n return \\\"default\\\"\\n return (\\n \\\"v\\\"\\n + str(KERNEL_KEY_VERSION)\\n + \\\"/\\\"\\n + \\\"|\\\".join([arg.to_native_string() for arg in self.arg_meta])\\n )\\n\\n\\n@dataclass(frozen=True)\\nclass ETKernelIndex:\\n index: dict[OperatorName, dict[ETKernelKey, BackendMetadata]]\\n\\n def has_kernels(self, g: NativeFunction | NativeFunctionsGroup) -> bool:\\n m = self.get_kernels(g)\\n return m is not None\\n\\n def get_kernels(\\n self, g: NativeFunction | NativeFunctionsGroup\\n ) -> dict[ETKernelKey, BackendMetadata]:\\n if isinstance(g, NativeFunction):\\n f = g\\n elif isinstance(g, NativeFunctionsGroup):\\n f = g.functional\\n else:\\n assert_never(g)\\n if f.func.name not in self.index:\\n return {}\\n return self.index[f.func.name]\\n\\n @staticmethod\\n def grow_from_backend_indices(\\n kernel_index: dict[OperatorName, dict[ETKernelKey, BackendMetadata]],\\n backend_indices: dict[DispatchKey, dict[OperatorName, BackendMetadata]],\\n ) -> None:\\n for dk in backend_indices:\\n index = backend_indices[dk]\\n for op, backend_metadata in index.items():\\n if op in kernel_index:\\n kernel_index[op][ETKernelKey(default=True)] = backend_metadata\\n else:\\n kernel_index[op] = {ETKernelKey(default=True): backend_metadata}\\n\\n @staticmethod\\n def from_backend_indices(\\n backend_indices: dict[DispatchKey, dict[OperatorName, BackendMetadata]]\\n ) -> ETKernelIndex:\\n kernel_index: dict[\\n OperatorName, dict[ETKernelKey, BackendMetadata]\\n ] = defaultdict(dict)\\n ETKernelIndex.grow_from_backend_indices(kernel_index, backend_indices)\\n return ETKernelIndex(kernel_index)\\n\\n def grow(\\n self, backend_indices: dict[DispatchKey, dict[OperatorName, BackendMetadata]]\\n ) -> ETKernelIndex:\\n ETKernelIndex.grow_from_backend_indices(self.index, backend_indices)\\n return self\\n\\n def _to_backend_index(self) -> BackendIndex:\\n \\\"\\\"\\\"\\n WARNING: this will be deprecated once all the codegen places know how to handle ETKernelIndex.\\n \\\"\\\"\\\"\\n index: dict[OperatorName, BackendMetadata] = {}\\n for op in self.index:\\n kernel_dict = self.index[op]\\n assert (\\n len(kernel_dict.values()) == 1\\n ), f\\\"Can't convert ETKernelIndex to BackendIndex because {op} has more than one kernels. Got {kernel_dict}\\\"\\n index[op] = kernel_dict.get(\\n ETKernelKey(default=True),\\n BackendMetadata(kernel=\\\"\\\", structured=False, cpp_namespace=\\\"\\\"),\\n )\\n return BackendIndex(\\n dispatch_key=DispatchKey.CPU,\\n use_out_as_primary=False,\\n device_guard=False,\\n external=False,\\n index=index,\\n )\\n\\n # Note duplicate ETKernelKey from index_b will clobber the metadata from index_a\\n @staticmethod\\n def merge_indices(index_a: ETKernelIndex, index_b: ETKernelIndex) -> ETKernelIndex:\\n combined = defaultdict(dict, index_a.index.copy())\\n\\n for op, entry in index_b.index.items():\\n for key, metadata in entry.items():\\n combined[op][key] = metadata\\n\\n return ETKernelIndex(combined)\\n\\n\\nfrom __future__ import annotations\\n\\nfrom collections import defaultdict, namedtuple\\nfrom typing import Any\\n\\nimport yaml\\n\\nfrom torchgen.executorch.model import ETKernelIndex, ETKernelKey\\nfrom torchgen.gen import LineLoader, parse_native_yaml\\nfrom torchgen.model import (\\n BackendMetadata,\\n DispatchKey,\\n FunctionSchema,\\n NativeFunction,\\n OperatorName,\\n)\\nfrom torchgen.utils import NamespaceHelper\\n\\n\\n# Parse native_functions.yaml into a sequence of NativeFunctions and ET Backend Indices.\\nETParsedYaml = namedtuple(\\\"ETParsedYaml\\\", [\\\"native_functions\\\", \\\"et_kernel_indices\\\"])\\n\\n# Fields in native_functions.yaml used to determine which kernels should be used\\nET_FIELDS = [\\\"kernels\\\", \\\"type_alias\\\", \\\"dim_order_alias\\\"]\\n\\n\\ndef parse_from_yaml(ei: dict[str, object]) -> dict[ETKernelKey, BackendMetadata]:\\n \\\"\\\"\\\"Given a loaded yaml representing kernel assignment information, extract the\\n mapping from `kernel keys` to `BackendMetadata` (the latter representing the kernel instance)\\n\\n Args:\\n ei: Dict keys {kernels, type_alias, dim_order_alias}\\n See ETKernelKey for description of arguments\\n \\\"\\\"\\\"\\n e = ei.copy()\\n if (kernels := e.pop(\\\"kernels\\\", None)) is None:\\n return {}\\n\\n type_alias: dict[str, list[str]] = e.pop(\\\"type_alias\\\", {}) # type: ignore[assignment]\\n dim_order_alias: dict[str, list[str]] = e.pop(\\\"dim_order_alias\\\", {}) # type: ignore[assignment]\\n dim_order_alias.pop(\\\"__line__\\\", None)\\n\\n kernel_mapping: dict[ETKernelKey, BackendMetadata] = {}\\n\\n for entry in kernels: # type: ignore[attr-defined]\\n arg_meta = entry.get(\\\"arg_meta\\\")\\n if arg_meta is not None:\\n arg_meta.pop(\\\"__line__\\\")\\n\\n kernel_name = entry.get(\\\"kernel_name\\\")\\n namespace_helper = NamespaceHelper.from_namespaced_entity(\\n kernel_name, max_level=3\\n )\\n kernel_namespace = namespace_helper.get_cpp_namespace(default=\\\"at\\\")\\n backend_metadata = BackendMetadata(\\n kernel=namespace_helper.entity_name,\\n structured=False,\\n cpp_namespace=(kernel_namespace + \\\"::native\\\"),\\n )\\n\\n kernel_keys = (\\n [ETKernelKey((), default=True)]\\n if arg_meta is None\\n else ETKernelKey.gen_from_yaml(arg_meta, type_alias, dim_order_alias) # type: ignore[arg-type]\\n )\\n\\n for kernel_key in kernel_keys:\\n assert kernel_key not in kernel_mapping, (\\n \\\"Duplicate kernel key: \\\" + str(kernel_key) + \\\" \\\" + str(e)\\n )\\n kernel_mapping[kernel_key] = backend_metadata\\n\\n return kernel_mapping\\n\\n\\ndef parse_et_yaml_struct(es: object) -> ETKernelIndex:\\n \\\"\\\"\\\"Given a loaded yaml representing a list of operators, for each op extract the mapping\\n of `kernel keys` to `BackendMetadata` (the latter representing the kernel instance\\n that should be used by the kernel key).\\n \\\"\\\"\\\"\\n indices: dict[OperatorName, dict[ETKernelKey, BackendMetadata]] = {}\\n for ei in es: # type: ignore[attr-defined]\\n e = ei.copy()\\n\\n funcs = e.pop(\\\"func\\\")\\n assert isinstance(funcs, str), f\\\"not a str: {funcs}\\\"\\n namespace_helper = NamespaceHelper.from_namespaced_entity(\\n namespaced_entity=funcs, max_level=1\\n )\\n opname = FunctionSchema.parse(namespace_helper.entity_name).name\\n\\n assert opname not in indices, f\\\"Duplicate func found in yaml: {opname} already\\\"\\n\\n if len(index := parse_from_yaml(e)) != 0:\\n indices[opname] = index\\n\\n return ETKernelIndex(indices)\\n\\n\\ndef extract_kernel_fields(es: object) -> dict[OperatorName, dict[str, Any]]:\\n \\\"\\\"\\\"Given a loaded yaml representing a list of operators, extract the\\n kernel key related fields indexed by the operator name.\\n \\\"\\\"\\\"\\n fields: dict[OperatorName, dict[str, Any]] = defaultdict(dict)\\n for ei in es: # type: ignore[attr-defined]\\n funcs = ei.get(\\\"func\\\")\\n assert isinstance(funcs, str), f\\\"not a str: {funcs}\\\"\\n namespace_helper = NamespaceHelper.from_namespaced_entity(\\n namespaced_entity=funcs, max_level=1\\n )\\n opname = FunctionSchema.parse(namespace_helper.entity_name).name\\n\\n for field in ET_FIELDS:\\n if (value := ei.get(field)) is not None:\\n fields[opname][field] = value\\n\\n return fields\\n\\n\\ndef parse_et_yaml(\\n path: str,\\n tags_yaml_path: str,\\n ignore_keys: set[DispatchKey] | None = None,\\n skip_native_fns_gen: bool = False,\\n) -> tuple[list[NativeFunction], dict[OperatorName, dict[str, Any]]]:\\n \\\"\\\"\\\"Parse native_functions.yaml into NativeFunctions and an Operator Indexed Dict\\n of fields to persist from native_functions.yaml to functions.yaml\\n \\\"\\\"\\\"\\n with open(path) as f:\\n es = yaml.load(f, Loader=LineLoader)\\n\\n et_kernel = extract_kernel_fields(es)\\n\\n # Remove ET specific fields from entries for BC compatibility\\n strip_et_fields(es)\\n\\n native_yaml = parse_native_yaml(\\n path,\\n tags_yaml_path,\\n ignore_keys,\\n skip_native_fns_gen=skip_native_fns_gen,\\n loaded_yaml=es,\\n )\\n return native_yaml.native_functions, et_kernel\\n\\n\\ndef strip_et_fields(es: object) -> None:\\n \\\"\\\"\\\"Given a loaded yaml representing a list of operators,\\n remove ET specific fields from every entries for BC compatibility\\n \\\"\\\"\\\"\\n for entry in es: # type: ignore[attr-defined]\\n for field in ET_FIELDS:\\n entry.pop(field, None)\\n\\n\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\nfrom typing import Callable, Sequence, TYPE_CHECKING\\n\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n ListType,\\n NativeFunction,\\n OptionalType,\\n Type,\\n)\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.api.types import Binding, CType, NamedCType\\n\\n\\nconnector = \\\"\\\\n\\\\t\\\"\\n\\n\\n# Return unboxing function name for a NativeFunction\\ndef name(f: NativeFunction) -> str:\\n return f.func.name.unambiguous_name()\\n\\n\\n@dataclass(frozen=True)\\nclass Unboxing:\\n \\\"\\\"\\\"\\n Takes a sequence of Bindings and unbox EValues to these Bindings. Return generated code that performs correct unboxing.\\n A sample generated code:\\n // aten::mul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)\\n void mul_out(EValue** stack) {\\n EValue& self = *stack[0];\\n EValue& other = *stack[1];\\n EValue& out = *stack[2];\\n const torch::executor::Tensor & self_base = self.to();\\n const torch::executor::Tensor & other_base = other.to();\\n torch::executor::Tensor & out_base = out.to();\\n\\n EXECUTORCH_SCOPE_PROF(\\\"native_call_mul.out\\\");\\n torch::executor::mul_outf(self_base, other_base, out_base);\\n\\n\\n }\\n \\\"\\\"\\\"\\n\\n # this is a callable that converts a JIT argument, into its C++ type.\\n # Translates (type, mutability, binds) to NamedCType. E.g., torchgen.api.cpp.argumenttype_type.\\n argument_type_gen: Callable[\\n ...,\\n NamedCType,\\n ]\\n\\n # Convert all the arguments in a NativeFunction to C++ code\\n def convert_arguments(\\n self, args: Sequence[Binding]\\n ) -> tuple[list[Binding], list[str]]:\\n code_list = [f\\\"EValue& {args[i].name} = *stack[{i}];\\\" for i in range(len(args))]\\n binding_list = []\\n for arg in args:\\n # expecting only Argument\\n if not isinstance(arg.argument, Argument):\\n raise Exception( # noqa: TRY002\\n f\\\"Unexpected argument type, expecting `Argument` but got {arg}\\\"\\n )\\n argument: Argument = arg.argument\\n unboxed_name, _, code, decl = self.argumenttype_evalue_convert(\\n argument.type, argument.name, mutable=argument.is_write\\n )\\n code_list.extend(decl)\\n code_list.extend(code)\\n binding_list.append(arg.with_name(unboxed_name))\\n return binding_list, code_list\\n\\n def argumenttype_evalue_convert(\\n self, t: Type, arg_name: str, *, mutable: bool = False\\n ) -> tuple[str, CType, list[str], list[str]]:\\n \\\"\\\"\\\"\\n Takes in the type, name and mutability corresponding to an argument, and generates a tuple of:\\n (1) the C++ code necessary to unbox the argument\\n (2) A Binding corresponding to the newly created unboxed variable, including variable name and its CType\\n :param t: a `Type` of an argument\\n :param arg_name: argument name\\n :param mutable: boolean for whether this argument type is mutable\\n :return: unboxed result\\n \\\"\\\"\\\"\\n ctype = self.argument_type_gen(t, mutable=mutable, binds=arg_name).type\\n\\n if isinstance(t, BaseType):\\n out_name = f\\\"{arg_name}_base\\\"\\n code, decl = self._gen_code_base_type(\\n arg_name=arg_name, out_name=out_name, ctype=ctype\\n )\\n elif isinstance(t, OptionalType):\\n out_name = f\\\"{arg_name}_opt_out\\\"\\n code, decl = self._gen_code_optional_type(\\n arg_name=arg_name, out_name=out_name, t=t, ctype=ctype\\n )\\n elif isinstance(t, ListType):\\n out_name = f\\\"{arg_name}_list_out\\\"\\n code, decl = self._gen_code_list_type(\\n arg_name=arg_name, out_name=out_name, t=t, ctype=ctype\\n )\\n else:\\n raise Exception( # noqa: TRY002\\n f\\\"Cannot handle type {t}. arg_name: {arg_name}\\\"\\n ) # noqa: TRY002\\n return out_name, ctype, code, decl\\n\\n def _gen_code_base_type(\\n self, arg_name: str, out_name: str, ctype: CType\\n ) -> tuple[list[str], list[str]]:\\n return [\\n f\\\"{ctype.cpp_type()} {out_name} = {arg_name}.to<{ctype.cpp_type(strip_ref=True)}>();\\\"\\n ], []\\n\\n def _gen_code_optional_type(\\n self, arg_name: str, out_name: str, t: OptionalType, ctype: CType\\n ) -> tuple[list[str], list[str]]:\\n in_name = f\\\"{arg_name}_opt_in\\\"\\n res_name, base_type, res_code, decl = self.argumenttype_evalue_convert(\\n t.elem, in_name\\n )\\n return (\\n f\\\"\\\"\\\"\\n auto {out_name} = {arg_name}.toOptional<{base_type.cpp_type(strip_ref=True)}>();\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n ),\\n decl,\\n )\\n\\n def _gen_code_list_type(\\n self, arg_name: str, out_name: str, t: ListType, ctype: CType\\n ) -> tuple[list[str], list[str]]:\\n in_name = f\\\"{arg_name}_list_in\\\"\\n elem_name = f\\\"{arg_name}_elem\\\"\\n code = []\\n res_name, res_ctype, res_code, decl = self.argumenttype_evalue_convert(\\n t.elem, elem_name\\n )\\n\\n if isinstance(t.elem, BaseType) and t.elem.name == BaseTy.Tensor:\\n code.extend(\\n f\\\"\\\"\\\"\\n auto {out_name} = {arg_name}.toTensorList();\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n elif isinstance(t.elem, BaseType) and (\\n t.elem.name == BaseTy.int or t.elem.name == BaseTy.SymInt\\n ):\\n code.extend(\\n f\\\"\\\"\\\"\\n auto {out_name} = {arg_name}.toIntList();\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n elif isinstance(t.elem, BaseType) and t.elem.name == BaseTy.float:\\n code.extend(\\n f\\\"\\\"\\\"\\n auto {out_name} = {arg_name}.toDoubleList();\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n elif isinstance(t.elem, BaseType) and t.elem.name == BaseTy.bool:\\n # handle list type with size, e.g., bool[4]\\n code.extend(\\n f\\\"\\\"\\\"\\n#ifdef USE_ATEN_LIB\\nstd::array {out_name};\\nauto {in_name} = {arg_name}.toBoolList();\\nsize_t _i = 0;\\nfor (auto {elem_name}: {in_name}) {{\\n {out_name}[_i++] = {elem_name};\\n}}\\n#else\\nauto {out_name} = {arg_name}.toBoolList();\\n#endif\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n # pytorch codegen:\\n # we have to use c10::List for optional element. e.g., Tensor?[] -> c10::List<::std::optional>\\n elif (\\n isinstance(t.elem, OptionalType)\\n and isinstance(t.elem.elem, BaseType)\\n and t.elem.elem.name == BaseTy.Tensor\\n ):\\n code.extend(\\n f\\\"\\\"\\\"\\n#ifdef USE_ATEN_LIB\\nauto {in_name} = {arg_name}.toListOptionalTensor();\\nc10::List<::std::optional> {out_name};\\nfor (auto {elem_name}: {in_name}) {{\\n {out_name}.push_back({elem_name});\\n}}\\n#else\\nauto {out_name} = {arg_name}.toListOptionalTensor();\\n#endif\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n else:\\n # use ArrayRef as default.\\n vec_name = arg_name + \\\"_vec\\\"\\n # need to bring vector instantiation out of scope so that ArrayRef has valid data\\n decl.append(\\n f\\\"std::vector<{res_ctype.cpp_type(strip_ref=True)}> {vec_name};\\\"\\n )\\n code.extend(\\n f\\\"\\\"\\\"\\n for (EValue {elem_name}: {in_name}) {{\\n {connector.join(res_code)}\\n {vec_name}.push_back({res_name});\\n }}\\n {ctype.cpp_type(strip_ref=True)} {out_name}({vec_name});\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n return code, decl\\n\\n\\nfrom __future__ import annotations\\n\\nfrom typing import Sequence\\n\\nfrom torchgen import local\\nfrom torchgen.api.types import (\\n ArgName,\\n BaseCType,\\n Binding,\\n ConstRefCType,\\n CType,\\n MutRefCType,\\n NamedCType,\\n SpecialArgName,\\n TupleCType,\\n VectorCType,\\n voidT,\\n)\\nfrom torchgen.executorch.api.types import (\\n ArrayRefCType,\\n BaseTypeToCppMapping,\\n OptionalCType,\\n scalarT,\\n tensorListT,\\n tensorT,\\n)\\nfrom torchgen.model import (\\n Argument,\\n Arguments,\\n BaseTy,\\n BaseType,\\n ListType,\\n NativeFunction,\\n OptionalType,\\n Return,\\n SelfArgument,\\n TensorOptionsArguments,\\n Type,\\n)\\nfrom torchgen.utils import assert_never\\n\\n\\n\\\"\\\"\\\"\\nThis file describes the translation of JIT schema to the public C++ API, which is what people use when they call\\nfunctions like at::add. It also serves as a native function API, which is the signature of kernels,\\nsince in Executorch CppSignature is the same as NativeSignature.\\n\\nDifference between this file and torchgen.api.cpp.py:\\n\\n - Executorch doesn't support TensorOptions, however in this file we still keep the logic here to be compatible with\\n torchgen.api.cpp, so that we can do stuff like ATen mode (running ATen kernels in Executorch).\\n\\n - Executorch doesn't support Dimname.\\n\\n - Executorch runtime doesn't support SymInt, will treat it as int.\\n\\\"\\\"\\\"\\n\\n\\n# Translation of \\\"value types\\\" in JIT schema to C++ API type. Value\\n# types look the same no matter if they are argument types or return\\n# types. Returns None if the type in question is not a value type.\\ndef valuetype_type(\\n t: Type,\\n *,\\n binds: ArgName,\\n remove_non_owning_ref_types: bool = False,\\n) -> NamedCType | None:\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Tensor or t.name == BaseTy.Scalar:\\n return None\\n # For SymInt we simply treat it as int.\\n elif str(t) == \\\"SymInt\\\":\\n return NamedCType(binds, BaseCType(BaseTypeToCppMapping[BaseTy.int]))\\n if remove_non_owning_ref_types:\\n if t.name == BaseTy.str:\\n raise AssertionError(\\n \\\"string ref->value conversion: not implemented yet\\\"\\n )\\n # All other BaseType currently map directly to BaseCppTypes.\\n return NamedCType(binds, BaseCType(BaseTypeToCppMapping[t.name]))\\n elif isinstance(t, OptionalType):\\n elem = valuetype_type(t.elem, binds=binds)\\n if elem is None:\\n return None\\n return NamedCType(binds, OptionalCType(elem.type))\\n elif isinstance(t, ListType):\\n if str(t.elem) == \\\"bool\\\":\\n assert t.size is not None\\n return NamedCType(\\n binds, ArrayRefCType(BaseCType(BaseTypeToCppMapping[BaseTy.bool]))\\n )\\n else:\\n return None\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\n# Translation of types occurring in JIT arguments to a C++ argument type.\\n# If remove_non_owning_ref_types is set, we'll guarantee that the outputed CType is not a non-owning reference type.\\n# For example, we'll return std::vector instead of IntArrayRef.\\n# See Note [translation from C++ reference to value types]\\ndef argumenttype_type(\\n t: Type,\\n *,\\n mutable: bool,\\n binds: ArgName,\\n remove_non_owning_ref_types: bool = False,\\n) -> NamedCType:\\n # If it's a value type, do the value type translation\\n r = valuetype_type(\\n t,\\n binds=binds,\\n remove_non_owning_ref_types=remove_non_owning_ref_types,\\n )\\n if r is not None:\\n return r\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Tensor:\\n if mutable and not local.use_const_ref_for_mutable_tensors():\\n return NamedCType(binds, MutRefCType(BaseCType(tensorT)))\\n else:\\n return NamedCType(binds, ConstRefCType(BaseCType(tensorT)))\\n elif t.name == BaseTy.Scalar:\\n return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))\\n else:\\n raise AssertionError(f\\\"base type should have been value type {t}\\\")\\n elif isinstance(t, OptionalType):\\n if str(t.elem) == \\\"Tensor\\\":\\n if mutable and not local.use_const_ref_for_mutable_tensors():\\n return NamedCType(\\n binds, MutRefCType(BaseCType(tensorT))\\n ) # TODO: fix this discrepancy\\n else:\\n return NamedCType(\\n binds, ConstRefCType(OptionalCType(BaseCType(tensorT)))\\n )\\n elif str(t.elem) == \\\"Scalar\\\":\\n return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT))))\\n elem = argumenttype_type(t.elem, mutable=mutable, binds=binds)\\n return NamedCType(binds, OptionalCType(elem.type))\\n elif isinstance(t, ListType):\\n # TODO: keeping these special cases for Tensor[] and Tensor?[] so that we can hookup with ATen kernels.\\n if str(t.elem) == \\\"Tensor\\\":\\n return NamedCType(binds, BaseCType(tensorListT))\\n elif str(t.elem) == \\\"Dimname\\\":\\n raise NotImplementedError(\\\"Executorch doesn't support Dimname\\\")\\n elif str(t.elem) == \\\"Tensor?\\\":\\n return NamedCType(binds, ArrayRefCType(OptionalCType(BaseCType(tensorT))))\\n elem = argumenttype_type(t.elem, mutable=mutable, binds=binds)\\n return NamedCType(binds, ArrayRefCType(elem.type))\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\n# Translate a JIT argument into its C++ type\\ndef argument_type(a: Argument, *, binds: ArgName) -> NamedCType:\\n return argumenttype_type(a.type, mutable=a.is_write, binds=binds)\\n\\n\\n# Translation of a (non-multi) return type from JIT to C++\\n# N.B: returntype_type returns a CType, not a NamedCType.\\n# This is mostly because of the mismatch between return types and return names.\\n# e.g. a function with a return type of 'void' has 0 return names,\\n# and a function with a return type of 'std::tuple' has >1 return name.\\ndef returntype_type(t: Type, *, mutable: bool) -> CType:\\n # placeholder is ignored\\n r = valuetype_type(t, binds=\\\"__placeholder__\\\")\\n if r is not None:\\n return r.type\\n\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Tensor:\\n if mutable:\\n if local.use_const_ref_for_mutable_tensors():\\n return ConstRefCType(BaseCType(tensorT))\\n else:\\n return MutRefCType(BaseCType(tensorT))\\n else:\\n # Note [Tensor Copy Returns]\\n # Currently, we use \\\"Argument.is_write\\\" to determine\\n # whether or not Tensor return types should be copies or references.\\n # If that ever changes, take a look at other locations of this note!\\n return BaseCType(tensorT)\\n elif t.name == BaseTy.Scalar:\\n return BaseCType(scalarT)\\n elif isinstance(t, ListType):\\n assert (\\n not mutable\\n ), \\\"Native functions should never return a mutable tensor list. They should return void.\\\"\\n elem = returntype_type(t.elem, mutable=False)\\n assert t.size is None, f\\\"fixed size list returns not supported: {t}\\\"\\n return VectorCType(elem)\\n\\n raise AssertionError(f\\\"unrecognized return type {t}\\\")\\n\\n\\n# Translation of a single return to its C++ type\\ndef return_type(r: Return) -> CType:\\n return returntype_type(r.type, mutable=r.is_write)\\n\\n\\n# Translation of a full (possibly multi) return from JIT to its C++ type\\ndef returns_type(rs: Sequence[Return]) -> CType:\\n if len(rs) == 0:\\n return BaseCType(voidT)\\n elif len(rs) == 1:\\n return return_type(rs[0])\\n else:\\n return TupleCType([return_type(r) for r in rs])\\n\\n\\ndef return_names(f: NativeFunction, *, fallback_name: str = \\\"result\\\") -> Sequence[str]:\\n returns: list[str] = []\\n for i, r in enumerate(f.func.returns):\\n # If we have an inplace function, the return argument is\\n # implicitly named self.\\n # TODO: Consider incorporating this into the data model\\n if f.func.name.name.inplace:\\n assert i == 0, \\\"illegal inplace function with multiple returns\\\"\\n name = \\\"self\\\"\\n # If we are out function, the name is the name of the\\n # corresponding output function (r.name will get recorded\\n # in field_name later.)\\n elif f.func.is_out_fn():\\n name = f.func.arguments.out[i].name\\n # If the return argument is explicitly named...\\n elif r.name:\\n name_conflict = any(\\n r.name == a.name for a in f.func.schema_order_arguments()\\n )\\n if name_conflict and not f.func.is_out_fn():\\n name = f\\\"{r.name}_return\\\"\\n else:\\n name = r.name\\n # If there is no explicit name and no fallback name was passed in, we just name the output result,\\n # unless it's a multi-return, in which case it's result0,\\n # result1, etc (zero-indexed)\\n else:\\n name = fallback_name if len(f.func.returns) == 1 else f\\\"{fallback_name}{i}\\\"\\n returns.append(name)\\n return returns\\n\\n\\nJIT_TO_CPP_DEFAULT = {\\n \\\"False\\\": \\\"false\\\",\\n \\\"True\\\": \\\"true\\\",\\n \\\"None\\\": \\\"torch::executorch::nullopt\\\", # UGH this one is type directed\\n \\\"[]\\\": \\\"{}\\\",\\n \\\"contiguous_format\\\": \\\"torch::executorch::MemoryFormat::Contiguous\\\",\\n \\\"long\\\": \\\"torch::executorch::kLong\\\",\\n}\\n\\n\\n# Convert a JIT default into C++ expression representing the default\\ndef default_expr(d: str, t: Type) -> str:\\n if d == \\\"None\\\" and str(t) == \\\"Tensor?\\\":\\n return \\\"{}\\\"\\n if isinstance(t, BaseType) and t.name is BaseTy.str:\\n # Schema allows single quotes but C++ needs double\\n if len(d) >= 2 and d[0] == \\\"'\\\" and d[-1] == \\\"'\\\":\\n s = \\\"\\\"\\n i = 1\\n while i + 1 < len(d):\\n if d[i] != \\\"\\\\\\\\\\\":\\n if d[i] == '\\\"':\\n s += '\\\\\\\\\\\"'\\n else:\\n s += d[i]\\n i += 1\\n else:\\n if d[i + 1] == \\\"'\\\":\\n s += \\\"'\\\"\\n else:\\n s += d[i : i + 2]\\n i += 2\\n\\n return f'\\\"{s}\\\"'\\n\\n if isinstance(t, OptionalType):\\n if d == \\\"None\\\":\\n return \\\"torch::executor::nullopt\\\"\\n\\n return default_expr(d, t.elem)\\n\\n if isinstance(t, ListType):\\n if d.startswith(\\\"[\\\") and d.endswith(\\\"]\\\"):\\n return \\\"{\\\" + d[1:-1] + \\\"}\\\"\\n elif t.size is None:\\n # NOTE: Sized lists can have scalar defaults\\n raise ValueError(f\\\"Expected a list default '[...]' but found: '{d}'\\\")\\n\\n return JIT_TO_CPP_DEFAULT.get(d, d)\\n\\n\\n# Convert an argument into its C++ API form\\n\\n\\ndef argument(\\n a: Argument | TensorOptionsArguments | SelfArgument,\\n *,\\n cpp_no_default_args: set[str],\\n method: bool,\\n faithful: bool,\\n has_tensor_options: bool,\\n) -> list[Binding]:\\n def sub_argument(\\n a: Argument | TensorOptionsArguments | SelfArgument,\\n ) -> list[Binding]:\\n return argument(\\n a,\\n cpp_no_default_args=cpp_no_default_args,\\n method=method,\\n faithful=faithful,\\n has_tensor_options=has_tensor_options,\\n )\\n\\n if isinstance(a, Argument):\\n binds: ArgName\\n if a.name == \\\"memory_format\\\" and has_tensor_options:\\n binds = SpecialArgName.possibly_redundant_memory_format\\n else:\\n binds = a.name\\n default: str | None = None\\n if a.name not in cpp_no_default_args and a.default is not None:\\n default = default_expr(a.default, a.type)\\n return [\\n Binding(\\n nctype=argument_type(a, binds=binds),\\n name=a.name,\\n default=default,\\n argument=a,\\n )\\n ]\\n elif isinstance(a, TensorOptionsArguments):\\n raise NotImplementedError(\\\"Need to implement type resolution for TensorOptions\\\")\\n elif isinstance(a, SelfArgument):\\n if method:\\n # Caller is responsible for installing implicit this in context!\\n return []\\n else:\\n return sub_argument(a.argument)\\n else:\\n assert_never(a)\\n\\n\\ndef arguments(\\n arguments: Arguments,\\n *,\\n faithful: bool,\\n method: bool,\\n cpp_no_default_args: set[str],\\n) -> list[Binding]:\\n args: list[Argument | TensorOptionsArguments | SelfArgument] = []\\n if faithful:\\n args.extend(arguments.non_out)\\n args.extend(arguments.out)\\n else:\\n args.extend(arguments.out)\\n args.extend(arguments.non_out)\\n return [\\n r.no_default() if faithful else r\\n for a in args\\n for r in argument(\\n a,\\n faithful=faithful,\\n method=method,\\n has_tensor_options=arguments.tensor_options is not None,\\n cpp_no_default_args=cpp_no_default_args,\\n )\\n ]\\n\\n\\nfrom __future__ import annotations\\n\\nfrom collections import defaultdict\\nfrom dataclasses import dataclass\\nfrom typing import Sequence, TYPE_CHECKING\\n\\nfrom torchgen import dest\\n\\n\\n# disable import sorting to avoid circular dependency.\\nfrom torchgen.api.types import DispatcherSignature # usort: skip\\nfrom torchgen.context import method_with_native_function\\nfrom torchgen.model import BaseTy, BaseType, DispatchKey, NativeFunction, Variant\\nfrom torchgen.utils import concatMap, Target\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.executorch.model import ETKernelIndex\\n from torchgen.selective_build.selector import SelectiveBuilder\\n\\n\\n# Generates RegisterKernelStub.cpp, which provides placeholder kernels for custom operators. This will be used at\\n# model authoring side.\\n@dataclass(frozen=True)\\nclass ComputeNativeFunctionStub:\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> str | None:\\n if Variant.function not in f.variants:\\n return None\\n\\n sig = DispatcherSignature.from_schema(\\n f.func, prefix=f\\\"wrapper_CPU_{f.func.name.overload_name}_\\\", symint=False\\n )\\n assert sig is not None\\n if len(f.func.returns) == 0:\\n ret_name = \\\"\\\"\\n elif len(f.func.returns) == 1:\\n if f.func.arguments.out:\\n ret_name = f.func.arguments.out[0].name\\n else:\\n ret_name = next(\\n (\\n a.name\\n for a in f.func.arguments.flat_non_out\\n if a.type == f.func.returns[0].type\\n ),\\n \\\"\\\",\\n )\\n if not ret_name:\\n # if return type is tensor\\n if f.func.returns[0].type == BaseType(BaseTy.Tensor):\\n # Returns an empty tensor\\n ret_name = \\\"at::Tensor()\\\"\\n else:\\n raise Exception( # noqa: TRY002\\n f\\\"Can't handle this return type {f.func}\\\"\\n ) # noqa: TRY002\\n elif len(f.func.arguments.out) == len(f.func.returns):\\n # Returns a tuple of out arguments\\n tensor_type = \\\"at::Tensor &\\\"\\n comma = \\\", \\\"\\n ret_name = f\\\"\\\"\\\"::std::tuple<{comma.join([tensor_type] * len(f.func.returns))}>(\\n {comma.join([r.name for r in f.func.arguments.out])}\\n )\\\"\\\"\\\"\\n else:\\n assert all(\\n a.type == BaseType(BaseTy.Tensor) for a in f.func.returns\\n ), f\\\"Only support tensor returns but got {f.func.returns}\\\"\\n # Returns a tuple of empty tensors\\n tensor_type = \\\"at::Tensor\\\"\\n comma = \\\", \\\"\\n ret_name = f\\\"\\\"\\\"::std::tuple<{comma.join([tensor_type] * len(f.func.returns))}>(\\n {comma.join([\\\"at::Tensor()\\\" for _ in f.func.returns])}\\n )\\\"\\\"\\\"\\n ret_str = f\\\"return {ret_name};\\\" if len(f.func.returns) > 0 else \\\"\\\"\\n return f\\\"\\\"\\\"\\n{sig.defn()} {{\\n {ret_str}\\n}}\\n \\\"\\\"\\\"\\n\\n\\ndef gen_custom_ops_registration(\\n *,\\n native_functions: Sequence[NativeFunction],\\n selector: SelectiveBuilder,\\n kernel_index: ETKernelIndex,\\n rocm: bool,\\n) -> tuple[str, str]:\\n \\\"\\\"\\\"\\n Generate custom ops registration code for dest.RegisterDispatchKey.\\n\\n :param native_functions: a sequence of `NativeFunction`\\n :param selector: for selective build.\\n :param kernel_index: kernels for all the ops.\\n :param rocm: bool for dest.RegisterDispatchKey.\\n :return: generated C++ code to register custom operators into PyTorch\\n \\\"\\\"\\\"\\n\\n # convert kernel index to BackendIndex. This is because we can't handle ETKernelIndex yet.\\n # TODO larryliu: evaluate if this code is still needed. If yes let it handle ETKernelIndex.\\n\\n dispatch_key = DispatchKey.CPU\\n backend_index = kernel_index._to_backend_index()\\n static_init_dispatch_registrations = \\\"\\\"\\n ns_grouped_native_functions: dict[str, list[NativeFunction]] = defaultdict(list)\\n for native_function in native_functions:\\n ns_grouped_native_functions[native_function.namespace].append(native_function)\\n\\n for namespace, functions in ns_grouped_native_functions.items():\\n if len(functions) == 0:\\n continue\\n dispatch_registrations_body = \\\"\\\\n\\\".join(\\n list(\\n concatMap(\\n dest.RegisterDispatchKey(\\n backend_index,\\n Target.REGISTRATION,\\n selector,\\n rocm=rocm,\\n symint=False,\\n class_method_name=None,\\n skip_dispatcher_op_registration=False,\\n ),\\n functions,\\n )\\n )\\n )\\n static_init_dispatch_registrations += f\\\"\\\"\\\"\\nTORCH_LIBRARY_IMPL({namespace}, {dispatch_key}, m) {{\\n{dispatch_registrations_body}\\n}};\\\"\\\"\\\"\\n anonymous_definition = \\\"\\\\n\\\".join(\\n list(\\n concatMap(\\n dest.RegisterDispatchKey(\\n backend_index,\\n Target.ANONYMOUS_DEFINITION,\\n selector,\\n rocm=rocm,\\n symint=False,\\n class_method_name=None,\\n skip_dispatcher_op_registration=False,\\n ),\\n native_functions,\\n )\\n )\\n )\\n return anonymous_definition, static_init_dispatch_registrations\\n\\n\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\nfrom typing import TYPE_CHECKING\\n\\nimport torchgen.api.cpp as aten_cpp\\nfrom torchgen.executorch.api.types.types import contextArg\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.api.types import Binding, CType\\n from torchgen.model import FunctionSchema, NativeFunction\\n\\n\\n@dataclass(frozen=True)\\nclass ExecutorchCppSignature:\\n \\\"\\\"\\\"\\n This signature is merely a CppSignature with Executorch types (optionally\\n contains KernelRuntimeContext as well). The inline definition of\\n CppSignature is generated in Functions.h and it's used by unboxing\\n functions.\\n \\\"\\\"\\\"\\n\\n # The schema this signature is derived from\\n func: FunctionSchema\\n\\n # The set of C++ arguments which should not have defaults applied to them\\n cpp_no_default_args: set[str]\\n\\n # Allows you to prepend an arbitrary prefix to the signature name.\\n # This is useful for parts of the codegen that generate wrappers around kernels,\\n # and need to avoid naming collisions.\\n prefix: str = \\\"\\\"\\n\\n def arguments(self, *, include_context: bool = True) -> list[Binding]:\\n return ([contextArg] if include_context else []) + et_cpp.arguments(\\n self.func.arguments,\\n faithful=True, # always faithful, out argument at the end\\n method=False, # method not supported\\n cpp_no_default_args=self.cpp_no_default_args,\\n )\\n\\n def name(self) -> str:\\n return self.prefix + aten_cpp.name(\\n self.func,\\n faithful_name_for_out_overloads=True,\\n )\\n\\n def decl(self, name: str | None = None, *, include_context: bool = True) -> str:\\n args_str = \\\", \\\".join(\\n a.decl() for a in self.arguments(include_context=include_context)\\n )\\n if name is None:\\n name = self.name()\\n return f\\\"{self.returns_type().cpp_type()} {name}({args_str})\\\"\\n\\n def defn(self, name: str | None = None) -> str:\\n args = [a.defn() for a in self.arguments()]\\n args_str = \\\", \\\".join(args)\\n if name is None:\\n name = self.name()\\n return f\\\"{self.returns_type().cpp_type()} {name}({args_str})\\\"\\n\\n def returns_type(self) -> CType:\\n return et_cpp.returns_type(self.func.returns)\\n\\n @staticmethod\\n def from_native_function(\\n f: NativeFunction, *, prefix: str = \\\"\\\"\\n ) -> ExecutorchCppSignature:\\n return ExecutorchCppSignature(\\n func=f.func, prefix=prefix, cpp_no_default_args=f.cpp_no_default_args\\n )\\n\\n\\nfrom torchgen.executorch.api import et_cpp\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\n\\nfrom torchgen.api.types import (\\n BaseCppType,\\n BaseCType,\\n Binding,\\n boolT,\\n CType,\\n doubleT,\\n Expr,\\n longT,\\n MutRefCType,\\n NamedCType,\\n)\\nfrom torchgen.model import BaseTy\\n\\n\\nhalfT = BaseCppType(\\\"torch::executor\\\", \\\"Half\\\")\\nbfloat16T = BaseCppType(\\\"torch::executor\\\", \\\"BFloat16\\\")\\nstringT = BaseCppType(\\\"torch::executor\\\", \\\"string_view\\\")\\nscalarTypeT = BaseCppType(\\\"torch::executor\\\", \\\"ScalarType\\\")\\ntensorT = BaseCppType(\\\"torch::executor\\\", \\\"Tensor\\\")\\ntensorListT = BaseCppType(\\\"torch::executor\\\", \\\"TensorList\\\")\\nscalarT = BaseCppType(\\\"torch::executor\\\", \\\"Scalar\\\")\\nmemoryFormatT = BaseCppType(\\\"torch::executor\\\", \\\"MemoryFormat\\\")\\nintArrayRefT = BaseCppType(\\\"torch::executor\\\", \\\"IntArrayRef\\\")\\noptionalT = BaseCppType(\\\"torch::executor\\\", \\\"optional\\\")\\ncontextT = BaseCppType(\\\"torch::executor\\\", \\\"KernelRuntimeContext\\\")\\n\\ncontextExpr = Expr(\\n expr=\\\"context\\\",\\n type=NamedCType(name=\\\"context\\\", type=MutRefCType(BaseCType(contextT))),\\n)\\n\\ncontextArg = Binding(\\n name=\\\"context\\\",\\n nctype=contextExpr.type,\\n argument=None, # type: ignore[arg-type]\\n default=None,\\n)\\n\\nBaseTypeToCppMapping: dict[BaseTy, BaseCppType] = {\\n BaseTy.int: longT,\\n BaseTy.float: doubleT,\\n BaseTy.bool: boolT,\\n BaseTy.str: stringT,\\n BaseTy.ScalarType: scalarTypeT,\\n BaseTy.Tensor: tensorT,\\n BaseTy.Scalar: scalarT,\\n BaseTy.MemoryFormat: memoryFormatT,\\n}\\n\\n\\n@dataclass(frozen=True)\\nclass OptionalCType(CType):\\n elem: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n # Do not pass `strip_ref` recursively.\\n return f\\\"torch::executor::optional<{self.elem.cpp_type()}>\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"torch::executor::optional<{self.elem.cpp_type_registration_declarations()}>\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return OptionalCType(self.elem.remove_const_ref())\\n\\n\\n@dataclass(frozen=True)\\nclass ArrayRefCType(CType):\\n elem: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n # Do not pass `strip_ref` recursively.\\n return f\\\"torch::executor::ArrayRef<{self.elem.cpp_type()}>\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"torch::executor::ArrayRef<{self.elem.cpp_type_registration_declarations()}>\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return ArrayRefCType(self.elem.remove_const_ref())\\n\\n\\nfrom torchgen.executorch.api.types.types import *\\n\\n\\nfrom torchgen.executorch.api.types.signatures import * # usort: skip\\n\\n\\nfrom __future__ import annotations\\n\\nfrom torchgen.api import dispatcher\\nfrom torchgen.api.types import (\\n BaseCppType,\\n BaseCType,\\n Binding,\\n boolT,\\n ConstRefCType,\\n CType,\\n longT,\\n NamedCType,\\n tensorT,\\n)\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n NativeFunction,\\n NativeFunctionsViewGroup,\\n)\\n\\n\\n# This file describes the translation of JIT schema to API's used\\n# when creating view lambdas that are used by the functionalization pass.\\n# There are two types of lambdas: forward lambdas and reverse lambdas.\\n# These API's mostly follow the dispatcher API, with a few quirks:\\n# - The lambda capture has to convert reference types to value types\\n# - While the forward lambda just directly calls into the at::_ops API\\n# (following the dispatcher convention), the logic here for the reverse lambda\\n# is responsible for generating both the call-site, and the declarations\\n# (which are implemented manually in the at::functionalization::impl namespace).\\n\\n# The lambdas generated for each view op in the functionalization pass are of the form\\n# [capture_arguments](outer_arguments) -> returns_type {\\n# return name(inner_arguments);\\n# }\\n\\n# Define some specific lambda input arguments.\\nbase_binding = Binding(\\n name=\\\"base\\\",\\n nctype=NamedCType(name=\\\"base\\\", type=ConstRefCType(BaseCType(tensorT))),\\n argument=Argument(\\n name=\\\"base\\\", type=BaseType(BaseTy.Tensor), default=None, annotation=None\\n ),\\n default=None,\\n)\\nmutated_view_binding = Binding(\\n name=\\\"mutated_view\\\",\\n nctype=NamedCType(name=\\\"mutated_view\\\", type=ConstRefCType(BaseCType(tensorT))),\\n argument=Argument(\\n name=\\\"base\\\", type=BaseType(BaseTy.Tensor), default=None, annotation=None\\n ),\\n default=None,\\n)\\nmutated_view_idx_binding = Binding(\\n name=\\\"mutated_view_idx\\\",\\n nctype=NamedCType(name=\\\"mutated_view_idx\\\", type=BaseCType(longT)),\\n argument=Argument(\\n name=\\\"base\\\", type=BaseType(BaseTy.Tensor), default=None, annotation=None\\n ),\\n default=None,\\n)\\nreapply_views_binding = Binding(\\n name=\\\"reapply_views\\\",\\n nctype=NamedCType(name=\\\"reapply_views\\\", type=BaseCType(boolT)),\\n argument=Argument(\\n name=\\\"reapply_views\\\", type=BaseType(BaseTy.bool), default=None, annotation=None\\n ),\\n default=None,\\n)\\n\\nInverseReturnModeT = BaseCppType(\\\"at::functionalization\\\", \\\"InverseReturnMode\\\")\\ninverse_return_mode_binding = Binding(\\n name=\\\"inverse_return_mode\\\",\\n nctype=NamedCType(name=\\\"inverse_return_mode\\\", type=BaseCType(InverseReturnModeT)),\\n argument=Argument(\\n name=\\\"inverse_return_mode\\\",\\n # NB: not actually a bool but it doesn't matter because this isn't used\\n type=BaseType(BaseTy.bool),\\n default=None,\\n annotation=None,\\n ),\\n default=None,\\n)\\n\\n\\n# The lambda capture itself doesn't have a name.\\n# The name returned here corresponds to the name of the inner function called by the lambda.\\ndef name(\\n g: NativeFunctionsViewGroup,\\n *,\\n is_reverse: bool,\\n include_namespace: bool,\\n reapply_views: bool | None = None,\\n) -> str:\\n if reapply_views is None:\\n # reapply_views is only important for the fwd lambda,\\n # since we always plumb the runtime \\\"reapply_views\\\" argument into the reverse function.\\n assert is_reverse\\n if is_reverse:\\n return reverse_name(g.view, include_namespace)\\n # in the forward case, we just directly call into the at::_ops API (so we always need the namespace)\\n assert include_namespace\\n assert g.view_copy is not None\\n api_name = (\\n g.view.func.name.unambiguous_name()\\n if reapply_views\\n else g.view_copy.func.name.unambiguous_name()\\n )\\n return f\\\"at::_ops::{api_name}::call\\\"\\n\\n\\ndef reverse_name(f: NativeFunction, include_namespace: bool) -> str:\\n # for the reverse: we plumb the \\\"reapply_views\\\" flag into that function and support\\n # both copy and non-copy variants. (We could avoid doing that, but that would require\\n # writing out twice as many view inverse functions).\\n api_name = f.func.name.unambiguous_name()\\n # in the reverse case, we codegen both the call-sites (which need the full namespace) and the declarations (which don't)\\n if include_namespace:\\n return f\\\"at::functionalization::FunctionalInverses::{api_name}_inverse\\\"\\n else:\\n return f\\\"{api_name}_inverse\\\"\\n\\n\\ndef capture_arguments(func: FunctionSchema, *, is_reverse: bool) -> list[Binding]:\\n # capture arguments include all arguments except `self`.\\n # Importantly, they don't include any C++ reference types (or else we'll get a dangling reference in the capture),\\n # So any reference types (IntArrayRef) need to be converted to value types (vector)\\n args = func.arguments.flat_all\\n assert args[0].type == BaseType(BaseTy.Tensor)\\n non_self_args = args[1:]\\n non_self_value_bindings = [\\n dispatcher.argument(a, remove_non_owning_ref_types=True) for a in non_self_args\\n ]\\n\\n all_bindings = [\\n inverse_return_mode_binding if is_reverse else reapply_views_binding\\n ]\\n all_bindings.extend(non_self_value_bindings)\\n return all_bindings\\n\\n\\ndef returns_type(func: FunctionSchema) -> CType:\\n # Assertion: all view ops return tensor-like outputs\\n assert len(func.returns) >= 1\\n for ret in func.returns:\\n assert ret.type.is_tensor_like()\\n # However, the return type of the lambda is always an individual tensor.\\n # For multi-tensor outputs, each tensor needs to be tracked individually.\\n return BaseCType(tensorT)\\n\\n\\ndef outer_arguments(*, is_reverse: bool) -> list[Binding]:\\n if is_reverse:\\n return [base_binding, mutated_view_binding, mutated_view_idx_binding]\\n else:\\n return [base_binding, mutated_view_idx_binding]\\n\\n\\ndef inner_call_index(func: FunctionSchema) -> Binding | None:\\n # For view ops that return multiple tensors (like `split`), we generate a separate lambda for each output.\\n # When we replay a view op that returns multiple tensors, we need to index into the output appropriately\\n if len(func.returns) > 1 or (\\n len(func.returns) == 1 and func.returns[0].type.is_list_like()\\n ):\\n return mutated_view_idx_binding\\n return None\\n\\n\\ndef inner_arguments(func: FunctionSchema, is_reverse: bool) -> list[Binding]:\\n args = func.arguments.flat_all\\n assert args[0].type == BaseType(BaseTy.Tensor)\\n non_self_args = args[1:]\\n # The forward lambda calls the at::_ops API, while the reverse lambda calls the view inverse API.\\n # Both of these follow the dispatcher API.\\n non_self_bindings = [dispatcher.argument(a) for a in non_self_args]\\n if not is_reverse:\\n # the forward lambda swaps out the original tensor argument with the lambd arg \\\"base\\\"\\n return [base_binding] + non_self_bindings\\n else:\\n # the reverse lambda does the same, but with an additional \\\"mutated_view\\\" arg\\n # additionally, we have a calling convention: for view ops that return multiple tensor outputs\\n # their corresponding view_inverse function takes in an additional index argument.\\n index_binding = inner_call_index(func)\\n if index_binding is not None:\\n return [\\n base_binding,\\n mutated_view_binding,\\n inverse_return_mode_binding,\\n index_binding,\\n ] + non_self_bindings\\n else:\\n return [\\n base_binding,\\n mutated_view_binding,\\n inverse_return_mode_binding,\\n ] + non_self_bindings\\n\\n\\nfrom __future__ import annotations\\n\\nfrom typing import Sequence\\n\\nfrom torchgen import local\\nfrom torchgen.api import cpp\\nfrom torchgen.api.types import (\\n ArgName,\\n BaseCType,\\n Binding,\\n boolT,\\n ConstRefCType,\\n CType,\\n deviceT,\\n layoutT,\\n ListCType,\\n MutRefCType,\\n NamedCType,\\n OptionalCType,\\n scalarT,\\n scalarTypeT,\\n tensorT,\\n)\\nfrom torchgen.model import (\\n Argument,\\n FunctionSchema,\\n Return,\\n SelfArgument,\\n TensorOptionsArguments,\\n Type,\\n)\\nfrom torchgen.utils import assert_never\\n\\n\\n# This file describes the translation of JIT schema to the native functions API.\\n# This looks a lot like the C++ API (which makes historical sense, because the\\n# idea was you wrote native functions to implement functions in the C++ API),\\n# but over time we have evolved the C++ API without actually changing our\\n# native:: kernels. The intention is to make native API and dispatcher API\\n# line up as closely as possible, since this results in the least overhead\\n# (no translation is needed from dispatcher API to native API).\\n#\\n# NB: this is symint aware, you will get the non-SymInt variant for some\\n# dispatch entries and SymInt for others.\\n\\n\\ndef name(func: FunctionSchema) -> str:\\n name = str(func.name.name)\\n # TODO: delete this!\\n if func.is_out_fn():\\n name += \\\"_out\\\"\\n if func.name.overload_name:\\n name += f\\\"_{func.name.overload_name}\\\"\\n return name\\n\\n\\ndef argumenttype_type(\\n t: Type, *, mutable: bool, binds: ArgName, symint: bool\\n) -> NamedCType:\\n if str(t) == \\\"Tensor?\\\":\\n tensor_type: OptionalCType = OptionalCType(BaseCType(tensorT))\\n if mutable and not local.use_const_ref_for_mutable_tensors():\\n return NamedCType(binds, MutRefCType(tensor_type))\\n else:\\n return NamedCType(binds, ConstRefCType(tensor_type))\\n elif str(t) == \\\"Tensor?[]\\\":\\n return NamedCType(\\n binds, ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT))))\\n )\\n elif str(t) == \\\"Scalar\\\":\\n return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))\\n elif str(t) == \\\"Scalar?\\\":\\n return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT))))\\n return cpp.argumenttype_type(t, mutable=mutable, binds=binds, symint=symint)\\n\\n\\ndef returns_type(rs: Sequence[Return], *, symint: bool) -> CType:\\n return cpp.returns_type(rs, symint=symint)\\n\\n\\ndef argument_type(a: Argument, *, binds: ArgName, symint: bool) -> NamedCType:\\n return argumenttype_type(a.type, mutable=a.is_write, binds=binds, symint=symint)\\n\\n\\ndef argument(\\n a: Argument | SelfArgument | TensorOptionsArguments,\\n *,\\n is_out: bool,\\n symint: bool,\\n) -> list[Binding]:\\n # Ideally, we NEVER default native functions. However, there are a number\\n # of functions that call native:: directly and rely on the defaulting\\n # existing. So for BC, we generate defaults for non-out variants (but not\\n # for out variants, where it is impossible to generate an appropriate\\n # default)\\n should_default = not is_out\\n if isinstance(a, Argument):\\n default: str | None = None\\n if should_default and a.default is not None:\\n default = cpp.default_expr(a.default, a.type, symint=symint)\\n return [\\n Binding(\\n nctype=argument_type(a, binds=a.name, symint=symint),\\n name=a.name,\\n default=default,\\n argument=a,\\n )\\n ]\\n elif isinstance(a, SelfArgument):\\n # Erase SelfArgument from the distinction\\n return argument(a.argument, is_out=is_out, symint=symint)\\n elif isinstance(a, TensorOptionsArguments):\\n default = None\\n if should_default:\\n default = \\\"{}\\\"\\n # TODO: Not sure why the arguments assigned here are for\\n # TensorOptionsArguments and not the constituent pieces. It seems\\n # to matter\\n return [\\n Binding(\\n nctype=NamedCType(\\\"dtype\\\", OptionalCType(BaseCType(scalarTypeT))),\\n name=\\\"dtype\\\",\\n default=default,\\n argument=a,\\n ),\\n Binding(\\n nctype=NamedCType(\\\"layout\\\", OptionalCType(BaseCType(layoutT))),\\n name=\\\"layout\\\",\\n default=default,\\n argument=a,\\n ),\\n Binding(\\n nctype=NamedCType(\\\"device\\\", OptionalCType(BaseCType(deviceT))),\\n name=\\\"device\\\",\\n default=default,\\n argument=a,\\n ),\\n Binding(\\n nctype=NamedCType(\\\"pin_memory\\\", OptionalCType(BaseCType(boolT))),\\n name=\\\"pin_memory\\\",\\n default=default,\\n argument=a,\\n ),\\n ]\\n else:\\n assert_never(a)\\n\\n\\ndef arguments(func: FunctionSchema, *, symint: bool) -> list[Binding]:\\n args: list[Argument | TensorOptionsArguments | SelfArgument] = []\\n args.extend(func.arguments.non_out)\\n args.extend(func.arguments.out)\\n return [\\n r for arg in args for r in argument(arg, symint=symint, is_out=func.is_out_fn())\\n ]\\n\\n\\nfrom __future__ import annotations\\n\\nfrom typing import Sequence\\n\\nfrom torchgen import local\\nfrom torchgen.api.types import (\\n ArgName,\\n ArrayCType,\\n ArrayRefCType,\\n BaseCType,\\n BaseTypeToCppMapping,\\n Binding,\\n boolT,\\n ConstRefCType,\\n CType,\\n dimnameListT,\\n intArrayRefT,\\n iTensorListRefT,\\n ListCType,\\n longT,\\n MutRefCType,\\n NamedCType,\\n OptionalCType,\\n optionalIntArrayRefT,\\n optionalSymIntArrayRefT,\\n scalarT,\\n SpecialArgName,\\n symIntArrayRefT,\\n SymIntT,\\n tensorListT,\\n tensorOptionsT,\\n tensorT,\\n TupleCType,\\n VectorCType,\\n voidT,\\n)\\nfrom torchgen.model import (\\n Argument,\\n Arguments,\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n ListType,\\n NativeFunction,\\n OptionalType,\\n Return,\\n SelfArgument,\\n TensorOptionsArguments,\\n Type,\\n)\\nfrom torchgen.utils import assert_never\\n\\n\\n# This file describes the translation of JIT schema to the public C++\\n# API, which is what people use when they call functions like at::add.\\n#\\n# Prominent characteristics of the C++ API:\\n#\\n# - dtype, layout, device and pin_memory are collected into\\n# a single C++ type TensorOptions (the native functions API\\n# also has this, but tensor options is really most relevant\\n# for the C++ API; it makes calling kwarg factory functions\\n# pleasant)\\n#\\n# - defaulting lives here (in fact, the dispatcher is completely\\n# oblivious of defaults!)\\n#\\n# BTW: policy on name collisions: we try not to have types with\\n# collisions, but functions are fair game to collide\\n\\n\\ndef name(\\n func: FunctionSchema,\\n *,\\n faithful_name_for_out_overloads: bool = False,\\n symint_overload: bool = False,\\n) -> str:\\n name = str(func.name.name)\\n if symint_overload:\\n name += \\\"_symint\\\"\\n if func.is_out_fn():\\n if faithful_name_for_out_overloads:\\n name += \\\"_outf\\\"\\n else:\\n name += \\\"_out\\\"\\n\\n return name\\n\\n\\n# Translation of \\\"value types\\\" in JIT schema to C++ API type. Value\\n# types look the same no matter if they are argument types or return\\n# types. Returns None if the type in question is not a value type.\\ndef valuetype_type(\\n t: Type,\\n *,\\n binds: ArgName,\\n mutable: bool = True,\\n remove_non_owning_ref_types: bool = False,\\n symint: bool = False,\\n) -> NamedCType | None:\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Tensor or t.name == BaseTy.Scalar:\\n return None\\n elif str(t) == \\\"SymInt\\\":\\n if symint:\\n return NamedCType(binds, BaseCType(SymIntT))\\n else:\\n return NamedCType(binds, BaseCType(longT))\\n if remove_non_owning_ref_types:\\n if t.name == BaseTy.str:\\n raise AssertionError(\\n \\\"string ref->value conversion: not implemented yet\\\"\\n )\\n # All other BaseType currently map directly to BaseCppTypes.\\n return NamedCType(binds, BaseCType(BaseTypeToCppMapping[t.name]))\\n elif isinstance(t, OptionalType):\\n elem = valuetype_type(t.elem, binds=binds, mutable=mutable, symint=symint)\\n if elem is None:\\n return None\\n return NamedCType(binds, OptionalCType(elem.type))\\n elif isinstance(t, ListType):\\n if str(t.elem) == \\\"bool\\\":\\n assert t.size is not None\\n return NamedCType(binds, ArrayCType(BaseCType(boolT), t.size))\\n else:\\n return None\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\n# Translation of types occurring in JIT arguments to a C++ argument type.\\n# If remove_non_owning_ref_types is set, we'll guarantee that the outputed CType is not a non-owning reference type.\\n# For example, we'll return std::vector instead of IntArrayRef.\\n# See Note [translation from C++ reference to value types]\\ndef argumenttype_type(\\n t: Type,\\n *,\\n mutable: bool,\\n binds: ArgName,\\n remove_non_owning_ref_types: bool = False,\\n symint: bool = False,\\n) -> NamedCType:\\n # If it's a value type, do the value type translation\\n r = valuetype_type(\\n t,\\n binds=binds,\\n mutable=mutable,\\n symint=symint,\\n remove_non_owning_ref_types=remove_non_owning_ref_types,\\n )\\n if r is not None:\\n return r\\n\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Tensor:\\n if mutable and not local.use_const_ref_for_mutable_tensors():\\n return NamedCType(binds, MutRefCType(BaseCType(tensorT)))\\n else:\\n return NamedCType(binds, ConstRefCType(BaseCType(tensorT)))\\n elif t.name == BaseTy.Scalar:\\n return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))\\n else:\\n raise AssertionError(f\\\"base type should have been value type {t}\\\")\\n elif isinstance(t, OptionalType):\\n if str(t.elem) == \\\"Tensor\\\":\\n if mutable and not local.use_const_ref_for_mutable_tensors():\\n return NamedCType(\\n binds, MutRefCType(BaseCType(tensorT))\\n ) # TODO: fix this discrepancy\\n else:\\n return NamedCType(\\n binds, ConstRefCType(OptionalCType(BaseCType(tensorT)))\\n )\\n elif str(t.elem) == \\\"Scalar\\\":\\n return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT))))\\n elif isinstance(t.elem, ListType) and str(t.elem.elem) == \\\"int\\\":\\n return NamedCType(binds, BaseCType(optionalIntArrayRefT))\\n elif isinstance(t.elem, ListType) and str(t.elem.elem) == \\\"SymInt\\\":\\n if symint:\\n return NamedCType(binds, BaseCType(optionalSymIntArrayRefT))\\n else:\\n return NamedCType(binds, BaseCType(optionalIntArrayRefT))\\n elem = argumenttype_type(t.elem, mutable=mutable, binds=binds, symint=symint)\\n return NamedCType(binds, OptionalCType(elem.type))\\n elif isinstance(t, ListType):\\n # TODO: remove these special cases, ArrayRef fallthrough works fine\\n if str(t.elem) == \\\"int\\\":\\n if remove_non_owning_ref_types:\\n return NamedCType(binds, VectorCType(BaseCType(longT)))\\n else:\\n return NamedCType(binds, BaseCType(intArrayRefT))\\n if str(t.elem) == \\\"SymInt\\\":\\n if remove_non_owning_ref_types:\\n if symint:\\n return NamedCType(binds, VectorCType(BaseCType(SymIntT)))\\n else:\\n return NamedCType(binds, VectorCType(BaseCType(longT)))\\n else:\\n if symint:\\n return NamedCType(binds, BaseCType(symIntArrayRefT))\\n else:\\n return NamedCType(binds, BaseCType(intArrayRefT))\\n if str(t.elem) == \\\"Tensor\\\":\\n if local.use_ilistref_for_tensor_lists():\\n return NamedCType(binds, ConstRefCType(BaseCType(iTensorListRefT)))\\n else:\\n return NamedCType(binds, BaseCType(tensorListT))\\n elif str(t.elem) == \\\"Scalar\\\":\\n return NamedCType(binds, ArrayRefCType(BaseCType(scalarT)))\\n elif str(t.elem) == \\\"Dimname\\\":\\n return NamedCType(binds, BaseCType(dimnameListT))\\n elif str(t.elem) == \\\"Tensor?\\\":\\n return NamedCType(\\n binds, ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT))))\\n )\\n elem = argumenttype_type(t.elem, mutable=mutable, binds=binds, symint=symint)\\n return NamedCType(binds, ArrayRefCType(elem.type))\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\n# Translate a JIT argument into its C++ type\\ndef argument_type(a: Argument, *, binds: ArgName, symint: bool = False) -> NamedCType:\\n return argumenttype_type(a.type, mutable=a.is_write, symint=symint, binds=binds)\\n\\n\\n# Translation of a (non-multi) return type from JIT to C++\\n# N.B: returntype_type returns a CType, not a NamedCType.\\n# This is mostly because of the mismatch between return types and return names.\\n# e.g. a function with a return type of 'void' has 0 return names,\\n# and a function with a return type of 'std::tuple' has >1 return name.\\ndef returntype_type(t: Type, *, mutable: bool, symint: bool = False) -> CType:\\n # placeholder is ignored\\n # NB: symint is ALWAYS respected for return types. So symint argument\\n # here is IGNORED\\n r = valuetype_type(t, binds=\\\"__placeholder__\\\", mutable=mutable, symint=True)\\n if r is not None:\\n return r.type\\n\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Tensor:\\n if mutable:\\n if local.use_const_ref_for_mutable_tensors():\\n return ConstRefCType(BaseCType(tensorT))\\n else:\\n return MutRefCType(BaseCType(tensorT))\\n else:\\n # Note [Tensor Copy Returns]\\n # Currently, we use \\\"Argument.is_write\\\" to determine\\n # whether or not Tensor return types should be copies or references.\\n # If that ever changes, take a look at other locations of this note!\\n return BaseCType(tensorT)\\n elif t.name == BaseTy.Scalar:\\n return BaseCType(scalarT)\\n elif isinstance(t, ListType):\\n assert (\\n not mutable\\n ), \\\"Native functions should never return a mutable tensor list. They should return void.\\\"\\n elem = returntype_type(t.elem, mutable=False)\\n assert t.size is None, f\\\"fixed size list returns not supported: {t}\\\"\\n return VectorCType(elem)\\n elif isinstance(t, OptionalType):\\n elem = returntype_type(t.elem, mutable=mutable)\\n if str(t.elem) == \\\"Tensor\\\":\\n return OptionalCType(elem)\\n\\n raise AssertionError(f\\\"unrecognized return type {t}\\\")\\n\\n\\n# Translation of a single return to its C++ type\\ndef return_type(r: Return, *, symint: bool = False) -> CType:\\n return returntype_type(r.type, mutable=r.is_write, symint=symint)\\n\\n\\n# Translation of a full (possibly multi) return from JIT to its C++ type\\ndef returns_type(rs: Sequence[Return], *, symint: bool = False) -> CType:\\n if len(rs) == 0:\\n return BaseCType(voidT)\\n elif len(rs) == 1:\\n return return_type(rs[0], symint=symint)\\n else:\\n return TupleCType([return_type(r, symint=symint) for r in rs])\\n\\n\\ndef return_names(f: NativeFunction, *, fallback_name: str = \\\"result\\\") -> Sequence[str]:\\n returns: list[str] = []\\n for i, r in enumerate(f.func.returns):\\n # If we have an inplace function, the return argument is\\n # implicitly named self.\\n # TODO: Consider incorporating this into the data model\\n if f.func.name.name.inplace:\\n assert i == 0, \\\"illegal inplace function with multiple returns\\\"\\n name = \\\"self\\\"\\n # If we are out function, the name is the name of the\\n # corresponding output function (r.name will get recorded\\n # in field_name later.)\\n elif f.func.is_out_fn():\\n name = f.func.arguments.out[i].name\\n # If the return argument is explicitly named...\\n elif r.name:\\n name_conflict = any(\\n r.name == a.name for a in f.func.schema_order_arguments()\\n )\\n if name_conflict and not f.func.is_out_fn():\\n name = f\\\"{r.name}_return\\\"\\n else:\\n name = r.name\\n # If there is no explicit name and no fallback name was passed in, we just name the output result,\\n # unless it's a multi-return, in which case it's result0,\\n # result1, etc (zero-indexed)\\n else:\\n name = fallback_name if len(f.func.returns) == 1 else f\\\"{fallback_name}{i}\\\"\\n returns.append(name)\\n return returns\\n\\n\\nJIT_TO_CPP_DEFAULT = {\\n \\\"False\\\": \\\"false\\\",\\n \\\"True\\\": \\\"true\\\",\\n \\\"None\\\": \\\"::std::nullopt\\\", # UGH this one is type directed\\n \\\"Mean\\\": \\\"at::Reduction::Mean\\\",\\n \\\"[]\\\": \\\"{}\\\",\\n \\\"contiguous_format\\\": \\\"c10::MemoryFormat::Contiguous\\\",\\n \\\"long\\\": \\\"at::kLong\\\",\\n}\\n\\n\\n# Convert a JIT default into C++ expression representing the default\\ndef default_expr(d: str, t: Type, *, symint: bool) -> str:\\n if d == \\\"None\\\" and str(t) == \\\"Tensor?\\\":\\n return \\\"{}\\\"\\n if isinstance(t, BaseType) and t.name is BaseTy.str:\\n # Schema allows single quotes but C++ needs double\\n if len(d) >= 2 and d[0] == \\\"'\\\" and d[-1] == \\\"'\\\":\\n s = \\\"\\\"\\n i = 1\\n while i + 1 < len(d):\\n if d[i] != \\\"\\\\\\\\\\\":\\n if d[i] == '\\\"':\\n s += '\\\\\\\\\\\"'\\n else:\\n s += d[i]\\n i += 1\\n else:\\n if d[i + 1] == \\\"'\\\":\\n s += \\\"'\\\"\\n else:\\n s += d[i : i + 2]\\n i += 2\\n\\n return f'\\\"{s}\\\"'\\n\\n if isinstance(t, OptionalType):\\n if d == \\\"None\\\":\\n return \\\"::std::nullopt\\\"\\n\\n return default_expr(d, t.elem, symint=symint)\\n\\n if isinstance(t, ListType):\\n if d.startswith(\\\"[\\\") and d.endswith(\\\"]\\\"):\\n return \\\"{\\\" + d[1:-1] + \\\"}\\\"\\n elif symint and d.isdigit() and str(t.elem) == \\\"SymInt\\\":\\n return f\\\"c10::SymInt({d})\\\"\\n elif t.size is None:\\n # NOTE: Sized lists can have scalar defaults\\n raise ValueError(f\\\"Expected a list default '[...]' but found: '{d}'\\\")\\n\\n return JIT_TO_CPP_DEFAULT.get(d, d)\\n\\n\\n# Convert an argument into its C++ API form\\n\\n\\ndef argument(\\n a: Argument | TensorOptionsArguments | SelfArgument,\\n *,\\n cpp_no_default_args: set[str],\\n method: bool,\\n faithful: bool,\\n symint: bool = False,\\n has_tensor_options: bool,\\n) -> list[Binding]:\\n def sub_argument(\\n a: Argument | TensorOptionsArguments | SelfArgument,\\n ) -> list[Binding]:\\n return argument(\\n a,\\n cpp_no_default_args=cpp_no_default_args,\\n method=method,\\n faithful=faithful,\\n symint=symint,\\n has_tensor_options=has_tensor_options,\\n )\\n\\n if isinstance(a, Argument):\\n binds: ArgName\\n if a.name == \\\"memory_format\\\" and has_tensor_options:\\n binds = SpecialArgName.possibly_redundant_memory_format\\n else:\\n binds = a.name\\n default: str | None = None\\n if a.name not in cpp_no_default_args and a.default is not None:\\n default = default_expr(a.default, a.type, symint=symint)\\n return [\\n Binding(\\n nctype=argument_type(a, binds=binds, symint=symint),\\n name=a.name,\\n default=default,\\n argument=a,\\n )\\n ]\\n elif isinstance(a, TensorOptionsArguments):\\n if faithful:\\n return (\\n sub_argument(a.dtype)\\n + sub_argument(a.layout)\\n + sub_argument(a.device)\\n + sub_argument(a.pin_memory)\\n )\\n else:\\n default = None\\n # Enforced by NativeFunction.__post_init__\\n assert \\\"options\\\" not in cpp_no_default_args\\n if all(x.default == \\\"None\\\" for x in a.all()):\\n default = \\\"{}\\\"\\n elif a.dtype.default == \\\"long\\\":\\n default = \\\"at::kLong\\\" # TODO: this is wrong\\n return [\\n Binding(\\n nctype=NamedCType(\\\"options\\\", BaseCType(tensorOptionsT)),\\n name=\\\"options\\\",\\n default=default,\\n argument=a,\\n )\\n ]\\n elif isinstance(a, SelfArgument):\\n if method:\\n # Caller is responsible for installing implicit this in context!\\n return []\\n else:\\n return sub_argument(a.argument)\\n else:\\n assert_never(a)\\n\\n\\ndef arguments(\\n arguments: Arguments,\\n *,\\n faithful: bool,\\n symint: bool = False,\\n method: bool,\\n cpp_no_default_args: set[str],\\n) -> list[Binding]:\\n args: list[Argument | TensorOptionsArguments | SelfArgument] = []\\n if faithful:\\n args.extend(arguments.non_out)\\n args.extend(arguments.out)\\n else:\\n args.extend(arguments.out)\\n args.extend(arguments.non_out)\\n return [\\n r.no_default() if faithful else r\\n for a in args\\n for r in argument(\\n a,\\n faithful=faithful,\\n symint=symint,\\n method=method,\\n has_tensor_options=arguments.tensor_options is not None,\\n cpp_no_default_args=cpp_no_default_args,\\n )\\n ]\\n\\n\\nfrom __future__ import annotations\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.types import Binding, CppSignatureGroup, CType\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n ListType,\\n NativeFunction,\\n OptionalType,\\n Type,\\n)\\n\\n\\n# This file generates the code for unboxing wrappers, i.e., the glue logic to unbox a boxed operator and convert the\\n# ivalues from stack to correct arguments to the unboxed kernel, based on corresponding JIT schema. This codegen is\\n# an alternative way to generate unboxing wrappers similar to the existing C++ metaprogramming approach but gets the\\n# job done statically. These generated unboxing wrappers will be useful under the scenario where we need to register\\n# a fixed set of operators known at compile time and thus can save some time in runtime initialization phase.\\n#\\n# Here's an example on how the codegen works:\\n#\\n# - Function Schema (source of truth)\\n#\\n# aten::empty.names(int[] size, *, Dimname[]? names,\\n# ScalarType? dtype=None, Layout? layout=None,\\n# Device? device=None, bool? pin_memory=None,\\n# MemoryFormat? memory_format=None) -> Tensor\\n# - Argument Conversion\\n# Generates C++ code to convert an ivalue (from stack) to its underlying C++ type.\\n# - int[] size\\n# ```cpp\\n# const c10::List size_list_in = (std::move(peek(stack, 0, 7))).toList();\\n#\\n# std::vector size_vec;\\n# for (c10::IValue size_elem: size_list_in) {\\n# int64_t size_base = size_elem.to();\\n# size_vec.push_back(size_base);\\n# }\\n# at::ArrayRef size_list_out(size_vec);\\n# ~~~~~~~~~~~~~ <-- The converted argument from ivalues in the stack.\\n# Will be passed to unboxed kernel.\\n# ```\\n# - Dimname[]? names\\n# ```cpp\\n# ::std::optional names_opt = (std::move(peek(stack, 1, 7))).toOptional();\\n# ::std::optional> names_opt_out;\\n# if (names_opt.has_value()) {\\n# ~~~~~~~~~~~ <-- Unwrapping optional shell\\n# const c10::IValue names_opt_in = names_opt.value();\\n# const c10::List names_list_in = names_opt_in.toList();\\n#\\n# std::vector names_vec;\\n# for (c10::IValue names_elem: names_list_in) {\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~ <-- Unrolling list, then convert elements one by one.\\n# at::Dimname names_base = names_elem.to();\\n# names_vec.push_back(names_base);\\n# }\\n# at::ArrayRef names_list_out(names_vec);\\n#\\n# names_opt_out = ::std::optional>(names_list_out);\\n# } else {\\n# names_opt_out = ::std::optional>();\\n# }\\n# ```\\n# - ScalarType? dtype (similarly for the rest of the arguments)\\n# ```cpp\\n# ::std::optional dtype_opt = (std::move(peek(stack, 2, 7))).toOptional();\\n# ::std::optional dtype_opt_out;\\n# if (dtype_opt.has_value()) {\\n# const c10::IValue dtype_opt_in = dtype_opt.value();\\n# at::ScalarType dtype_base = dtype_opt_in.to();\\n# ~~~~~~~~~~~~~~~~~~~~ <-- For base types, convert ivalue to it\\n# directly using \\\".to()\\\" API.\\n# dtype_opt_out = ::std::optional(dtype_base);\\n# } else {\\n# dtype_opt_out = ::std::optional();\\n# }\\n# ```\\n#\\n# - Unboxed Kernel Call\\n# ```cpp\\n# auto result_ = torch::empty(\\n# size_list_out,\\n# names_opt_out,\\n# options,\\n# memory_format_opt_out\\n# );\\n# ```\\n#\\n# - Push Result Back to Stack\\n# ```cpp\\n# drop(stack, 7);\\n# pack(stack, std::move(result_));\\n# ```\\nconnector = \\\"\\\\n\\\\t\\\"\\n\\n\\n# Return unboxing function name for a NativeFunction\\ndef name(f: NativeFunction) -> str:\\n return f.func.name.unambiguous_name()\\n\\n\\n# Convert all the arguments in a NativeFunction to C++ code\\ndef convert_arguments(f: NativeFunction) -> tuple[list[Binding], list[str]]:\\n # we need the 'self' argument so method needs to be False\\n args = (\\n CppSignatureGroup.from_native_function(f, method=False)\\n .most_faithful_signature()\\n .arguments()\\n )\\n code_list = [\\n f\\\"c10::IValue {args[i].name} = std::move(peek(stack, {i}, {len(args)}));\\\"\\n for i in range(len(args))\\n ] + [\\\"\\\"]\\n binding_list = []\\n for arg in args:\\n # expecting only Argument\\n if not isinstance(arg.argument, Argument):\\n raise Exception( # noqa: TRY002\\n f\\\"Unexpected argument type, expecting `Argument` but got {arg}\\\"\\n )\\n argument: Argument = arg.argument\\n unboxed_name, _, code, decl = argumenttype_ivalue_convert(\\n argument.type,\\n argument.name,\\n mutable=argument.is_write,\\n )\\n code_list.extend(decl)\\n code_list.extend(code)\\n binding_list.append(arg.with_name(unboxed_name))\\n return binding_list, code_list\\n\\n\\n# Takes in the type, name and mutability corresponding to an argument, and generates a tuple of:\\n# (1) the C++ code necessary to unbox the argument\\n# (2) A Binding corresponding to the newly created unboxed variable, including variable name and its CType\\ndef argumenttype_ivalue_convert(\\n t: Type, arg_name: str, *, mutable: bool = False\\n) -> tuple[str, CType, list[str], list[str]]:\\n # Unboxing is for mobile, which doesn't care about SymInts\\n ctype = cpp.argumenttype_type(\\n t=t, mutable=mutable, binds=arg_name, symint=False\\n ).type\\n\\n if isinstance(t, BaseType):\\n out_name = f\\\"{arg_name}_base\\\"\\n code, decl = _gen_code_base_type(\\n arg_name=arg_name, out_name=out_name, ctype=ctype\\n )\\n elif isinstance(t, OptionalType):\\n out_name = f\\\"{arg_name}_opt_out\\\"\\n code, decl = _gen_code_optional_type(\\n arg_name=arg_name,\\n out_name=out_name,\\n t=t,\\n ctype=ctype,\\n )\\n elif isinstance(t, ListType):\\n out_name = f\\\"{arg_name}_list_out\\\"\\n code, decl = _gen_code_list_type(\\n arg_name=arg_name,\\n out_name=out_name,\\n t=t,\\n ctype=ctype,\\n )\\n else:\\n raise Exception(f\\\"Cannot handle type {t}. arg_name: {arg_name}\\\") # noqa: TRY002\\n return out_name, ctype, code, decl\\n\\n\\ndef _gen_code_base_type(\\n arg_name: str, out_name: str, ctype: CType\\n) -> tuple[list[str], list[str]]:\\n return [\\n f\\\"{ctype.cpp_type(strip_ref=True)} {out_name} = {arg_name}.to<{ctype.cpp_type(strip_ref=True)}>();\\\"\\n ], []\\n\\n\\ndef _gen_code_optional_type(\\n arg_name: str, out_name: str, t: OptionalType, ctype: CType\\n) -> tuple[list[str], list[str]]:\\n in_name = f\\\"{arg_name}_opt_in\\\"\\n res_name, _, res_code, decl = argumenttype_ivalue_convert(t.elem, in_name)\\n return (\\n f\\\"\\\"\\\"\\nauto {arg_name}_opt = {arg_name}.toOptional();\\n{ctype.cpp_type(strip_ref=True)} {out_name};\\nif ({arg_name}_opt.has_value()) {{\\n const c10::IValue {in_name} = {arg_name}_opt.value();\\n {connector.join(res_code)}\\n {out_name} = {ctype.cpp_type(strip_ref=True)}({res_name});\\n}} else {{\\n {out_name} = {ctype.cpp_type(strip_ref=True)}();\\n}}\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n ),\\n decl,\\n )\\n\\n\\ndef _gen_code_list_type(\\n arg_name: str, out_name: str, t: ListType, ctype: CType\\n) -> tuple[list[str], list[str]]:\\n in_name = f\\\"{arg_name}_list_in\\\"\\n elem_name = f\\\"{arg_name}_elem\\\"\\n code = [f\\\"const c10::List {in_name} = {arg_name}.toList();\\\"]\\n res_name, res_ctype, res_code, decl = argumenttype_ivalue_convert(t.elem, elem_name)\\n # handle list type with size, e.g., bool[4]\\n if isinstance(t.elem, BaseType) and t.elem.name == BaseTy.bool and t.size:\\n code.extend(\\n f\\\"\\\"\\\"\\n{ctype.cpp_type(strip_ref=True)} {out_name} = as_array<{res_ctype.cpp_type(strip_ref=True)}, {t.size}>({in_name});\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n # we have to use c10::List for optional element. e.g., Tensor?[] -> c10::List<::std::optional>\\n elif isinstance(t.elem, OptionalType):\\n code.extend(\\n f\\\"\\\"\\\"\\n{ctype.cpp_type(strip_ref=True)} {out_name};\\nfor (c10::IValue {elem_name}: {in_name}) {{\\n {connector.join(res_code)}\\n {out_name}.push_back({res_name});\\n}}\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n else:\\n # use ArrayRef as default.\\n vec_name = arg_name + \\\"_vec\\\"\\n # need to bring vector instantiation out of scope so that ArrayRef has valid data\\n decl.append(f\\\"std::vector<{res_ctype.cpp_type(strip_ref=True)}> {vec_name};\\\")\\n code.extend(\\n f\\\"\\\"\\\"\\nfor (c10::IValue {elem_name}: {in_name}) {{\\n {connector.join(res_code)}\\n {vec_name}.push_back({res_name});\\n}}\\n{ctype.cpp_type(strip_ref=True)} {out_name}({vec_name});\\n \\\"\\\"\\\".split(\\n \\\"\\\\n\\\"\\n )\\n )\\n return code, decl\\n\\n\\nfrom torchgen.model import NativeFunctionsGroup\\n\\n\\n# Follows dispatcher calling convention, but:\\n# - Mutable arguments not allowed. Meta functions are always\\n# written in functional form. Look at FunctionSchema.signature()\\n# - No tensor returns; instead we return a TensorMeta describing\\n# the tensor in question\\n\\n\\ndef name(g: NativeFunctionsGroup) -> str:\\n # use the overload name from the functional version\\n return str(g.functional.func.name).replace(\\\".\\\", \\\"_\\\")\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\n\\nimport torchgen.api.types as api_types\\nfrom torchgen.api import cpp, structured\\nfrom torchgen.api.types import (\\n ArgName,\\n BaseCppType,\\n BaseCType,\\n Binding,\\n ConstRefCType,\\n CType,\\n NamedCType,\\n scalarT,\\n)\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n DispatchKey,\\n FunctionSchema,\\n NativeFunctionsGroup,\\n Type,\\n)\\n\\n\\ndef schema_kernel_name(func: FunctionSchema, dispatch_key: DispatchKey) -> str:\\n assert func.is_out_fn(), \\\"ufunc.kernel_name should only be invoked on out schemas\\\"\\n return f\\\"ufunc_{func.name.name}_{dispatch_key}\\\"\\n\\n\\ndef kernel_name(g: NativeFunctionsGroup, dispatch_key: DispatchKey) -> str:\\n return schema_kernel_name(g.out.func, dispatch_key)\\n\\n\\n# Tensors are omitted (as they are stored in TensorIterator), everything else is\\n# passed along (technically, we can pass tensors along too, it just wastes\\n# argument registers)\\n#\\n# NB: used for CPU only\\ndef dispatchstub_type(t: Type, *, binds: ArgName) -> NamedCType | None:\\n # Dispatch stubs are always plain ints\\n r = cpp.valuetype_type(t, binds=binds, symint=False)\\n if r is not None:\\n return r\\n\\n if t == BaseType(BaseTy.Scalar):\\n return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))\\n elif t == BaseType(BaseTy.Tensor):\\n return None\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\ndef opmath_type(scalar_t: BaseCppType) -> BaseCppType:\\n if scalar_t == api_types.scalar_t:\\n return api_types.opmath_t\\n raise NotImplementedError\\n\\n\\n# NB: Tensors in constructor are stored in opmath_t, not scalar_t\\n# because Tensor in constructor = its a scalar tensor partially applied =\\n# it can be higher precision and we want to compute in that higher precision\\n#\\n# NB: CUDA only\\ndef ufunctor_ctor_type(t: Type, *, binds: ArgName, scalar_t: BaseCppType) -> NamedCType:\\n r = cpp.valuetype_type(t, binds=binds, symint=False)\\n if r is not None:\\n return r\\n\\n if t == BaseType(BaseTy.Scalar):\\n return NamedCType(binds, BaseCType(opmath_type(scalar_t)))\\n elif t == BaseType(BaseTy.Tensor):\\n return NamedCType(binds, BaseCType(opmath_type(scalar_t)))\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\n# Only Tensors ever get passed directly to operator()\\n#\\n# NB: CUDA only\\n# (Actually, this works for CPU too)\\ndef ufunctor_apply_type(\\n t: Type, *, binds: ArgName, scalar_t: BaseCppType\\n) -> NamedCType:\\n if t == BaseType(BaseTy.Tensor):\\n return NamedCType(binds, BaseCType(scalar_t))\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\n# The actual ufunc template function the user writes. Everything here\\n# is done in the computation type. compute_t is opmath_t in CUDA and scalar_t\\n# in CPU\\ndef ufunc_type(t: Type, *, binds: ArgName, compute_t: CType) -> NamedCType:\\n r = cpp.valuetype_type(t, binds=binds, symint=False)\\n if r is not None:\\n return r\\n\\n if t == BaseType(BaseTy.Scalar):\\n return NamedCType(binds, compute_t)\\n elif t == BaseType(BaseTy.Tensor):\\n return NamedCType(binds, compute_t)\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\ndef ufunctor_ctor_argument(a: Argument, scalar_t: BaseCppType) -> Binding:\\n return Binding(\\n nctype=ufunctor_ctor_type(a.type, binds=a.name, scalar_t=scalar_t),\\n name=a.name,\\n default=None,\\n argument=a,\\n )\\n\\n\\ndef ufunctor_apply_argument(a: Argument, scalar_t: BaseCppType) -> Binding:\\n return Binding(\\n nctype=ufunctor_apply_type(a.type, binds=a.name, scalar_t=scalar_t),\\n name=a.name,\\n default=None,\\n argument=a,\\n )\\n\\n\\ndef ufunc_argument(a: Argument, compute_t: CType) -> Binding:\\n return Binding(\\n nctype=ufunc_type(a.type, binds=a.name, compute_t=compute_t),\\n name=a.name,\\n default=None,\\n argument=a,\\n )\\n\\n\\n@dataclass(frozen=True)\\nclass UfunctorBindings:\\n ctor: list[Binding]\\n apply: list[Binding]\\n\\n\\n# ufunctors are a CUDA-only concept representing functors that take some of\\n# their arguments on a host-side constructor, and the rest in the device-side\\n# apply. E.g.,\\n#\\n# template \\n# struct CUDAFunctorOnSelf_add {\\n# using opmath_t = at::opmath_type;\\n# opmath_t other_;\\n# opmath_t alpha_;\\n# CUDAFunctorOnSelf_add(opmath_t other, opmath_t alpha) : other_(other), alpha_(alpha) {}\\n# __device__ scalar_t operator()(scalar_t self) {\\n# return ufunc::add(static_cast(self), other_, alpha_);\\n# }\\n# };\\n#\\n# The ctor refers to the constructor CUDAFunctorOnSelf_add, while apply refers\\n# to the operator() definition\\ndef ufunctor_arguments(\\n g: NativeFunctionsGroup, *, scalar_tensor_idx: int | None, scalar_t: BaseCppType\\n) -> UfunctorBindings:\\n ctor = []\\n apply = []\\n for a in g.functional.func.arguments.flat_non_out:\\n if a.type.is_tensor_like():\\n if scalar_tensor_idx == 0:\\n # put it in the ctor anyway\\n ctor.append(ufunctor_ctor_argument(a, scalar_t=scalar_t))\\n scalar_tensor_idx = None\\n else:\\n if scalar_tensor_idx is not None:\\n scalar_tensor_idx -= 1\\n apply.append(ufunctor_apply_argument(a, scalar_t=scalar_t))\\n else:\\n ctor.append(ufunctor_ctor_argument(a, scalar_t=scalar_t))\\n assert scalar_tensor_idx is None\\n return UfunctorBindings(ctor=ctor, apply=apply)\\n\\n\\n# ufuncs are the inner loop template functions that you wrote in ufunc/add.h\\n# which do the actual computation in question. E.g.,\\n#\\n# template \\n# C10_HOST_DEVICE T add(T self, T other, T alpha) __ubsan_ignore_undefined__ {\\n# return self + alpha * other;\\n# }\\n#\\n# In this file, we refer to T as compute_t which is bound by caller\\ndef ufunc_arguments(g: NativeFunctionsGroup, *, compute_t: CType) -> list[Binding]:\\n return [\\n ufunc_argument(a, compute_t=compute_t)\\n for a in g.functional.func.arguments.flat_non_out\\n ]\\n\\n\\n# Stubs are the DispatchStub trampolines that CPU kernels use to get to their\\n# vectorized versions. E.g.,\\n#\\n# using structured_binary_fn_alpha = void(*)(TensorIteratorBase&, const Scalar& alpha);\\n# DECLARE_DISPATCH(structured_binary_fn_alpha, add_stub);\\ndef stub_arguments(g: NativeFunctionsGroup) -> list[Binding]:\\n # stubs drop all tensor arguments (they are implicit in the TensorIterator\\n # argument and keep everything else)\\n return [\\n r\\n for a in g.out.func.arguments.flat_non_out\\n if not a.type.is_tensor_like()\\n for r in structured.argument(a)\\n ]\\n\\n\\nfrom __future__ import annotations\\n\\nfrom typing import NoReturn, Sequence\\n\\nfrom torchgen.api.types import (\\n ArrayRefCType,\\n BaseCType,\\n Binding,\\n boolT,\\n ConstRefCType,\\n deviceT,\\n Expr,\\n intArrayRefT,\\n iOptTensorListRefT,\\n layoutT,\\n ListCType,\\n longT,\\n memoryFormatT,\\n MutRefCType,\\n NamedCType,\\n opmath_t,\\n OptionalCType,\\n optionalIntArrayRefT,\\n optionalScalarRefT,\\n optionalSymIntArrayRefT,\\n optionalTensorRefT,\\n scalar_t,\\n scalarT,\\n scalarTypeT,\\n SpecialArgName,\\n symIntArrayRefT,\\n SymIntT,\\n tensorOptionsT,\\n tensorT,\\n VectorCType,\\n)\\n\\n\\n# This file implements a small program synthesis engine that implements\\n# conversions between one API to another.\\n#\\n# The key data type in this file in NamedCType, short for Named C++ semantic type. A NamedCType\\n# represents a C++ type, plus semantic information about what it represents.\\n# For example, consider the argument \\\"bool pin_memory\\\"; its normal C++ type is\\n# \\\"bool\\\", but its C++ semantic type also keeps track that this represents a\\n# \\\"pin_memory\\\"; you can't just use a random other boolean in a context where you\\n# need a \\\"pin_memory\\\"!\\n#\\n# The translator takes a list of needed NamedCTypes, and then figures out how\\n# to construct expressions with these NamedCTypes from the given bindings. Many\\n# of these expressions are trivial (I need a Tensor other; there's a Tensor\\n# other scope); others are more nontrivial and may require packing/unpacking.\\n# Some examples of non-trivial action:\\n#\\n# - Need the \\\"dtype\\\" binding? Well, maybe \\\"dtype\\\" isn't available\\n# in the context, instead, \\\"options\\\" is, and you need to extract\\n# it from there. (Gather)\\n#\\n# - Need the \\\"context\\\" binding? Well, maybe \\\"context\\\" isn't available\\n# in the context, and you need to construct it from \\\"dtype\\\", \\\"device\\\",\\n# etc. (Scatter)\\n#\\n# - Need the \\\"memory_format\\\" binding? Well, actually, it's available\\n# from both \\\"memory_format\\\" and \\\"options\\\", so you had better make sure\\n# they are consistent. (Join)\\n\\noptions_ctype = NamedCType(\\\"options\\\", ConstRefCType(BaseCType(tensorOptionsT)))\\n\\nout_tensor_ctype = NamedCType(\\\"out\\\", ConstRefCType(BaseCType(tensorT)))\\n\\nlongVec_ctype = VectorCType(BaseCType(longT))\\nlongSymVec_ctype = VectorCType(BaseCType(SymIntT))\\noptionalLongVec_ctype = OptionalCType(VectorCType(BaseCType(longT)))\\noptionalScalar_ctype = OptionalCType(BaseCType(scalarT))\\noptionalTensor_ctype = OptionalCType(BaseCType(tensorT))\\n\\n\\nclass UnsatError(RuntimeError):\\n pass\\n\\n\\n# Given a set of in-scope bindings and a set of target bindings, synthesize\\n# a list of expressions that uses only the in-scope bindings (bindings) that\\n# have all of the types of goals. You may want to use this function if\\n# you're generating code for a function like:\\n#\\n# void f({args}) {\\n# g({exprs}); // g is a different API\\n# }\\n#\\n# and you need to generate \\\"exprs\\\".\\n#\\n# Typically, a list of Bindings is convenient to get (you usually call something\\n# like arguments() to get them); but technically you only need less information:\\n# for 'bindings' an (un-ordered) list of Exprs is sufficient; similarly, for\\n# 'goals', an (ordered) list of NamedCType goals is sufficient. If you are doing\\n# something more complicated, e.g., tracking the set of bindings in a context,\\n# you may find using these smaller types more convenient.\\ndef translate(\\n bindings: Sequence[Expr | Binding],\\n goals: Sequence[NamedCType | Binding],\\n *,\\n method: bool = False,\\n allow_expensive_conversions: bool = False,\\n) -> list[Expr]:\\n binding_exprs: list[Expr] = []\\n for b in bindings:\\n if isinstance(b, Binding):\\n binding_exprs.append(\\n Expr(\\n expr=b.name,\\n type=b.nctype,\\n )\\n )\\n else:\\n binding_exprs.append(b)\\n\\n goal_ctypes: list[NamedCType] = []\\n for g in goals:\\n if isinstance(g, Binding):\\n goal_ctypes.append(g.nctype)\\n else:\\n goal_ctypes.append(g)\\n\\n # Add all the bindings to the context\\n ctx: dict[NamedCType, str] = {}\\n for b in binding_exprs:\\n ctx[b.type] = b.expr\\n\\n # While we're at it, do some simple forward inference, looking through\\n # constructors.\\n #\\n # NB: When should you do forward inference versus backward inference?\\n # The general idea:\\n #\\n # - Backward inference WHEN the goal gets smaller\\n # - Forward inference WHEN the hypothesis gets smaller\\n #\\n # This helps ensure termination: backward inference starts with a goal\\n # and tries to make it simpler and simpler until it's trivial; if the\\n # goal can grow in size, we blow up to a really huge goal size.\\n # Similarly, with forward inference we take hypotheses and decompose\\n # them into simpler hypotheses; if hypotheses could expand in size,\\n # we also have potential nontermination. (In the code below, forward\\n # inference is only ever carried out at a single step, but you could\\n # imagine repeated application of forward inference being profitable.)\\n #\\n # A good starting point in the literature for exploring more about proof\\n # search are these lecture notes\\n # https://www.cs.cmu.edu/~fp/courses/oregon-m10/04-focusing.pdf\\n #\\n # TODO: My kingdom for a pattern matcher\\n # https://www.python.org/dev/peps/pep-0634/\\n #\\n # TODO: This could get us in recomputation trouble if b.expr is nontrivial.\\n # Fix this by implementing some sort of sharing so that if multiple\\n # goals share the same expression, we only compute it once. This seems\\n # to matter in practice as compiler is often unwilling to CSE nontrivial\\n # expressions like scalar.to()\\n t = b.type\\n if (\\n isinstance(t, ConstRefCType)\\n and isinstance(t.elem, OptionalCType)\\n and isinstance(t.elem.elem, BaseCType)\\n and str(t.elem.elem.type) == \\\"at::Tensor\\\"\\n ):\\n ctx[\\n NamedCType(t.elem.elem.name, ConstRefCType(BaseCType(tensorT)))\\n ] = f\\\"({b.expr}.has_value() ? *{b.expr} : at::Tensor())\\\"\\n\\n if t.type == ConstRefCType(OptionalCType(BaseCType(tensorT))):\\n ctx[\\n NamedCType(t.name, BaseCType(optionalTensorRefT))\\n ] = f\\\"(({b.expr}.has_value() && (*{b.expr}).defined()) ? at::OptionalTensorRef(*{b.expr}) : at::OptionalTensorRef())\\\"\\n\\n if t.type == ConstRefCType(BaseCType(scalarT)):\\n ctx[NamedCType(t.name, BaseCType(opmath_t))] = f\\\"({b.expr}).to()\\\"\\n\\n if t.type == ConstRefCType(OptionalCType(BaseCType(scalarT))):\\n ctx[\\n NamedCType(t.name, BaseCType(optionalScalarRefT))\\n ] = f\\\"({b.expr}.has_value() ? at::OptionalScalarRef(&({b.expr}.value())) : at::OptionalScalarRef())\\\"\\n\\n if t.type == BaseCType(scalar_t):\\n ctx[\\n NamedCType(t.name, BaseCType(opmath_t))\\n ] = f\\\"static_cast({b.expr})\\\"\\n\\n # [Note: IOptTensorListRef]\\n if t.type == ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT)))):\\n ctx[\\n NamedCType(t.name, BaseCType(iOptTensorListRefT))\\n ] = f\\\"at::IOptTensorListRef({b.expr})\\\"\\n\\n # Add implicit bindings if the generated code is inside a Tensor method\\n if method:\\n ctx[\\n NamedCType(\\\"self\\\", MutRefCType(BaseCType(tensorT)))\\n ] = \\\"const_cast(*this)\\\"\\n ctx[\\n NamedCType(\\\"self\\\", ConstRefCType(BaseCType(tensorT)))\\n ] = \\\"const_cast(*this)\\\"\\n # This is better! Byte-for-byte compat\\n # ctx[NamedCType(\\\"self\\\", ConstRefCType(BaseCType(tensorT)))] = \\\"*this\\\"\\n\\n def unsat(goal: NamedCType) -> NoReturn:\\n ctx_desc = \\\"\\\\n\\\".join(\\n f\\\" {t.cpp_type()} {t.name}; // {e}\\\" for t, e in ctx.items()\\n )\\n raise UnsatError(\\n f\\\"\\\"\\\"\\nFailed to synthesize the expression \\\"{goal.cpp_type()} {goal.name}\\\".\\nWhen I failed, the following bindings were available in the context:\\n\\n{ctx_desc}\\n\\nThis probably means there is a missing rule in the rules of torchgen.api.translate.\\nCheck this module for more information.\\n\\\"\\\"\\\"\\n )\\n\\n # A shitty backtracking search implementation. It's shitty because it\\n # does backtracking via stack (bad idea!) and for the most part tries to\\n # avoid backtracking. In particular, if\\n # direct=True, we won't try to do any fancy synthesis, just trivial\\n # conversions (e.g., \\\"T a\\\" is OK for \\\"const T& a\\\"). So all of the\\n # existing rules in this function simply try to solve immediately,\\n # and bail if things don't work out.\\n def solve(goal: NamedCType, *, direct: bool) -> str:\\n def direct_solve(goal: NamedCType) -> str:\\n return solve(goal, direct=True)\\n\\n if goal in ctx:\\n # Trivial\\n return ctx[goal]\\n\\n # const & is satisfied with mutable &\\n if isinstance(goal.type, ConstRefCType):\\n try:\\n # WARNING: not strictly decreasing; be careful not\\n # to add a direct conversion that goes satisfies\\n # mutable& with const&\\n return solve(\\n NamedCType(goal.name, MutRefCType(goal.type.elem)), direct=direct\\n )\\n except UnsatError:\\n pass\\n\\n # mutable & is satisfied with value\\n if isinstance(goal.type, MutRefCType):\\n try:\\n return solve(NamedCType(goal.name, goal.type.elem), direct=direct)\\n except UnsatError:\\n pass\\n\\n # TODO: These are referentially equal, shouldn't have to do this;\\n # ensuring we don't use type synonym IntArrayRef in codegen would\\n # help\\n if goal.type == ArrayRefCType(BaseCType(longT)):\\n return solve(NamedCType(goal.name, BaseCType(intArrayRefT)), direct=direct)\\n\\n if direct:\\n unsat(goal)\\n\\n # For now, all of these rules are mutually exclusive.\\n if goal == NamedCType(\\\"memory_format\\\", OptionalCType(BaseCType(memoryFormatT))):\\n memory_format = direct_solve(\\n NamedCType(\\n SpecialArgName.possibly_redundant_memory_format,\\n OptionalCType(BaseCType(memoryFormatT)),\\n )\\n )\\n # No need to join \\\"memory_format\\\" and \\\"options\\\" if the target API takes \\\"options\\\" directly.\\n # Otherwise it will cause the redundant memory_format error.\\n if options_ctype in goal_ctypes:\\n return memory_format\\n try:\\n options = direct_solve(options_ctype)\\n return f\\\"c10::impl::check_tensor_options_and_extract_memory_format({options}, {memory_format})\\\"\\n except UnsatError:\\n return memory_format\\n elif goal == NamedCType(\\\"options\\\", BaseCType(tensorOptionsT)):\\n dtype = direct_solve(\\n NamedCType(\\\"dtype\\\", OptionalCType(BaseCType(scalarTypeT)))\\n )\\n pin_memory = direct_solve(\\n NamedCType(\\\"pin_memory\\\", OptionalCType(BaseCType(boolT)))\\n )\\n device = direct_solve(\\n NamedCType(\\\"device\\\", OptionalCType(BaseCType(deviceT)))\\n )\\n layout = direct_solve(\\n NamedCType(\\\"layout\\\", OptionalCType(BaseCType(layoutT)))\\n )\\n return f\\\"TensorOptions().dtype({dtype}).layout({layout}).device({device}).pinned_memory({pin_memory})\\\"\\n\\n elif goal == NamedCType(\\\"dtype\\\", OptionalCType(BaseCType(scalarTypeT))):\\n try:\\n options = direct_solve(options_ctype)\\n return f\\\"c10::optTypeMetaToScalarType({options}.dtype_opt())\\\"\\n except UnsatError:\\n out_tensor = direct_solve(out_tensor_ctype)\\n return f\\\"{out_tensor}.scalar_type()\\\"\\n\\n elif goal == NamedCType(\\\"layout\\\", OptionalCType(BaseCType(layoutT))):\\n try:\\n options = direct_solve(options_ctype)\\n return f\\\"{options}.layout_opt()\\\"\\n except UnsatError:\\n out_tensor = direct_solve(out_tensor_ctype)\\n return f\\\"{out_tensor}.layout()\\\"\\n\\n elif goal == NamedCType(\\\"device\\\", OptionalCType(BaseCType(deviceT))):\\n try:\\n options = direct_solve(options_ctype)\\n return f\\\"{options}.device_opt()\\\"\\n except UnsatError:\\n out_tensor = direct_solve(out_tensor_ctype)\\n return f\\\"{out_tensor}.device()\\\"\\n\\n elif goal == NamedCType(\\\"pin_memory\\\", OptionalCType(BaseCType(boolT))):\\n try:\\n options = direct_solve(options_ctype)\\n return f\\\"{options}.pinned_memory_opt()\\\"\\n except UnsatError:\\n # If we're calling a factory op from its out= variant,\\n # We don't actually care about the value of pin_memory.\\n out_tensor = direct_solve(out_tensor_ctype)\\n return \\\"::std::nullopt\\\"\\n\\n # We can always do translations from value types to reference types, like vector -> IntArrayRef\\n elif goal.type == BaseCType(intArrayRefT):\\n try:\\n return direct_solve(NamedCType(goal.name, longVec_ctype))\\n except UnsatError:\\n # We can also go SymIntArrayRef -> IntArrayRef\\n symIntArrayRef_type = direct_solve(\\n NamedCType(goal.name, BaseCType(symIntArrayRefT))\\n )\\n return f\\\"C10_AS_INTARRAYREF_SLOW({symIntArrayRef_type})\\\"\\n elif goal.type == BaseCType(symIntArrayRefT):\\n try:\\n r = direct_solve(NamedCType(goal.name, BaseCType(intArrayRefT)))\\n return f\\\"c10::fromIntArrayRefSlow({r})\\\"\\n except UnsatError:\\n return direct_solve(NamedCType(goal.name, longSymVec_ctype))\\n elif goal.type == BaseCType(SymIntT):\\n return direct_solve(NamedCType(goal.name, BaseCType(longT)))\\n elif goal.type == OptionalCType(BaseCType(SymIntT)):\\n argname = direct_solve(\\n NamedCType(goal.name, OptionalCType(BaseCType(longT)))\\n )\\n return f\\\"{argname}.has_value() ? ::std::make_optional(c10::SymInt(*{argname})) : ::std::nullopt\\\"\\n elif goal.type == BaseCType(longT):\\n symInt_type = direct_solve(NamedCType(goal.name, BaseCType(SymIntT)))\\n return f\\\"{symInt_type}.guard_int(__FILE__, __LINE__)\\\"\\n elif goal.type == OptionalCType(BaseCType(longT)):\\n argname = direct_solve(\\n NamedCType(goal.name, OptionalCType(BaseCType(SymIntT)))\\n )\\n return f\\\"{argname}.has_value() ? ::std::make_optional({argname}->guard_int(__FILE__, __LINE__)) : ::std::nullopt\\\"\\n elif goal.type == BaseCType(optionalIntArrayRefT):\\n try:\\n return direct_solve(NamedCType(goal.name, optionalLongVec_ctype))\\n except UnsatError:\\n argname = direct_solve(\\n NamedCType(goal.name, BaseCType(optionalSymIntArrayRefT))\\n )\\n return f\\\"{argname}.has_value() ? ::std::make_optional(C10_AS_INTARRAYREF_SLOW(*{argname})) : ::std::nullopt\\\"\\n elif goal.type == BaseCType(optionalSymIntArrayRefT):\\n # TODO: You might also want to solve this from longSymVec_ctype or\\n # an optional version of it\\n argname = direct_solve(\\n NamedCType(goal.name, BaseCType(optionalIntArrayRefT))\\n )\\n return f\\\"{argname}.has_value() ? ::std::make_optional(c10::fromIntArrayRefSlow(*{argname})) : ::std::nullopt\\\"\\n elif goal.type == BaseCType(optionalScalarRefT):\\n return direct_solve(NamedCType(goal.name, optionalScalar_ctype))\\n elif goal.type == BaseCType(optionalTensorRefT):\\n return direct_solve(NamedCType(goal.name, optionalTensor_ctype))\\n\\n # Note [translation from C++ reference to value types]\\n # The below cases are all for when we have an argument with a reference type,\\n # and a corresponding goal with a value type.\\n # These are needed when we populate the inputs to a lambda capture and we need\\n # to guarantee the lifetime of each captured argument.\\n # We guard it with an explicit kwarg because converting to a value type is expensive\\n # (O(n)) to convert from IntArrayRef to vector),\\n # so the caller of translate() should be explicit that they need it.\\n if allow_expensive_conversions:\\n if goal.type == VectorCType(BaseCType(longT)):\\n intArrayRef_ctype = NamedCType(goal.name, BaseCType(intArrayRefT))\\n argname = direct_solve(intArrayRef_ctype)\\n return f\\\"{argname}.vec()\\\"\\n if goal.type == VectorCType(BaseCType(SymIntT)):\\n symIntArrayRef_ctype = NamedCType(goal.name, BaseCType(symIntArrayRefT))\\n argname = direct_solve(symIntArrayRef_ctype)\\n return f\\\"{argname}.vec()\\\"\\n elif goal.type == OptionalCType(VectorCType(BaseCType(longT))):\\n optionalIntArrayRef_ctype = NamedCType(\\n goal.name, BaseCType(optionalIntArrayRefT)\\n )\\n argname = direct_solve(optionalIntArrayRef_ctype)\\n return f\\\"{argname}.has_value() ? ::std::make_optional({argname}->vec()) : ::std::nullopt\\\"\\n elif goal.type == OptionalCType(BaseCType(scalarT)):\\n optionalScalarRef_ctype = NamedCType(\\n goal.name, BaseCType(optionalScalarRefT)\\n )\\n argname = direct_solve(optionalScalarRef_ctype)\\n return f\\\"{argname}.has_value() ? ::std::make_optional({argname}) : ::std::nullopt\\\"\\n elif goal.type == OptionalCType(BaseCType(scalarT)):\\n optionalTensorRef_ctype = NamedCType(\\n goal.name, BaseCType(optionalTensorRefT)\\n )\\n argname = direct_solve(optionalTensorRef_ctype)\\n return f\\\"{argname}.has_value() ? ::std::make_optional({argname}) : ::std::nullopt\\\"\\n # Technically, we also need to handle cases of C++ containers holding reference types.\\n # But there currently aren't any ops that require lambda capture codegen\\n # With arguments like ::std::vector.\\n # If that changes, we'll have to add the translation here.\\n\\n # We allow const casting on tensors, since const-correctness is a bit broken for at::Tensor.\\n # We could probably generalize this to non-tensor types too.\\n if goal.type == MutRefCType(BaseCType(tensorT)):\\n const_ref_tensor_ctype = NamedCType(\\n goal.name, ConstRefCType(BaseCType(tensorT))\\n )\\n argname = direct_solve(const_ref_tensor_ctype)\\n return f\\\"const_cast({argname})\\\"\\n\\n unsat(goal)\\n\\n return [Expr(solve(g, direct=False), g) for g in goal_ctypes]\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\nfrom typing import Sequence\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.types import Binding, CppSignature, CppSignatureGroup\\nfrom torchgen.gen import pythonify_default\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n ListType,\\n NativeFunction,\\n OptionalType,\\n Return,\\n Type,\\n Variant,\\n)\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Data Models\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# [Notes] python binding codegen\\n#\\n# The Python binding codegen produces code that takes the input list of\\n# PyObjects, finds the matching ATen C++ function using PythonArgParser,\\n# converts the PyObjects into C++ types and calls the ATen C++ function:\\n#\\n# +--------+ parsing +------------------------+ binding +-----------------------+\\n# | PyObjs | ---------> | PythonArgParser Output | ---------> | Cpp Function Dispatch |\\n# +--------+ +------------------------+ +-----------------------+\\n#\\n# The following examples demonstrate the data models the Python binding\\n# codegen needs to deal with and the tasks it needs to accomplish. It\\n# helps understand the purpose of the new data types we introduced below.\\n#\\n# - Function Schema (source of truth)\\n#\\n# aten::empty.names(int[] size, *, Dimname[]? names,\\n# ScalarType? dtype=None, Layout? layout=None,\\n# Device? device=None, bool? pin_memory=None,\\n# MemoryFormat? memory_format=None) -> Tensor\\n#\\n# - Python Signature\\n#\\n# It's used to generate input schema string for PythonArgParser.\\n# Note: TensorOptions fields are reordered and the additional\\n# 'requires_grad' field is added:\\n#\\n# empty(IntArrayRef size, *, DimnameList? names,\\n# MemoryFormat? memory_format=None, ScalarType dtype=None,\\n# Layout layout=torch.strided, Device device=None,\\n# bool pin_memory=False, bool requires_grad=False)\\n#\\n# - C++ Signature\\n#\\n# It's used to generate C++ lambda formals & dispatch call.\\n# Note: the scattered TensorOptions fields are packed into 'options'.\\n#\\n# auto dispatch_empty =\\n# [](IntArrayRef size, std::optional names,\\n# const TensorOptions & options,\\n# std::optional memory_format) -> Tensor {\\n# pybind11::gil_scoped_release no_gil;\\n# return torch::empty(size, names, options, memory_format);\\n# };\\n#\\n# - Binding between Python Arguments and C++ Arguments\\n#\\n# Given a set of Python Arguments in scope, we need produce the\\n# binding expressions that translate the Python API into C++ API:\\n#\\n# Python Args Cpp Args Binding Exprs\\n# -----------------------------------------------------------------\\n# 0: size size '_r.intlist(0)'\\n# 1: names names 'names' [special init]\\n# 2: memory_format -------+\\n# 3: dtype -----+-|--> options 'options' [special packing]\\n# 4: layout / |\\n# 5: device / +--> memory_format '_r.memoryformatOptional(2)'\\n# 6: pin_memory /\\n# 7: requires_grad -+\\n#\\n# So the full dispatch expression would look like:\\n#\\n# dispatch_empty(_r.intlist(0), names, options,\\n# _r.memoryformatOptional(2))\\n#\\n# Where does 'names' come from? It involves special local init:\\n#\\n# auto __names = _r.toDimnameListOptional(1);\\n# std::optional names =\\n# __names ? std::make_optional(DimnameList(__names.value()))\\n# : std::nullopt;\\n#\\n# Where does 'options' come from? It involves special local init\\n# for TensorOptions. Note that Python side has the additional\\n# 'requires_grad' field:\\n#\\n# const auto options = TensorOptions()\\n# .dtype(_r.scalartype(3))\\n# .device(_r.device(5))\\n# .layout(_r.layoutOptional(4))\\n# .requires_grad(_r.toBool(7))\\n# .pinned_memory(_r.toBool(6));\\n#\\n# In some other cases one Python Argument can map to multiple C++\\n# Arguments. For example:\\n#\\n# aten::max.names_dim(Tensor self, Dimname dim, bool keepdim=False)\\n# -> (Tensor values, Tensor indices)\\n#\\n# Python Args Cpp Args Binding Exprs\\n# ---------------------------------------------------------------------\\n# +----> max 'out[0]'\\n# /-----> max_values 'out[1]\\n# 0: input / self '_r.tensor(0)'\\n# 1: dim / dim '_r.dimname(1)'\\n# 2: keepdim / keepdim '_r.toBool(2)'\\n# 3: out -----+ [local init] out '_r.tensorlist_n<2>(3)'\\n#\\n# As demonstrated above, the binding can involve reordering,\\n# packing, unpacking and special local inits.\\n#\\n#\\n# Let's look at a concrete example:\\n#\\n# static PythonArgParser parser({\\n# \\\"abs(Tensor input, *, Tensor out=None)\\\",\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n# ^\\n# +--- Python Schema, represented by PythonSignature and PythonArgument\\n#\\n# }, /*traceable=*/true);\\n#\\n# ParsedArgs<2> parsed_args;\\n# auto _r = parser.parse(nullptr, args, kwargs, parsed_args);\\n#\\n# ...\\n#\\n# if (_r.isNone(1)) {\\n# ~~~~~~~~~~~~ <--- Scattered PythonArgParser output (arg name = 'out')\\n# represented by PythonArgParserOutputExpr\\n#\\n# // aten::abs(Tensor self) -> Tensor\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n# ^\\n# +--- NativeFunction schema, base version\\n#\\n# auto dispatch_abs = [](const Tensor & self) -> Tensor {\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n# ^\\n# +--- dispatch_lambda_args / dispatch_lambda_return_str\\n# generated from NativeFunction / CppSignature\\n# (deprecated PythonSignature is special)\\n# arguments are represented by DispatchLambdaArgument\\n#\\n# pybind11::gil_scoped_release no_gil;\\n# return self.abs();\\n# ~~~~~~~~~~~ <--- cpp_dispatch_target / cpp_dispatch_exprs\\n# generated from NativeFunction / CppSignature\\n# };\\n# return wrap(dispatch_abs(_r.tensor(0)));\\n# ~~~~~~~~~~~~~\\n# ^\\n# +--- dispatch_lambda_exprs\\n# binding PythonArgParserOutputExpr (python args)\\n# and DispatchLambdaArgument (c++ args)\\n#\\n# } else {\\n# // aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!)\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n# ^\\n# +--- NativeFunction schema, out-variant\\n#\\n# auto dispatch_abs_out = [](Tensor out, const Tensor & self) -> Tensor {\\n# pybind11::gil_scoped_release no_gil;\\n# return at::abs_out(out, self);\\n# };\\n# return wrap(dispatch_abs_out(_r.tensor(1), _r.tensor(0)));\\n# }\\n#\\n#\\n# [Notes] python interface codegen\\n# The python dataclasses below are used used to generate both python binding code\\n# and pyi type hint signatures.\\n# In theory these two should look very similar, but there are number of differences\\n# in how pyi signatures vs. python_arg_parser signatures are generated.\\n# These differences have been encapsulated in signature_str() vs. signature_str_pyi()\\n# to display the full signatures, and argument_str() vs argument_str_pyi() to display arguments.\\n# For examples, only pyi signatures include return types.\\n\\n\\n@dataclass(frozen=True)\\nclass PythonReturns:\\n returns: tuple[Return, ...]\\n\\n\\n@dataclass(frozen=True)\\nclass PythonArgument:\\n name: str\\n type: Type\\n default: str | None\\n\\n # Used to generate the default init expr for some PythonArgParser outputs, e.g.:\\n #\\n # _r.layoutWithDefault(3, layout_from_backend(self.options().backend())))\\n # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n # ^\\n # +--- default_init str\\n default_init: str | None\\n\\n # Compute argument formal for python argument parsing.\\n # Needs to be consistent with torch/csrc/utils/python_arg_parser.h.\\n def argument_str(self, *, method: bool = False, symint: bool = True) -> str:\\n type_str = (\\n argument_type_str(self.type, symint=symint)\\n .replace(\\\"const \\\", \\\"\\\")\\n .replace(\\\" &\\\", \\\"\\\")\\n )\\n\\n name = self.name\\n # s/self/input/ outside method bindings\\n # [old codegen] TODO: remove this? doesn't rename in codegen, it's just\\n # for the parse string\\n if name == \\\"self\\\" and type_str in [\\\"Tensor\\\", \\\"Number\\\"] and not method:\\n name = \\\"input\\\"\\n\\n # add default\\n if self.default is not None:\\n default = {\\n \\\"nullptr\\\": \\\"None\\\",\\n \\\"::std::nullopt\\\": \\\"None\\\",\\n \\\"std::nullopt\\\": \\\"None\\\",\\n \\\"{}\\\": \\\"None\\\",\\n }.get(self.default, self.default)\\n return f\\\"{type_str} {name}={default}\\\"\\n else:\\n return f\\\"{type_str} {name}\\\"\\n\\n def argument_str_pyi(\\n self, *, method: bool = False, deprecated: bool = False\\n ) -> str:\\n type_str = argument_type_str_pyi(self.type)\\n\\n name = self.name\\n # s/self/input/ outside method bindings\\n # [old codegen] TODO: remove this? doesn't rename in codegen, it's just\\n # for the parse string\\n if name == \\\"self\\\" and type_str == \\\"Tensor\\\" and not method and not deprecated:\\n name = \\\"input\\\"\\n\\n if name == \\\"from\\\": # from is a Python keyword...\\n name += \\\"_\\\"\\n\\n # pyi merges the _out and functional variants into the same signature, with an optional out arg\\n if name == \\\"out\\\" and type_str == \\\"Tensor\\\" and not deprecated:\\n type_str = \\\"Optional[\\\" + type_str + \\\"]\\\"\\n\\n # pyi deprecated signatures don't get defaults for their out arg\\n treat_as_no_default = (\\n deprecated\\n and isinstance(self, PythonOutArgument)\\n and self.default == \\\"None\\\"\\n )\\n\\n # add default\\n if self.default is not None and not treat_as_no_default:\\n if (\\n isinstance(self.type, ListType)\\n and self.type.elem == BaseType(BaseTy.int)\\n and self.default.startswith(\\\"{\\\")\\n and self.default.endswith(\\\"}\\\")\\n ):\\n default = (\\n \\\"(\\\" + \\\", \\\".join(map(str.strip, self.default[1:-1].split(\\\",\\\"))) + \\\")\\\"\\n )\\n else:\\n default = {\\n \\\"nullptr\\\": \\\"None\\\",\\n \\\"::std::nullopt\\\": \\\"None\\\",\\n \\\"std::nullopt\\\": \\\"None\\\",\\n \\\"{}\\\": \\\"None\\\",\\n \\\"c10::MemoryFormat::Contiguous\\\": \\\"contiguous_format\\\",\\n \\\"QScheme::PER_TENSOR_AFFINE\\\": \\\"per_tensor_affine\\\",\\n }.get(self.default, self.default)\\n return f\\\"{name}: {type_str} = {default}\\\"\\n else:\\n return f\\\"{name}: {type_str}\\\"\\n\\n\\n@dataclass(frozen=True)\\nclass PythonOutArgument(PythonArgument):\\n # In Python signature multiple output fields are packed into one 'out' argument.\\n # When binding to C++, it's first binded to a local 'out' variable:\\n # 'auto out = _r.tensorlist_n<2>(2);',\\n # then binded to scattered C++ output arguments as 'out[0]', 'out[1]', and etc.\\n # TODO: maybe don't need keep scattered out fields for python signature?\\n outputs: tuple[PythonArgument, ...]\\n\\n @staticmethod\\n def from_outputs(outputs: tuple[PythonArgument, ...]) -> PythonOutArgument | None:\\n if not outputs:\\n return None\\n\\n size = len(outputs)\\n if size == 1:\\n return PythonOutArgument(\\n name=outputs[0].name,\\n type=outputs[0].type,\\n default=\\\"None\\\",\\n default_init=None,\\n outputs=outputs,\\n )\\n elif size > 1:\\n if any(not a.type.is_tensor_like() for a in outputs):\\n raise RuntimeError(f\\\"Unsupported output type: {outputs}\\\")\\n return PythonOutArgument(\\n name=\\\"out\\\",\\n # TODO: shouldn't this be OptionalType[ListType[...]], since it defaults to None?\\n type=ListType(BaseType(BaseTy.Tensor), size),\\n default=\\\"None\\\",\\n default_init=None,\\n outputs=outputs,\\n )\\n raise AssertionError(r\\\"Unexpected PythonOutArgument size\\\")\\n\\n\\n@dataclass(frozen=True)\\nclass PythonSignature:\\n # Base operator name, without inplace/outplace suffix.\\n name: str\\n\\n # Positional arguments.\\n # TODO: create a dedicated SelfArgument type for 'self'?\\n input_args: tuple[PythonArgument, ...]\\n\\n # Keyword arguments excluding the 'out' argument and scattered kwargs belonging\\n # to TensorOptions (dtype, layout, device, pin_memory, requires_grad, etc).\\n input_kwargs: tuple[PythonArgument, ...]\\n\\n output_args: PythonOutArgument | None\\n\\n # Return types, which are only used by pyi\\n returns: PythonReturns\\n\\n # These are scattered kwargs arguments belonging to TensorOptions.\\n # When binding to C++, they are packed into a TensorOptions object 'options'.\\n # It's possible that the C++ signature doesn't take TensorOptions object (e.g.\\n # for out variant), in which case they will be used as scattered fields without\\n # being packed into 'options'.\\n # TODO: maybe create a PythonTensorOptionsArgument?\\n tensor_options_args: tuple[PythonArgument, ...]\\n\\n # method or function signature?\\n method: bool\\n\\n @property\\n def deprecated(self) -> bool:\\n return False\\n\\n def arguments(\\n self, *, skip_outputs: bool = False, skip_tensor_options: bool = False\\n ) -> tuple[PythonArgument | PythonOutArgument, ...]:\\n result: list[PythonArgument | PythonOutArgument] = []\\n result.extend(self.input_args)\\n result.extend(self.input_kwargs)\\n if self.output_args is not None and not skip_outputs:\\n result.append(self.output_args)\\n if not skip_tensor_options:\\n result.extend(self.tensor_options_args)\\n return tuple(result)\\n\\n def arguments_count(self) -> int:\\n return len(self.arguments())\\n\\n def output_idx(self) -> int:\\n return len(self.input_args) + len(self.input_kwargs)\\n\\n # [old codegen] Compute the Python function signature for argument parsing,\\n # as specified in torch/csrc/utils/python_arg_parser.h. WARNING:\\n # this is NOT the same type signature as specified by PEP 484\\n # as understood by mypy; our format was independently developed\\n # and has some quirks to make it more suitable specifically\\n # for error parsing.\\n #\\n # For a translation to mypy-valid type signatures, see\\n # signature_str_pyi().\\n def signature_str(self, *, skip_outputs: bool = False, symint: bool = True) -> str:\\n args = self.arguments(skip_outputs=skip_outputs)\\n schema_formals: list[str] = [\\n a.argument_str(method=self.method, symint=symint) for a in args\\n ]\\n positional_argc = len(self.input_args)\\n if len(schema_formals) > positional_argc:\\n schema_formals.insert(positional_argc, \\\"*\\\")\\n\\n return f'{self.name}({\\\", \\\".join(schema_formals)})'\\n\\n def signature_str_pyi(self, *, skip_outputs: bool = False) -> str:\\n args = self.arguments(skip_outputs=skip_outputs)\\n schema_formals: list[str] = [\\n a.argument_str_pyi(method=self.method) for a in args\\n ]\\n positional_argc = len(self.input_args)\\n if len(schema_formals) > positional_argc:\\n schema_formals.insert(positional_argc, \\\"*\\\")\\n\\n # only pyi signatures include returns\\n returns_str = returns_str_pyi(self)\\n # pyi also includes self (with no typing/defaults) for methods\\n if self.method:\\n schema_formals.insert(0, \\\"self\\\")\\n return f'def {self.name}({\\\", \\\".join(schema_formals)}) -> {returns_str}: ...'\\n\\n def signature_str_pyi_vararg(self, *, skip_outputs: bool = False) -> str | None:\\n # only pyi uses vararg signatures\\n args = self.arguments(skip_outputs=skip_outputs)\\n schema_formals: list[str] = [\\n a.argument_str_pyi(method=self.method) for a in args\\n ]\\n # vararg only applies to pyi signatures. vararg variants are not generated for all signatures\\n num_args = self.arguments_count()\\n num_positionalargs = len(self.input_args)\\n\\n have_vararg_version = False\\n if num_args > 0:\\n vararg_type = args[0].type\\n if (\\n isinstance(vararg_type, ListType)\\n and str(vararg_type.elem) in [\\\"int\\\", \\\"SymInt\\\"]\\n and num_positionalargs == 1\\n ):\\n have_vararg_version = True\\n\\n if not have_vararg_version:\\n return None\\n\\n # Below are the major changes in vararg vs. regular pyi signatures\\n # vararg signatures also omit the asterix\\n assert isinstance(vararg_type, ListType)\\n schema_formals[0] = (\\n \\\"*\\\" + args[0].name + \\\": \\\" + argument_type_str_pyi(vararg_type.elem)\\n )\\n\\n returns_str = returns_str_pyi(self)\\n # pyi also includes self (with no typing/defaults) for methods\\n if self.method:\\n schema_formals.insert(0, \\\"self\\\")\\n return f'def {self.name}({\\\", \\\".join(schema_formals)}) -> {returns_str}: ...'\\n\\n\\n# The deprecated python signature involves some special logic, so create a\\n# dedicated data model to store these extra properties.\\n@dataclass(frozen=True)\\nclass PythonSignatureDeprecated(PythonSignature):\\n # Schema for the deprecated function\\n deprecated_schema: FunctionSchema\\n\\n # The deprecated signature might miss some arguments that the corresponding\\n # C++ signature expects. We need store the constant default values to pass in.\\n # For example:\\n # [deprecate signature]: addmm(Scalar beta, Tensor self, Tensor mat1, Tensor mat2)\\n # [func schema]: aten::addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor\\n # [func call]: self.addmm(mat1, mat2, beta, 1)\\n # We store ['self', 'mat1', 'mat2', 'beta', '1'] in this case.\\n deprecated_args_exprs: tuple[str, ...]\\n\\n @property\\n def deprecated(self) -> bool:\\n return True\\n\\n def signature_str(self, *, skip_outputs: bool = False, symint: bool = True) -> str:\\n return (\\n PythonSignature.signature_str(\\n self, skip_outputs=skip_outputs, symint=symint\\n )\\n + \\\"|deprecated\\\"\\n )\\n\\n def signature_str_pyi(self, *, skip_outputs: bool = False) -> str:\\n args = self.arguments(skip_outputs=skip_outputs)\\n schema_formals: list[str] = [\\n a.argument_str_pyi(method=self.method, deprecated=True) for a in args\\n ]\\n positional_argc = len(self.input_args)\\n if len(schema_formals) > positional_argc:\\n schema_formals.insert(positional_argc, \\\"*\\\")\\n\\n returns_str = returns_str_pyi(self)\\n return f'def {self.name}({\\\", \\\".join(schema_formals)}) -> {returns_str}: ...'\\n\\n def signature_str_pyi_vararg(self, *, skip_outputs: bool = False) -> str | None:\\n # the codegen doesn't include vararg variants for deprecated signatures\\n return None\\n\\n\\n# This struct is used to hold the PythonSignature and its corresponding\\n# NativeFunction BEFORE grouping base and out-variant functions.\\n# Why not store NativeFunction in PythonSignature or construct PythonSignature\\n# from NativeFunction? Because they are not 1-1 mapped.\\n# One native function could have both deprecated and non-deprecated python\\n# signatures - NativeFunction doesn't contain information to construct the\\n# deprecated python signature.\\n# One python signature is used to handle both the base and the out-variant\\n# function - see 'PythonSignatureGroup'.\\n@dataclass(frozen=True)\\nclass PythonSignatureNativeFunctionPair:\\n signature: PythonSignature\\n function: NativeFunction\\n\\n\\n# We merge pairs of functions with signatures that are equivalent mod\\n# output arguments, and use a single entry in the python_arg_parser sig\\n# list for both (output arguments become optional).\\n@dataclass(frozen=True)\\nclass PythonSignatureGroup:\\n # The signature used for Python argument parsing. The outplace signature\\n # is preferred if exists, because it can be used to parse inputs for both\\n # the out-place variant and the base version (with output omitted).\\n signature: PythonSignature\\n\\n # The regular ATen declaration (e.g. conv2d)\\n base: NativeFunction\\n\\n # The out variant (e.g. conv2d_out)\\n outplace: NativeFunction | None\\n\\n @classmethod\\n def from_pairs(\\n cls,\\n functional: PythonSignatureNativeFunctionPair,\\n out: PythonSignatureNativeFunctionPair | None,\\n ) -> PythonSignatureGroup:\\n if out is None:\\n return PythonSignatureGroup(\\n signature=functional.signature,\\n base=functional.function,\\n outplace=None,\\n )\\n\\n # prefer the signature with optional out=... arguments because it's the\\n # superset that can be used to parse input for both base and outplace.\\n signature_kwargs = out.signature.__dict__.copy()\\n\\n # Out overloads in C++ don't have TensorOptions arguments,\\n # so take these from the functional variant\\n signature_kwargs[\\n \\\"tensor_options_args\\\"\\n ] = functional.signature.tensor_options_args\\n\\n return PythonSignatureGroup(\\n signature=type(out.signature)(**signature_kwargs),\\n base=functional.function,\\n outplace=out.function,\\n )\\n\\n\\n# C++ function dispatch is wrapped in a lambda function. The lambda function\\n# has almost the same signature as the C++ function, only with some small\\n# variants - see details below.\\n# This data model is used to represent arguments of the lambda function\\n# signature.\\n@dataclass(frozen=True)\\nclass DispatchLambdaArgument:\\n name: str\\n type_str: str\\n is_out_arg: bool\\n\\n\\n# To pass PyObjects arguments to C++ function (via the lambda wrapper),\\n# we need first convert PyObjects into simple C++ objects. This work\\n# is done by PythonArgParser.\\n# This data model is used to represent the output of PythonArgParser.\\n# It has 1-1 mapping with PythonArgument in PythonSignature.\\n@dataclass(frozen=True)\\nclass PythonArgParserOutputExpr:\\n # argument name\\n name: str\\n\\n # RHS expression to reference PythonArgParser output.\\n expr: str\\n\\n # In some special cases we need create different expr, e.g.:\\n # '_r.isNone(1)' instead of '_r.tensor(1)'.\\n index: int\\n\\n # The python argument it maps to.\\n argument: PythonArgument\\n\\n @property\\n def is_none_expr(self) -> str:\\n return f\\\"_r.isNone({self.index})\\\"\\n\\n\\n# To pass PythonArgParser output to the lambda wrapper, we need bind\\n# PythonArgParserOutputExpr to DispatchLambdaArgument.\\n# They are not always 1-1 mapped, e.g. scattered TensorOptions fields\\n# need be packed into a TensorOptions object, which is the argument\\n# that the lambda function wrapper takes.\\n@dataclass(frozen=True)\\nclass DispatchLambdaArgumentExprs:\\n # The exprs that provide the binding for lambda arguments, e.g.:\\n #\\n # 'self' -> '_r.tensor(0)'\\n # 'min' -> 'out[0]' / 'min_indices' -> 'out[1]'\\n # 'options' -> 'options'\\n #\\n # It has 1-1 mapping with DispatchLambdaArgument.\\n exprs: Sequence[str]\\n\\n # Special local inits, which might introduce new variables that\\n # the 'exprs' above reference, e.g.:\\n #\\n # 'auto out = _r.tensorlist_n<2>(2);'\\n #\\n inits: Sequence[str]\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Helper Functions\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef _cpp_signature(f: NativeFunction, *, method: bool = False) -> CppSignature:\\n return CppSignatureGroup.from_native_function(f, method=method).signature\\n\\n\\ndef has_tensor_options(f: NativeFunction) -> bool:\\n return f.func.arguments.tensor_options is not None\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Python Signature\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\n# 'simple_type' was introduced by the old codegen, which is slightly\\n# different from the python schema type, e.g.: doesn't have '?' suffix\\n# for optional Tensor/TensorList; doesn't have '[size]' suffix for list type.\\ndef argument_type_str(\\n t: Type, *, simple_type: bool = False, symint: bool = True\\n) -> str:\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Tensor:\\n return \\\"Tensor\\\"\\n elif t.name == BaseTy.int:\\n return \\\"int64_t\\\"\\n elif t.name == BaseTy.float:\\n return \\\"double\\\"\\n elif t.name == BaseTy.str:\\n return \\\"c10::string_view\\\"\\n elif t.name in [\\n BaseTy.bool,\\n BaseTy.QScheme,\\n BaseTy.Scalar,\\n BaseTy.ScalarType,\\n BaseTy.Generator,\\n BaseTy.Storage,\\n BaseTy.Layout,\\n BaseTy.Device,\\n BaseTy.DeviceIndex,\\n BaseTy.MemoryFormat,\\n BaseTy.Dimname,\\n BaseTy.Stream,\\n BaseTy.ConstQuantizerPtr,\\n BaseTy.SymInt,\\n ]:\\n # These python schema type names line up with their function schema names\\n return t.name.name\\n\\n elif isinstance(t, OptionalType):\\n if str(t.elem) == \\\"Tensor\\\":\\n # Is it desired to keep '?' for simple_type with new style dispatcher?\\n return \\\"Tensor?\\\"\\n elem = argument_type_str(t.elem, simple_type=simple_type, symint=symint)\\n return f\\\"{elem}?\\\"\\n elif isinstance(t, ListType):\\n size = t.size if not simple_type else None\\n if str(t.elem) == \\\"bool\\\":\\n assert t.size is not None\\n return f\\\"::std::array\\\"\\n elif str(t.elem) == \\\"int\\\":\\n return f\\\"IntArrayRef[{size}]\\\" if size is not None else \\\"IntArrayRef\\\"\\n elif str(t.elem) == \\\"SymInt\\\":\\n if symint:\\n return (\\n f\\\"SymIntArrayRef[{size}]\\\" if size is not None else \\\"SymIntArrayRef\\\"\\n )\\n else:\\n return f\\\"IntArrayRef[{size}]\\\" if size is not None else \\\"IntArrayRef\\\"\\n elif str(t.elem) == \\\"Tensor\\\":\\n return f\\\"TensorList[{size}]\\\" if size is not None else \\\"TensorList\\\"\\n elif str(t.elem) == \\\"Scalar\\\":\\n return f\\\"ScalarList[{size}]\\\" if size is not None else \\\"ScalarList\\\"\\n elif str(t.elem) == \\\"Tensor?\\\":\\n if simple_type:\\n return \\\"c10::List<::std::optional>\\\"\\n else:\\n return \\\"const c10::List<::std::optional> &\\\"\\n elif str(t.elem) == \\\"Dimname\\\":\\n return f\\\"DimnameList[{size}]\\\" if size is not None else \\\"DimnameList\\\"\\n elem = argument_type_str(t.elem, simple_type=simple_type, symint=symint)\\n return f\\\"ArrayRef<{elem}>\\\"\\n\\n raise RuntimeError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\ndef argument_type_size(t: Type) -> int | None:\\n l = t.is_list_like()\\n if l is not None and str(l.elem) != \\\"bool\\\":\\n return l.size\\n else:\\n return None\\n\\n\\ndef argument(a: Argument) -> PythonArgument:\\n return PythonArgument(\\n name=a.name,\\n type=a.type,\\n # TODO: directly translate a.default to python default\\n default=(\\n str(pythonify_default(cpp.default_expr(a.default, a.type, symint=False)))\\n if a.default is not None\\n else None\\n ),\\n default_init=None,\\n )\\n\\n\\n# Generates a PythonSignature that can be used for either .pyi or PythonArgParser codegen\\ndef signature(\\n f: NativeFunction, *, method: bool = False, pyi: bool = False\\n) -> PythonSignature:\\n return signature_from_schema(\\n f.func, category_override=f.category_override, method=method, pyi=pyi\\n )\\n\\n\\ndef signature_from_schema(\\n func: FunctionSchema,\\n *,\\n category_override: str | None,\\n method: bool = False,\\n pyi: bool = False,\\n) -> PythonSignature:\\n args: list[Argument] = []\\n args.extend(func.arguments.pre_self_positional)\\n # Skip SelfArgument if this is method.\\n if not method and func.arguments.self_arg is not None:\\n args.append(func.arguments.self_arg.argument)\\n args.extend(func.arguments.post_self_positional)\\n args.extend(func.arguments.pre_tensor_options_kwarg_only)\\n # Skip TensorOptionsArguments. Python side TensorOptions\\n # arguments are created based on different rules - see below.\\n args.extend(func.arguments.post_tensor_options_kwarg_only)\\n args.extend(func.arguments.out)\\n\\n input_arg_set = {a.name for a in func.arguments.flat_positional}\\n kwarg_only_set = {a.name for a in func.arguments.flat_kwarg_only}\\n out_arg_set = {a.name for a in func.arguments.out}\\n\\n input_args = tuple(map(argument, filter(lambda a: a.name in input_arg_set, args)))\\n input_kwargs = tuple(\\n map(argument, filter(lambda a: a.name in kwarg_only_set, args))\\n )\\n outputs = tuple(map(argument, filter(lambda a: a.name in out_arg_set, args)))\\n\\n # Reintroduce the scattered fields of TensorOptions for Python.\\n # Compared to the cpp counterpart, the python arguments have new property\\n # (default_init) and a new argument 'requires_grad', which require some\\n # special handlings.\\n # [old codegen] TODO: because these aren't guaranteed to be 100% faithful\\n # to the original versions in the yaml, this recreation is a potential\\n # source of drift between eager and JIT. Pull this logic out to a shared place.\\n\\n has_tensor_input_arg = any(\\n a.type.is_tensor_like() for a in func.arguments.flat_non_out\\n )\\n if any(a.name == \\\"requires_grad\\\" for a in func.schema_order_arguments()):\\n raise ValueError(\\n \\\"argument named requires_grad is reserved, should not explicitly add it in the schema\\\"\\n )\\n\\n # [old codegen] this probably won't work if one of the returns is not a tensor,\\n # but it will produce a compile-time error that is obvious.\\n has_tensor_return = any(r.type.is_tensor_like() for r in func.returns)\\n\\n name: str = cpp.name(func)\\n is_factory_function = category_override == \\\"factory\\\" or (\\n has_tensor_return and not has_tensor_input_arg\\n )\\n is_like_or_new_function = (\\n category_override in (\\\"new\\\", \\\"like\\\")\\n or name.startswith(\\\"new_\\\")\\n or name.endswith(\\\"_like\\\")\\n )\\n is_dummy_function = category_override == \\\"dummy\\\"\\n\\n tensor_options_args: list[PythonArgument] = []\\n if (is_factory_function or is_like_or_new_function) and not is_dummy_function:\\n\\n def topt_default_init(name: str) -> str | None:\\n topt_args = func.arguments.tensor_options\\n if topt_args is None:\\n return None\\n a = getattr(topt_args, name)\\n if a.default is None or a.default == \\\"None\\\":\\n return None\\n return cpp.default_expr(a.default, a.type, symint=False)\\n\\n tensor_options_args.append(\\n PythonArgument(\\n name=\\\"dtype\\\",\\n type=OptionalType(BaseType(BaseTy.ScalarType)),\\n default=\\\"None\\\",\\n default_init=(\\n None if is_like_or_new_function else topt_default_init(\\\"dtype\\\")\\n ),\\n )\\n )\\n tensor_options_args.append(\\n PythonArgument(\\n name=\\\"layout\\\",\\n type=OptionalType(BaseType(BaseTy.Layout)),\\n default=\\\"None\\\",\\n default_init=(\\n None if is_like_or_new_function else topt_default_init(\\\"layout\\\")\\n ),\\n )\\n )\\n tensor_options_args.append(\\n PythonArgument(\\n name=\\\"device\\\",\\n type=OptionalType(BaseType(BaseTy.Device)),\\n default=\\\"None\\\",\\n default_init=(\\n None\\n if is_like_or_new_function\\n else (\\n topt_default_init(\\\"device\\\")\\n or \\\"torch::tensors::get_default_device()\\\"\\n )\\n ),\\n )\\n )\\n tensor_options_args.append(\\n PythonArgument(\\n name=\\\"pin_memory\\\",\\n type=OptionalType(BaseType(BaseTy.bool)),\\n default=\\\"False\\\",\\n default_init=None,\\n )\\n )\\n tensor_options_args.append(\\n PythonArgument(\\n name=\\\"requires_grad\\\",\\n type=OptionalType(BaseType(BaseTy.bool)),\\n default=\\\"False\\\",\\n default_init=None,\\n )\\n )\\n\\n returns = PythonReturns(returns=func.returns)\\n\\n return PythonSignature(\\n name=str(func.name.name),\\n input_args=input_args,\\n input_kwargs=input_kwargs,\\n output_args=PythonOutArgument.from_outputs(outputs),\\n tensor_options_args=tuple(tensor_options_args),\\n returns=returns,\\n method=method,\\n )\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Python Interface\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef structseq_fieldnames(returns: tuple[Return, ...]) -> list[str]:\\n if len(returns) <= 1 or all(r.name is None for r in returns):\\n return []\\n else:\\n if any(r.name is None for r in returns):\\n # When building on Windows, `PyStructSequence_UnnamedField` could not be\\n # resolved by the linker for some reason, which cause error in building:\\n #\\n # python_nn_functions.cpp.obj : error LNK2001: unresolved external symbol\\n # PyStructSequence_UnnamedField\\n #\\n # Thus, at this point in time, we do not support unnamed\\n # fields in structseq; you must either name all fields,\\n # or none of them.\\n raise ValueError(\\\"Unnamed field is not supported by codegen\\\")\\n\\n return [str(r.name) for r in returns]\\n\\n\\ndef argument_type_str_pyi(t: Type) -> str:\\n add_optional = False\\n if isinstance(t, OptionalType):\\n t = t.elem\\n add_optional = True\\n\\n if isinstance(t, BaseType):\\n if t.name in [BaseTy.int, BaseTy.DeviceIndex]:\\n ret = \\\"_int\\\"\\n if t.name == BaseTy.SymInt:\\n ret = \\\"Union[_int, SymInt]\\\"\\n elif t.name == BaseTy.float:\\n ret = \\\"_float\\\"\\n elif t.name == BaseTy.str:\\n ret = \\\"str\\\"\\n elif t.name == BaseTy.Scalar:\\n ret = \\\"Union[Number, _complex]\\\"\\n elif t.name == BaseTy.ScalarType:\\n ret = \\\"_dtype\\\"\\n elif t.name == BaseTy.bool:\\n ret = \\\"_bool\\\"\\n elif t.name == BaseTy.QScheme:\\n ret = \\\"_qscheme\\\"\\n elif t.name == BaseTy.Layout:\\n ret = \\\"_layout\\\"\\n elif t.name == BaseTy.Device:\\n ret = \\\"Optional[DeviceLikeType]\\\"\\n elif t.name == BaseTy.MemoryFormat:\\n ret = \\\"memory_format\\\"\\n elif t.name == BaseTy.Dimname:\\n ret = \\\"Union[str, ellipsis, None]\\\"\\n elif t.name == BaseTy.Storage:\\n ret = \\\"Union[Storage, UntypedStorage]\\\"\\n elif t.name in [BaseTy.Tensor, BaseTy.Generator, BaseTy.Stream]:\\n # These python schema type names line up with their function schema names\\n ret = t.name.name\\n\\n elif isinstance(t, ListType):\\n if str(t.elem) == \\\"int\\\":\\n ret = \\\"Union[_int, _size]\\\" if t.size is not None else \\\"_size\\\"\\n elif t.is_tensor_like():\\n # TODO: this doesn't seem right...\\n # Tensor?[] currently translates to Optional[Union[Tuple[Tensor, ...], List[Tensor]]]\\n # It should probably translate to Union[Tuple[Optional[Tensor], ...], List[Optional[Tensor]]]\\n if isinstance(t.elem, OptionalType):\\n add_optional = True\\n ret = (\\n \\\"Union[Tensor, Tuple[Tensor, ...], List[Tensor]]\\\"\\n if t.size is not None\\n else \\\"Union[Tuple[Tensor, ...], List[Tensor]]\\\"\\n )\\n elif str(t.elem) == \\\"float\\\":\\n ret = \\\"Sequence[_float]\\\"\\n elif str(t.elem) == \\\"SymInt\\\" and t.size is not None:\\n elem = argument_type_str_pyi(t.elem)\\n ret = f\\\"Union[{elem}, Sequence[{elem}]]\\\"\\n else:\\n elem = argument_type_str_pyi(t.elem)\\n ret = f\\\"Sequence[{elem}]\\\"\\n\\n else:\\n raise RuntimeError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n if add_optional:\\n ret = \\\"Optional[\\\" + ret + \\\"]\\\"\\n\\n return ret\\n\\n\\ndef return_type_str_pyi(t: Type) -> str:\\n # Where arguments are open to accepting Union, return types should return\\n # concrete types\\n\\n if isinstance(t, OptionalType):\\n inner = return_type_str_pyi(t.elem)\\n return f\\\"Optional[{inner}]\\\"\\n\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Device:\\n return \\\"_device\\\"\\n elif t.name == BaseTy.Dimname:\\n ret = \\\"Optional[str]\\\"\\n else:\\n return argument_type_str_pyi(t)\\n\\n if isinstance(t, ListType):\\n inner = return_type_str_pyi(t.elem)\\n return f\\\"Tuple[{inner}, ...]\\\"\\n\\n return argument_type_str_pyi(t)\\n\\n\\ndef returns_structseq_pyi(signature: PythonSignature) -> tuple[str, str] | None:\\n python_returns = [return_type_str_pyi(r.type) for r in signature.returns.returns]\\n structseq_name = signature.name\\n field_names = structseq_fieldnames(signature.returns.returns)\\n if field_names:\\n # These types are structseq objects which act like named NamedTuples, but\\n # the constructor acts like the constructor of tuple. Using typing.NamedTuple\\n # does not allow us to override __init__.\\n seq_type = f\\\"Tuple[{', '.join(python_returns)}]\\\"\\n structseq_def_lines = [\\n f\\\"class {structseq_name}({seq_type}):\\\",\\n ]\\n for name, typ in zip(field_names, python_returns):\\n structseq_def_lines.extend(\\n [\\n \\\" @property\\\",\\n f\\\" def {name}(self) -> {typ}: ...\\\",\\n ]\\n )\\n structseq_def_lines.extend(\\n [\\n f\\\" def __new__(cls, sequence: {seq_type}): ...\\\",\\n f\\\" n_fields: _int = {len(field_names)}\\\",\\n f\\\" n_sequeunce_fields: _int = {len(field_names)}\\\",\\n \\\" n_unnamed_fields: _int = 0\\\",\\n \\\" def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing\\\",\\n \\\"\\\", # add an extra newline\\n ]\\n )\\n structseq_def = \\\"\\\\n\\\".join(structseq_def_lines)\\n # Example:\\n # structseq_def = (\\n # \\\"class max(Tuple[Tensor, Tensor]):\\\\n\\\"\\n # \\\" @property\\\\n\\\"\\n # \\\" def values(self) -> Tensor: ...\\\\n\\\"\\n # \\\" @property\\\\n\\\"\\n # \\\" def indices(self) -> Tensor: ...\\\\n\\\"\\n # \\\" def __new__(cls, sequence: Tuple[Tensor, Tensor]): ...\\\\n\\\"\\n # \\\" n_fields: _int = 2\\\",\\n # \\\" n_sequeunce_fields: _int = 2\\\",\\n # \\\" n_unnamed_fields: _int = 0\\\",\\n # \\\" def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing\\\",\\n # )\\n return structseq_name, structseq_def\\n return None\\n\\n\\ndef returns_str_pyi(signature: PythonSignature) -> str:\\n field_names = structseq_fieldnames(signature.returns.returns)\\n if field_names:\\n return f\\\"torch.return_types.{signature.name}\\\"\\n\\n python_returns = [return_type_str_pyi(r.type) for r in signature.returns.returns]\\n if len(python_returns) > 1:\\n return \\\"Tuple[\\\" + \\\", \\\".join(python_returns) + \\\"]\\\"\\n if len(python_returns) == 1:\\n return python_returns[0]\\n return \\\"None\\\"\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# C++ Function Dispatch\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n# This section provides APIs to generate the code that does C++ function\\n# dispatch. The C++ function call is wrapped by a lambda function.\\n# For example:\\n#\\n# // aten::selu_(Tensor(a!) self) -> Tensor(a!)\\n# auto dispatch_selu_ = [](Tensor self) -> Tensor {\\n# pybind11::gil_scoped_release no_gil;\\n# return at::selu_(self);\\n# };\\n#\\n# The lambda function's signature follows the C++ signature in common\\n# cases, e.g.:\\n#\\n# // aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor\\n# [](const Tensor & self, const Tensor & other, Scalar alpha) -> Tensor\\n#\\n# For out variant the 'out' argument's type is changed from 'Tensor &'\\n# to 'Tensor'. It's because when calling the lambda it passes in the\\n# PythonArgParser output '_r.tensor(3)', which is stack allocated object\\n# and needs to pass by value. Also see comments in 'dispatch_lambda_return_str()'.\\n#\\n# // aten::add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!)\\n# [](Tensor out, const Tensor & self, const Tensor & other, Scalar alpha) -> Tensor\\n#\\n# For multi-output case it can keep using reference type because the\\n# PythonArgParser output has been unpacked to local variables, e.g.:\\n#\\n# // aten::max.names_dim_max(Tensor self, Dimname dim, bool keepdim=False, *,\\n# // Tensor(a!) max, Tensor(b!) max_values) -> (Tensor(a!) values, Tensor(b!) indices)\\n# [](Tensor & max, Tensor & max_values, const Tensor & self, Dimname dim, bool keepdim) -> std::tuple\\n#\\n# For deprecated python signature, it should follow deprecated python arg order.\\n# TODO: This is to keep same byte-for-byte result as the old codegen - maybe unnecessary?\\n\\n\\ndef dispatch_lambda_args(\\n ps: PythonSignature, f: NativeFunction, symint: bool = True\\n) -> tuple[DispatchLambdaArgument, ...]:\\n if isinstance(ps, PythonSignatureDeprecated):\\n schema = ps.deprecated_schema\\n else:\\n schema = f.func\\n\\n # Start with cpp arguments - dispatch lambda signature always include 'self'\\n cpp_args = cpp.arguments(\\n arguments=schema.arguments,\\n faithful=False,\\n symint=symint,\\n method=False,\\n cpp_no_default_args=f.cpp_no_default_args,\\n )\\n out_args: set[str] = {a.name for a in schema.arguments.out}\\n\\n # Convert from cpp argument to lambda argument\\n def dispatch_lambda_arg(cpp_arg: Binding) -> DispatchLambdaArgument:\\n type_str = cpp_arg.type\\n is_out_arg = cpp_arg.name in out_args\\n if ps.method and cpp_arg.name == \\\"self\\\":\\n # For method's 'self', we can use 'const Tensor &' and simply ignore mutability!\\n type_str = \\\"const at::Tensor &\\\"\\n else:\\n # For other cases we need prevent dangling refs to temps (unless it's\\n # unpacked scattered output)\\n # The reason is explained in the comments above and in 'dispatch_lambda_return_str()'.\\n # TODO: avoid this special handling?\\n ensure_temp_safe = len(out_args) <= 1 or not is_out_arg\\n if ensure_temp_safe:\\n type_str = {\\n \\\"at::Tensor &\\\": \\\"at::Tensor\\\",\\n }.get(type_str, type_str)\\n return DispatchLambdaArgument(\\n name=cpp_arg.name,\\n type_str=type_str,\\n is_out_arg=is_out_arg,\\n )\\n\\n return tuple(map(dispatch_lambda_arg, cpp_args))\\n\\n\\n# [old codegen] XXX: if you got here because of an assertion failure, it doesn't mean\\n# it's enough to just extend the list here. Before you do this, make sure\\n# to add an appropriate wrap() overload in torch/csrc/autograd/utils/wrap_outputs.h.\\nSUPPORTED_RETURN_TYPES = {\\n \\\"at::Tensor\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple\\\",\\n \\\"::std::tuple>\\\",\\n \\\"::std::vector\\\",\\n # Needed for flash attention forw/backward\\n \\\"::std::tuple\\\",\\n \\\"at::Scalar\\\",\\n \\\"bool\\\",\\n \\\"int64_t\\\",\\n \\\"void*\\\",\\n \\\"void\\\",\\n \\\"at::QScheme\\\",\\n \\\"double\\\",\\n \\\"at::IntArrayRef\\\",\\n \\\"at::ScalarType\\\",\\n \\\"at::Stream\\\",\\n}\\n\\n\\ndef dispatch_lambda_return_str(f: NativeFunction) -> str:\\n # [old codegen] Remove type annotation (e.g. 'Tensor' rather than 'Tensor &')\\n # because the dispatch lambdas take mutable arguments *by value*, not\\n # by reference. If you then return a reference to such an argument, you\\n # will now have a pointer to a dangling stack entry. Not good.\\n #\\n # You want:\\n #\\n # auto dispatch_selu_ = [](Tensor self) -> Tensor { ...; return at::selu_(self); };\\n # ^^^^^^\\n #\\n # *not*\\n #\\n # auto dispatch_selu_ = [](Tensor self) -> Tensor& { ...; return at::selu_(self); };\\n # ^^^^^^^\\n #\\n # (NB: We can't make dispatch_selu_ take Tensor&, because the enclosing\\n # codegen looks like dispatch_selu_(_r.tensor(0)), and you can't take a\\n # mutable reference to temporary. Maybe we could assign it to a\\n # variable itself.)\\n returns_without_annotation = tuple(\\n Return(r.name, r.type, None) for r in f.func.returns\\n )\\n return_str = cpp.returns_type(returns_without_annotation, symint=True).cpp_type()\\n if return_str not in SUPPORTED_RETURN_TYPES:\\n raise RuntimeError(f\\\"{f.func.name} returns unsupported type {return_str}\\\")\\n return return_str\\n\\n\\ndef cpp_dispatch_target(f: NativeFunction) -> str:\\n symint = f.func.has_symint()\\n name = cpp.name(f.func, symint_overload=symint)\\n if Variant.method in f.variants:\\n return f\\\"self.{name}\\\"\\n if Variant.function in f.variants:\\n if has_tensor_options(f) or f.func.name.name.base.endswith(\\\"_like\\\"):\\n namespace = \\\"torch\\\"\\n else:\\n namespace = \\\"at\\\"\\n return f\\\"{namespace}::{name}\\\"\\n raise RuntimeError(f\\\"could not dispatch, neither function nor method: {f.func}\\\")\\n\\n\\ndef cpp_dispatch_exprs(\\n f: NativeFunction,\\n *,\\n python_signature: PythonSignature | None = None,\\n) -> tuple[str, ...]:\\n cpp_args: Sequence[Binding] = _cpp_signature(f, method=False).arguments()\\n\\n exprs: tuple[str, ...] = ()\\n if not isinstance(python_signature, PythonSignatureDeprecated):\\n # By default the exprs are consistent with the C++ signature.\\n exprs = tuple(a.name for a in cpp_args)\\n else:\\n # For deprecated python signature we may need fill in some constants.\\n exprs = tuple(\\n filter(\\n lambda n: n != \\\"out\\\" or f.func.is_out_fn(),\\n python_signature.deprecated_args_exprs,\\n )\\n )\\n\\n if Variant.method in f.variants:\\n exprs = tuple(filter(\\\"self\\\".__ne__, exprs))\\n\\n return exprs\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Python / C++ Args Binding\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\n# We explicitly enumerate the PythonArgParser unpacking methods for all\\n# supported types. This might be more verbose than necessary, partially\\n# because of the irregularity of unpacking method naming, partially\\n# because we want to mimic the old codegen behavior - to reject\\n# unexpected and/or unsupported cases which the old codegen rejects.\\n# For certain cases it is intentionally more restrictive than necessary,\\n# e.g.: it doesn't accepts doublelist with definite size.\\ndef arg_parser_unpack_method(\\n t: Type, default: str | None, default_init: str | None, *, symint: bool = True\\n) -> str:\\n has_default_init = default_init is not None\\n if has_default_init and str(t) not in (\\n \\\"ScalarType?\\\",\\n \\\"ScalarType\\\",\\n \\\"Device\\\",\\n \\\"Device?\\\",\\n \\\"Layout\\\",\\n \\\"Layout?\\\",\\n \\\"bool\\\",\\n \\\"bool?\\\",\\n ):\\n raise RuntimeError(f\\\"type '{t}' does not supported unpacking with default\\\")\\n\\n if isinstance(t, BaseType):\\n if t.name in [\\n BaseTy.Tensor,\\n BaseTy.Stream,\\n BaseTy.Storage,\\n BaseTy.Scalar,\\n BaseTy.Dimname,\\n ]:\\n # These unpack methods line up with their schema names\\n return t.name.name.lower()\\n elif t.name == BaseTy.ScalarType:\\n return \\\"scalartypeWithDefault\\\" if has_default_init else \\\"scalartype\\\"\\n elif t.name == BaseTy.Device:\\n return \\\"deviceWithDefault\\\" if has_default_init else \\\"device\\\"\\n elif t.name == BaseTy.DeviceIndex:\\n return \\\"toInt64\\\"\\n elif t.name == BaseTy.int:\\n return \\\"toInt64\\\"\\n elif t.name == BaseTy.SymInt:\\n return \\\"toSymInt\\\" if symint else \\\"toInt64\\\"\\n elif t.name == BaseTy.bool:\\n return \\\"toBoolWithDefault\\\" if has_default_init else \\\"toBool\\\"\\n elif t.name == BaseTy.float:\\n return \\\"toDouble\\\"\\n elif t.name == BaseTy.str:\\n return \\\"stringView\\\"\\n elif t.name == BaseTy.Layout:\\n return \\\"layoutWithDefault\\\" if has_default_init else \\\"layout\\\"\\n elif t.name == BaseTy.MemoryFormat:\\n return \\\"memoryformat\\\"\\n\\n elif isinstance(t, OptionalType):\\n if str(t.elem) == \\\"Tensor\\\":\\n return \\\"optionalTensor\\\"\\n elif str(t.elem) == \\\"Generator\\\":\\n return \\\"generator\\\"\\n elif str(t.elem) == \\\"Dimname[]\\\":\\n return \\\"toDimnameListOptional\\\"\\n elif not has_default_init and default in (\\n None,\\n \\\"None\\\",\\n \\\"::std::nullopt\\\",\\n \\\"std::nullopt\\\",\\n ):\\n # If default is None: append 'Optional' to elem's unpacking method\\n return (\\n arg_parser_unpack_method(t.elem, None, None, symint=symint) + \\\"Optional\\\"\\n )\\n else:\\n # Otherwise, load as underlying type with default\\n return arg_parser_unpack_method(\\n t.elem, default, default_init, symint=symint\\n )\\n\\n elif isinstance(t, ListType):\\n if str(t.elem) == \\\"Tensor\\\":\\n # accept and use definite size\\n return f\\\"tensorlist_n<{t.size}>\\\" if t.size is not None else \\\"tensorlist\\\"\\n elif str(t.elem) == \\\"Tensor?\\\":\\n return \\\"list_of_optional_tensors\\\"\\n elif str(t.elem) == \\\"Dimname\\\":\\n # accept definite size\\n return \\\"dimnamelist\\\"\\n elif str(t.elem) == \\\"int\\\":\\n # accept definite size\\n return \\\"intlist\\\"\\n elif str(t.elem) == \\\"float\\\":\\n return \\\"doublelist\\\"\\n elif str(t.elem) == \\\"SymInt\\\":\\n # accept definite size\\n return \\\"symintlist\\\" if symint else \\\"intlist\\\"\\n elif str(t.elem) == \\\"Scalar\\\":\\n return \\\"scalarlist\\\"\\n raise RuntimeError(f\\\"type '{t}' is not supported by PythonArgParser\\\")\\n\\n\\n# Return RHS expression for python argument using PythonArgParser output.\\n# e.g. for arg name 'foo', arg type 'bool', arg_index = 2, returns '_r.toBool(2)'\\ndef arg_parser_output_expr(\\n arg_index: int, a: PythonArgument, *, symint: bool = True\\n) -> PythonArgParserOutputExpr:\\n has_default = a.default_init is not None\\n unpack_method = arg_parser_unpack_method(\\n t=a.type, default=a.default, default_init=a.default_init, symint=symint\\n )\\n default = f\\\", {a.default_init}\\\" if has_default else \\\"\\\"\\n expr = f\\\"_r.{unpack_method}({arg_index}{default})\\\"\\n\\n return PythonArgParserOutputExpr(\\n name=a.name,\\n expr=expr,\\n index=arg_index,\\n argument=a,\\n )\\n\\n\\n# Returns a map with key = arg_name and value = PythonArgParserOutputExpr.\\ndef arg_parser_output_exprs(\\n ps: PythonSignature, f: NativeFunction, *, symint: bool = True\\n) -> dict[str, PythonArgParserOutputExpr]:\\n return {\\n e.name: e\\n for i, a in enumerate(ps.arguments())\\n for e in (arg_parser_output_expr(i, a, symint=symint),)\\n }\\n\\n\\n# argument name to type for scattered tensor options fields\\nTENSOR_OPTIONS_FIELDS = {\\n \\\"dtype\\\": \\\"ScalarType?\\\",\\n \\\"device\\\": \\\"Device?\\\",\\n \\\"layout\\\": \\\"Layout?\\\",\\n \\\"pin_memory\\\": \\\"bool?\\\",\\n \\\"requires_grad\\\": \\\"bool?\\\",\\n}\\n\\n\\n# bind arg parser outputs (python args) with dispatch lambda arguments (c++ args).\\ndef dispatch_lambda_exprs(\\n ps: PythonSignature, f: NativeFunction, *, symint: bool = True\\n) -> DispatchLambdaArgumentExprs:\\n # This method is to bind 'arg_parser_outputs' and 'lambda_args' by producing\\n # 'inits' and 'lambda_args_exprs' for each lambda argument using arg parser\\n # outputs.\\n arg_parser_outputs = arg_parser_output_exprs(ps, f, symint=symint)\\n lambda_args = dispatch_lambda_args(ps, f, symint=symint)\\n inits: list[str] = []\\n lambda_args_exprs: dict[str, str] = {}\\n\\n has_toptions = has_tensor_options(f)\\n\\n # 1. special inits/unpacking to provide binding exprs for lambda arguments.\\n for a in ps.arguments(skip_tensor_options=True):\\n name = a.name\\n arg_parser_expr = arg_parser_outputs[a.name].expr\\n\\n if has_toptions and name == \\\"self\\\":\\n # TODO: why this needs to be special case?\\n inits.extend(\\n [\\n f\\\"auto self = {arg_parser_expr};\\\",\\n ]\\n )\\n lambda_args_exprs[name] = name\\n elif (\\n isinstance(a, PythonOutArgument)\\n and len(a.outputs) > 1\\n and f.func.is_out_fn()\\n ):\\n inits.extend(\\n [\\n f\\\"auto out = {arg_parser_expr};\\\",\\n ]\\n )\\n for i, out_arg in enumerate(a.outputs):\\n lambda_args_exprs[out_arg.name] = f\\\"out[{i}]\\\"\\n elif str(a.type) == \\\"Dimname[]?\\\":\\n # [old codegen]\\n # TODO: make this part of something more general, or get rid of it.\\n # optional> are special. The PythonArgParser returns an\\n # optional>, which cannot be implicitly converted to\\n # optional>. One needs to unwrap the optional and rewrap.\\n inits.extend(\\n [\\n f\\\"auto __{name} = {arg_parser_expr};\\\",\\n f\\\"::std::optional {name} = __{name} ? ::std::make_optional(DimnameList(__{name}.value())) : ::std::nullopt;\\\", # noqa: B950\\n ]\\n )\\n lambda_args_exprs[name] = name\\n else:\\n # default case - directly using PythonArgParser output expr\\n lambda_args_exprs[name] = arg_parser_expr\\n\\n # method's self is passed directly to python binding, rather than parsed\\n if ps.method:\\n lambda_args_exprs[\\\"self\\\"] = \\\"self\\\"\\n\\n # 2. special packing/checking for TensorOptions.\\n tensor_options_args_names = [a.name for a in ps.tensor_options_args]\\n if has_toptions:\\n if f.func.is_out_fn():\\n raise RuntimeError(f\\\"{f.func}: tensor options with output arg\\\")\\n for a in ps.tensor_options_args:\\n if a.name not in TENSOR_OPTIONS_FIELDS:\\n raise RuntimeError(\\n f\\\"{f.func}: unrecognized tensor options field '{a.name}' in python binding arguments\\\"\\n )\\n if str(a.type) != TENSOR_OPTIONS_FIELDS.get(a.name):\\n raise RuntimeError(\\n f\\\"{f.func}: unrecognized type '{str(a.type)}' for tensor options field '{a.name}'\\\"\\n )\\n if not all(a in tensor_options_args_names for a in TENSOR_OPTIONS_FIELDS):\\n raise RuntimeError(\\n f\\\"{f.func}: incomplete tensor options args: {tensor_options_args_names}\\\"\\n )\\n\\n inits.append(\\n f\\\"\\\"\\\"\\\\\\nconst auto options = TensorOptions()\\n .dtype({arg_parser_outputs['dtype'].expr})\\n .device({arg_parser_outputs['device'].expr})\\n .layout({arg_parser_outputs['layout'].expr})\\n .requires_grad({arg_parser_outputs['requires_grad'].expr})\\n .pinned_memory({arg_parser_outputs['pin_memory'].expr});\\ntorch::utils::maybe_initialize_device(options);\\n\\\"\\\"\\\"\\n )\\n lambda_args_exprs[\\\"options\\\"] = \\\"options\\\"\\n\\n # 3. special case - access scattered TensorOptions fields without packing\\n # TODO: maybe move to the generator side as it's not related to binding.\\n if not has_toptions and tensor_options_args_names:\\n if \\\"dtype\\\" in tensor_options_args_names:\\n # we're an output-arg variant, check these args against output tensor\\n if not f.func.is_out_fn():\\n raise RuntimeError(\\n f\\\"{f.func}: dtype in tensor_options_args without output arg, {ps} {ps.arguments}\\\"\\n )\\n if not all(a in tensor_options_args_names for a in (\\\"layout\\\", \\\"device\\\")):\\n raise RuntimeError(\\n f\\\"{f.func}: incomplete tensor options for output check\\\"\\n )\\n\\n inits.append(\\n f\\\"\\\"\\\"\\\\\\ncheck_out_type_matches({arg_parser_outputs['out'].expr}, {arg_parser_outputs['dtype'].expr},\\n {arg_parser_outputs['dtype'].is_none_expr}, {arg_parser_outputs['layout'].expr},\\n {arg_parser_outputs['device'].expr}, {arg_parser_outputs['device'].is_none_expr});\\n\\\"\\\"\\\"\\n )\\n # we'll set requires_grad on outgoing tensor\\n if \\\"requires_grad\\\" not in tensor_options_args_names:\\n raise RuntimeError(\\n f'{f.func}: expected \\\"requires_grad\\\" in tensor_options_args absent, but found [{tensor_options_args_names}]'\\n )\\n\\n return DispatchLambdaArgumentExprs(\\n exprs=tuple(lambda_args_exprs[a.name] for a in lambda_args),\\n inits=inits,\\n )\\n\\n\\nfrom __future__ import annotations\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.types import (\\n ArgName,\\n ArrayRefCType,\\n BaseCType,\\n Binding,\\n ConstRefCType,\\n dimnameListT,\\n intArrayRefT,\\n iOptTensorListRefT,\\n iTensorListRefT,\\n NamedCType,\\n OptionalCType,\\n optionalIntArrayRefT,\\n optionalScalarRefT,\\n optionalTensorRefT,\\n scalarT,\\n tensorT,\\n)\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n ListType,\\n NativeFunctionsGroup,\\n OptionalType,\\n SelfArgument,\\n TensorOptionsArguments,\\n Type,\\n)\\nfrom torchgen.utils import assert_never\\n\\n\\n# This file describes the translation of JIT schema to the structured functions API.\\n# This is similar to native API, but a number of historical problems with native\\n# API have been fixed.\\n\\n\\n# Translation of types occurring in JIT arguments to a C++ argument type.\\n# NB: For now, mutable doesn't do anything; but it could if we make\\n# some more nominal types\\ndef argumenttype_type(t: Type, *, mutable: bool, binds: ArgName) -> NamedCType:\\n # If it's a value type, do the value type translation\\n # NB: structured kernels ALWAYS have symint off, since they involve actual\\n # kernels that require real ints. The one exception is the\\n # CompositeExplicitAutograd and the meta function (which could\\n # hypothetically be SymInt), but for simplicity we plan for these to just\\n # be handled in Python\\n r = cpp.valuetype_type(t, symint=False, binds=binds, mutable=mutable)\\n if r is not None:\\n return r\\n\\n if isinstance(t, BaseType):\\n if t.name == BaseTy.Tensor:\\n return NamedCType(binds, ConstRefCType(BaseCType(tensorT)))\\n elif t.name == BaseTy.Scalar:\\n return NamedCType(binds, ConstRefCType(BaseCType(scalarT)))\\n else:\\n raise AssertionError(f\\\"base type should have been value type {t}\\\")\\n elif isinstance(t, OptionalType):\\n if t.elem == BaseType(BaseTy.Tensor):\\n return NamedCType(binds, BaseCType(optionalTensorRefT))\\n elif t.elem == BaseType(BaseTy.Scalar):\\n return NamedCType(binds, BaseCType(optionalScalarRefT))\\n elif isinstance(t.elem, ListType) and str(t.elem.elem) == \\\"int\\\":\\n return NamedCType(binds, BaseCType(optionalIntArrayRefT))\\n elem = argumenttype_type(t.elem, mutable=mutable, binds=binds)\\n return NamedCType(binds, OptionalCType(elem.type))\\n elif isinstance(t, ListType):\\n if t.elem == BaseType(BaseTy.Tensor):\\n return NamedCType(binds, ConstRefCType(BaseCType(iTensorListRefT)))\\n elif t.elem == OptionalType(BaseType(BaseTy.Tensor)):\\n return NamedCType(binds, BaseCType(iOptTensorListRefT))\\n # TODO: delete these special cases; see torchgen.api.cpp--these\\n # must be changed in tandem, but there are problems; see\\n # https://github.com/pytorch/pytorch/pull/51485\\n elif str(t.elem) == \\\"int\\\":\\n return NamedCType(binds, BaseCType(intArrayRefT))\\n elif str(t.elem) == \\\"Dimname\\\":\\n return NamedCType(binds, BaseCType(dimnameListT))\\n elem = argumenttype_type(t.elem, mutable=mutable, binds=binds)\\n return NamedCType(binds, ArrayRefCType(elem.type))\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(t)}\\\")\\n\\n\\ndef argument_type(a: Argument, *, binds: ArgName) -> NamedCType:\\n return argumenttype_type(a.type, mutable=a.is_write, binds=binds)\\n\\n\\n# returns_type intentionally omitted, because structured kernels never \\\"return\\\";\\n# instead, they always indirectly report their outputs (in the case of a meta\\n# function, by calling set_output; in the case of an impl function, by writing\\n# directly into the provided out argument).\\n\\n\\n# Structured kernels are never defaulted\\ndef argument(a: Argument | SelfArgument | TensorOptionsArguments) -> list[Binding]:\\n if isinstance(a, Argument):\\n return [\\n Binding(\\n nctype=argument_type(a, binds=a.name),\\n name=a.name,\\n default=None,\\n argument=a,\\n )\\n ]\\n elif isinstance(a, SelfArgument):\\n return argument(a.argument)\\n elif isinstance(a, TensorOptionsArguments):\\n raise AssertionError(\\\"structured kernels don't support TensorOptions yet\\\")\\n else:\\n assert_never(a)\\n\\n\\ndef impl_arguments(g: NativeFunctionsGroup) -> list[Binding]:\\n args: list[Argument | TensorOptionsArguments | SelfArgument] = []\\n\\n if g.out.precomputed:\\n # A list of parameters for the impl function with\\n # certain parameters replaced with precomputed counterparts\\n # as specified in native_functions.yaml.\\n non_out_args_replaced: list[\\n Argument | TensorOptionsArguments | SelfArgument\\n ] = []\\n for a in g.out.func.arguments.non_out:\\n if isinstance(a, Argument) and a.name in g.out.precomputed.replace:\\n # If a is in precompute.replace, append the parameters\\n # that should replace it onto non_out_args_replaced.\\n non_out_args_replaced.extend(g.out.precomputed.replace[a.name])\\n else:\\n # If not, push a as it is.\\n non_out_args_replaced.append(a)\\n\\n args.extend(non_out_args_replaced)\\n # g.out.precomputed.add is the list of parameters that are added\\n # without replacement after the non out args and just before the out args\\n args.extend(g.out.precomputed.add)\\n else:\\n args.extend(g.out.func.arguments.non_out)\\n\\n args.extend(g.out.func.arguments.out)\\n return [r for arg in args for r in argument(arg)]\\n\\n\\ndef meta_arguments(g: NativeFunctionsGroup) -> list[Binding]:\\n args: list[Argument | TensorOptionsArguments | SelfArgument] = []\\n args.extend(g.functional.func.arguments.non_out)\\n return [r for arg in args for r in argument(arg)]\\n\\n\\ndef out_arguments(g: NativeFunctionsGroup) -> list[Binding]:\\n args: list[Argument | TensorOptionsArguments | SelfArgument] = []\\n args.extend(g.out.func.arguments.out)\\n return [r for arg in args for r in argument(arg)]\\n\\n\\nfrom __future__ import annotations\\n\\nimport itertools\\nfrom typing import Sequence\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.types import ArgName, Binding, CType, NamedCType\\nfrom torchgen.model import (\\n Argument,\\n FunctionSchema,\\n Return,\\n SelfArgument,\\n TensorOptionsArguments,\\n Type,\\n)\\nfrom torchgen.utils import assert_never, concatMap\\n\\n\\n# This file describes the translation of JIT schema to the dispatcher\\n# API, the *unboxed* calling convention by which invocations through\\n# the dispatcher are made. Historically, the dispatcher API matched\\n# the C++ API, but with the establishment of the boxed API, we've\\n# made changes to the dispatcher API to so that the unboxed API\\n# better aligns with the boxed API. The dispatcher API hooks heavily\\n# into our template based boxing/unboxing machinery, so changes\\n# to this convention will usually need template updates too.\\n#\\n# Prominent characteristics of the dispatcher API:\\n#\\n# - dtype, layout, device and pin_memory are represented as separate\\n# arguments.\\n#\\n\\n\\ndef name(func: FunctionSchema) -> str:\\n return cpp.name(func)\\n\\n\\ndef argumenttype_type(\\n t: Type,\\n *,\\n mutable: bool,\\n binds: ArgName,\\n remove_non_owning_ref_types: bool = False,\\n symint: bool = True,\\n) -> NamedCType:\\n # This is a faux amis. If it makes sense in the future to add\\n # more special cases here, or invert things so cpp.argument_type\\n # calls this, or just completely inline the function, please do\\n # it.\\n return cpp.argumenttype_type(\\n t,\\n mutable=mutable,\\n binds=binds,\\n symint=symint,\\n remove_non_owning_ref_types=remove_non_owning_ref_types,\\n )\\n\\n\\ndef argument_type(\\n a: Argument,\\n *,\\n binds: ArgName,\\n remove_non_owning_ref_types: bool = False,\\n symint: bool = True,\\n) -> NamedCType:\\n return argumenttype_type(\\n a.type,\\n mutable=a.is_write,\\n binds=binds,\\n remove_non_owning_ref_types=remove_non_owning_ref_types,\\n symint=symint,\\n )\\n\\n\\ndef returns_type(rs: Sequence[Return], *, symint: bool = True) -> CType:\\n # At present, there is no difference. But there could be!\\n return cpp.returns_type(rs, symint=symint)\\n\\n\\ndef jit_arguments(func: FunctionSchema) -> list[Argument]:\\n def to_argument(\\n a: Argument | TensorOptionsArguments | SelfArgument,\\n ) -> list[Argument]:\\n if isinstance(a, Argument):\\n return [a]\\n elif isinstance(a, SelfArgument):\\n return [a.argument]\\n elif isinstance(a, TensorOptionsArguments):\\n return [a.dtype, a.layout, a.device, a.pin_memory]\\n else:\\n assert_never(a)\\n\\n return list(\\n concatMap(\\n to_argument,\\n itertools.chain(\\n func.arguments.positional, func.arguments.kwarg_only, func.arguments.out\\n ),\\n )\\n )\\n\\n\\ndef argument(\\n a: Argument, *, remove_non_owning_ref_types: bool = False, symint: bool = True\\n) -> Binding:\\n return Binding(\\n nctype=argument_type(\\n a,\\n binds=a.name,\\n remove_non_owning_ref_types=remove_non_owning_ref_types,\\n symint=symint,\\n ),\\n name=a.name,\\n argument=a,\\n )\\n\\n\\ndef arguments(func: FunctionSchema, *, symint: bool = True) -> list[Binding]:\\n return [argument(a, symint=symint) for a in jit_arguments(func)]\\n\\n\\nfrom __future__ import annotations\\n\\nimport re\\nfrom dataclasses import dataclass\\nfrom typing import cast, Sequence\\n\\nfrom torchgen import local\\nfrom torchgen.api import cpp\\nfrom torchgen.api.types import BaseCType, Binding, NamedCType, tensorListT\\nfrom torchgen.model import (\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n ListType,\\n NativeFunction,\\n NativeFunctionsViewGroup,\\n SchemaKind,\\n Type,\\n)\\nfrom torchgen.utils import IDENT_REGEX\\n\\n\\n# Represents a saved attribute involved in backward calculation.\\n# Note that it can be a derived property of an input argument, e.g.:\\n# we could save `other.scalar_type()` instead of the entire `other` tensor.\\n@dataclass(frozen=True)\\nclass SavedAttribute:\\n # The NamedCType holds the updated name and cpp type of the attribute\\n # for the name, Suffix is appended if it's derived property, e.g.: `other_scalar_type`\\n nctype: NamedCType\\n\\n # The expression to read the derived property at save time, e.g.:\\n # `other.scalar_type()`.\\n expr: str\\n\\n\\n# Represents a backward formula that calculates derivatives for one\\n# or more tensors.\\n@dataclass(frozen=True)\\nclass Derivative:\\n # The formula string (legit C++ expression).\\n # Note that expressions against input arguments have been replaced with the\\n # corresponding saved attributes.\\n # E.g.:\\n # raw formula: `mul_tensor_backward(grad, self, other.scalar_type())`\\n # here: `mul_tensor_backward(grad, self, other_scalar_type)`\\n formula: str\\n\\n # The formula string before input argument replacement\\n original_formula: str\\n\\n # Names of the arguments for which this formula calculates derivatives.\\n var_names: tuple[str, ...]\\n\\n # Saved inputs that are referenced by the formula.\\n saved_inputs: tuple[SavedAttribute, ...]\\n\\n # Saved outputs that are referenced by the formula.\\n saved_outputs: tuple[SavedAttribute, ...]\\n\\n # Gradients that are referenced by name in the formula.\\n named_gradients: set[str]\\n\\n\\n# Represents a forward formula that calculates forward derivatives\\n# for one tensor.\\n@dataclass(frozen=True)\\nclass ForwardDerivative:\\n # The formula string (legit C++ expression).\\n # Note that special keywords such as \\\"linear\\\" or \\\"element_wise\\\" have been\\n # replaced by the automatically generated formula.\\n formula: str\\n\\n # Name of the output arguments for which this formula calculates forward\\n # derivatives\\n var_names: tuple[str, ...]\\n\\n # Type of the output arguments for which this formula calculates forward\\n # derivatives\\n var_types: tuple[Type, ...]\\n\\n # Inputs for which the forward derivatives are required for this formula\\n required_inputs_fw_grad: tuple[str, ...] | None\\n\\n # Inputs for which the primal is required for this formula\\n required_inputs_primal: tuple[str, ...] | None\\n\\n # Flag to specify if this formula requires the original value of self\\n # This is only used by inplace operations\\n required_original_self_value: bool\\n\\n # If this formula is specified in derivatives.yaml or if we are re-using the\\n # out of place formula for inplace\\n is_reusing_outplace_formula: bool\\n\\n\\n# Represents differentiability info for a NativeFunction.\\n@dataclass(frozen=True)\\nclass DifferentiabilityInfo:\\n # The base name read from derivatives.yaml.\\n name: str\\n\\n # The matching native function.\\n #\\n # There can be multiple NativeFunction having the same base name:\\n # - different overloads with different types of input arguments;\\n # - in-place/out/functional variants of the same function;\\n #\\n # We first use the schema string (under the 'name' key) in derivatives.yaml\\n # to find the NativeFunction having the same schema string.\\n # Then we find the in-place/out/functional variants of the matching function.\\n # Among these variants, we choose the one having the same name as the\\n # derivatives.yaml entry. If there is no exact match, then we choose the\\n # in-place variant.\\n # TODO: maybe the logic to search for all variants is no longer necessary?\\n func: NativeFunction\\n\\n # The name of the generated autograd function.\\n # It's set only if we will calculate a derivative, i.e.\\n # 'args_with_derivatives' is not empty.\\n op: str | None\\n\\n # The derivatives formulae for this function.\\n # Note that the length of this sequence is the number of differentiable inputs\\n derivatives: Sequence[Derivative]\\n\\n # The forward derivatives formulae for this function.\\n # Note that the length of this sequence is the number of differentiable outputs\\n forward_derivatives: Sequence[ForwardDerivative]\\n\\n # The union of 'saved_inputs' of all 'derivatives'.\\n all_saved_inputs: Sequence[SavedAttribute]\\n\\n # The union of 'saved_outputs' of all 'derivatives'.\\n all_saved_outputs: Sequence[SavedAttribute]\\n\\n # All named gradients that are available for use, in the same\\n # order as in the grads vector.\\n available_named_gradients: Sequence[str]\\n\\n # The named gradients that are used in any of the derivatives.\\n # Invariant: all(name in available_named_gradients for name in used_named_gradients)\\n used_named_gradients: set[str]\\n\\n # The function's input arguments for which it calculates derivatives.\\n # It's the union of 'var_names' of all 'derivatives', sorted by the\\n # argument order in the function schema.\\n args_with_derivatives: Sequence[Binding]\\n\\n # Names of arguments whose derivative formula is 'non_differentiable'.\\n non_differentiable_arg_names: Sequence[str]\\n\\n # Raw data read from derivatives.yaml.\\n output_differentiability: list[bool] | None\\n\\n # output_differentiability in derivatives.yaml can be a list of\\n # conditions that express if the output is differentiable. In this case,\\n # the number of conditions must match the number of outputs\\n # (NB: we only support one condition right now).\\n # output_differentiability gets populated with True for each condition,\\n # while output_differentiability_conditions gets populated with the conditions\\n output_differentiability_conditions: list[str] | None\\n\\n @property\\n def has_derivatives(self) -> bool:\\n return len(self.args_with_derivatives) > 0\\n\\n # Generates a new DifferentiabilityInfo using the exact same set of derivative information,\\n # but with a new operator name.\\n # This is used when generating \\\"copy\\\" variants of view ops,\\n # which are able to use the exact same derivative formula as the original view op\\n # See Note [Codegen'd {view}_copy Operators]\\n def create_view_copy_from_view_derivative(\\n self, g: NativeFunctionsViewGroup\\n ) -> DifferentiabilityInfo | None:\\n if g.view_copy is None:\\n return None\\n f = g.view_copy\\n\\n name_split_by_period = self.name.split(\\\".\\\", maxsplit=2)\\n # Append a \\\"_copy\\\" to the base name of the operator (but keep the overload name the same)\\n view_copy_name = f\\\"{name_split_by_period[0]}_copy.\\\" + \\\".\\\".join(\\n name_split_by_period[1:]\\n )\\n view_copy_op_name = None if self.op is None else f\\\"{self.op}_copy\\\"\\n\\n return DifferentiabilityInfo(\\n # Use the \\\"_copy\\\" version of name/func/op\\n name=view_copy_name,\\n func=f,\\n op=view_copy_op_name,\\n # But keep all derivative info the same\\n derivatives=self.derivatives,\\n forward_derivatives=self.forward_derivatives,\\n all_saved_inputs=self.all_saved_inputs,\\n all_saved_outputs=self.all_saved_outputs,\\n available_named_gradients=self.available_named_gradients,\\n used_named_gradients=self.used_named_gradients,\\n args_with_derivatives=self.args_with_derivatives,\\n non_differentiable_arg_names=self.non_differentiable_arg_names,\\n output_differentiability=self.output_differentiability,\\n output_differentiability_conditions=self.output_differentiability_conditions,\\n )\\n\\n\\ndef uses_ident(info: DifferentiabilityInfo | None, ident: str) -> bool:\\n if info is None:\\n return False\\n for derivative in info.derivatives:\\n formula = derivative.formula\\n if re.search(IDENT_REGEX.format(ident), formula):\\n return True\\n return False\\n\\n\\ndef uses_retain_variables(info: DifferentiabilityInfo | None) -> bool:\\n return uses_ident(info, \\\"retain_variables\\\")\\n\\n\\ndef uses_single_grad(info: DifferentiabilityInfo | None) -> bool:\\n return uses_ident(info, \\\"grad\\\")\\n\\n\\n# Represents a differentiable `Argument`.\\n# How is it different from the `Argument` type?\\n# - It's processed Arguments which are differentiable and only used in the\\n# context of the autograd codegen;\\n# - It can represent SelfArgument or regular Argument but not TensorOptionsArgument;\\n@dataclass(frozen=True)\\nclass DifferentiableInput:\\n name: str\\n type: Type\\n\\n # TODO: only to keep it byte-for-byte compatible with the old codegen, should remove.\\n cpp_type: str\\n\\n\\n# Represents a differentiable `Return`.\\n# How it it different from the `Return` type?\\n# - The name in `Return` is optional. Here it is always populated using the same\\n# `cpp.return_names()` method.\\n# TODO: some cpp naming logic (e.g. resolving name conflict) might be irrelevant?\\n# - It's processed Returns which are differentiable, in compliance with the\\n# `output_differentiability` field defined in derivatives.yaml (if specified),\\n# and are only used in the context of the autograd codegen;\\n@dataclass(frozen=True)\\nclass DifferentiableOutput:\\n name: str\\n type: Type\\n\\n # TODO: only to keep it byte-for-byte compatible with the old codegen, should remove.\\n cpp_type: str\\n\\n\\n@dataclass(frozen=True)\\nclass NativeFunctionWithDifferentiabilityInfo:\\n func: NativeFunction\\n info: dict[str, DifferentiabilityInfo] | None\\n fw_derivatives: dict[str, Sequence[ForwardDerivative]] | None\\n\\n\\n# TODO: Update comment below since it is out of date.\\ndef dispatch_strategy(fn: NativeFunctionWithDifferentiabilityInfo) -> str:\\n \\\"\\\"\\\"How are we going to call the underlying implementation of a\\n declaration? There are two strategies:\\n - use_derived: we want to call the implementation on CPUDoubleType\\n (or a similar, derived Type instance). Because these derived\\n instances deal in Tensors, not Variables (it's a completely different\\n object, so it doesn't dispatch back to VariableType), code on\\n this dispatch path needs to wrap/unwrap tensors. If the\\n derived implementation takes and returns tensors, the\\n implementation is usually differentiable (although we also use\\n the derived dispatch path for non-differentiable functions\\n that we still want to dispatch on the derived Type instance;\\n e.g., size())\\n - use_type: we want to call the implementation on Type, because\\n it is implemented concretely, and the functions it invokes will\\n get dispatched back to VariableType (which will ensure that they\\n are differentiable.)\\n \\\"\\\"\\\"\\n # fn is derived as long as any of its per-key differentiability infos\\n # has_derivatives. dispatch_strategy() is used to guard generation of fns in VariableType\\n # and ADInplaceOrViewType. We want to generate these functions as long as a\\n # derivative is defined for ANY dispatch key.\\n if fn.func.is_abstract or (\\n fn.info is not None and any(info.has_derivatives for info in fn.info.values())\\n ):\\n # If the function is abstract (not implemented on at::Type), we must\\n # call the implementation on the derived type with unpacked tensors.\\n\\n # If the function has a derivative specified and is concrete, we could\\n # call either implementation. We prefer the calling the derived\\n # type's implementation with unpacked tensors because it is more\\n # performant in some cases: any internal calls to other ATen functions\\n # won't have the history tracked.\\n\\n # If the function has a type dispatched argument (i.e. is a factory),\\n # we prefer calling the derived type's implementation both because it is\\n # more performant and to ensure factory functions return tensors with _version\\n # of 0 (probably not strictly necessary, but nice to have to keeps versions simple\\n # to understand.\\n\\n return \\\"use_derived\\\"\\n else:\\n # If the function is concrete (we don't have to override it) and we\\n # didn't declare it in derivatives.yaml, we'll assume that it is\\n # actually implemented out of differentiable functions. (This\\n # assumption might not hold, but then you'll see gradcheck fail.)\\n return \\\"use_type\\\"\\n\\n\\ndef is_foreach_func(f: NativeFunction) -> bool:\\n return f.func.name.name.base.startswith(\\\"_foreach_\\\")\\n\\n\\n# note(crcrpar): Most foreach functions can reference an out-place `torch` function whose schema kind\\n# is functional for their backward derivatives (and forward derivatives in the future), i.e.,\\n# they would find such one in `functional_info_by_signature`. There however are some exceptions:\\n_foreach_with_inplace_ref = {\\\"_foreach_zero_\\\"}\\n_foreach_with_tensor_overload = {\\n \\\"_foreach_add.Tensor\\\",\\n \\\"_foreach_mul.Tensor\\\",\\n \\\"_foreach_div.Tensor\\\",\\n}\\n# The following do not support the alpha kwarg, which the nonforeach versions support.\\n_skip_argument_len_check = {\\n \\\"_foreach_add.Scalar\\\",\\n \\\"_foreach_add_.Scalar\\\",\\n \\\"_foreach_add.ScalarList\\\",\\n \\\"_foreach_add_.ScalarList\\\",\\n \\\"_foreach_sub.Scalar\\\",\\n \\\"_foreach_sub_.Scalar\\\",\\n \\\"_foreach_sub.ScalarList\\\",\\n \\\"_foreach_sub_.ScalarList\\\",\\n}\\n\\n\\n# Checks if `function_schema` is a native, non-foreach function which `f`, a foreach function\\n# reference to generate derivatives.\\ndef is_reference_for_foreach(\\n f: NativeFunction,\\n function_schema: FunctionSchema,\\n) -> bool:\\n return (\\n f.func.name.name.base.split(\\\"_foreach_\\\")[-1] == function_schema.name.name.base\\n and (\\n not function_schema.name.name.inplace\\n or str(f.func.name) in _foreach_with_inplace_ref\\n )\\n and (\\n str(f.func.name) in _skip_argument_len_check\\n or len(f.func.arguments.flat_non_out)\\n == len(function_schema.arguments.flat_non_out)\\n )\\n and all(\\n ref_arg.type in (arg.type, getattr(arg.type, \\\"elem\\\", None))\\n for arg, ref_arg in zip(\\n f.func.arguments.flat_non_out,\\n function_schema.arguments.flat_non_out,\\n )\\n )\\n )\\n\\n\\n# TODO(crcrpar): Avoid hard coding \\\"Default\\\" ideally.\\ndef gen_foreach_derivativeinfo(\\n foreach_function: NativeFunction,\\n functional_info_by_signature: dict[\\n FunctionSchema, dict[str, DifferentiabilityInfo]\\n ],\\n non_functional_info_by_signature: dict[\\n FunctionSchema, dict[str, DifferentiabilityInfo]\\n ],\\n dispatch_key: str = \\\"Default\\\",\\n) -> tuple[DifferentiabilityInfo | None, bool]:\\n \\\"\\\"\\\"Generate DifferentiabilityInfo for out-place foreach function, return the existing one for in-place.\\n\\n The second return value indicates whether the info is generated in this function.\\n \\\"\\\"\\\"\\n ref_diff_info: DifferentiabilityInfo | None = None\\n\\n for function_schema, diff_info in functional_info_by_signature.items():\\n if not is_reference_for_foreach(foreach_function, function_schema):\\n continue\\n ref_diff_info = diff_info[dispatch_key]\\n if ref_diff_info is not None:\\n break\\n # note(crcrpar): It seems like `zero`'s info isn't available in functional_info_by_signature\\n # while the info of `zero_` is in non_functional_info_by_signature\\n if (\\n ref_diff_info is None\\n and foreach_function.func.kind() == SchemaKind.inplace\\n and str(foreach_function.func.name) in _foreach_with_inplace_ref\\n ):\\n for function_schema, diff_info in non_functional_info_by_signature.items():\\n if not is_reference_for_foreach(foreach_function, function_schema):\\n continue\\n ref_diff_info = diff_info[dispatch_key]\\n if ref_diff_info is not None:\\n break\\n if ref_diff_info is None:\\n return None, False\\n\\n # non out-place uses the existing Derivative.\\n if foreach_function.func.kind() == SchemaKind.inplace:\\n return ref_diff_info, False\\n\\n map_refarg2foreacharg, map_name2arg = {}, {}\\n for i, (arg, ref_arg) in enumerate(\\n zip(\\n foreach_function.func.arguments.flat_non_out,\\n function_schema.arguments.flat_non_out,\\n )\\n ):\\n map_refarg2foreacharg[ref_arg.name] = arg.name\\n map_name2arg[arg.name] = arg\\n\\n all_saved_inputs, all_saved_outputs, all_var_names = [], [], []\\n modified_derivative_formulas = []\\n for i, derivative in enumerate(ref_diff_info.derivatives):\\n modified_formula = derivative.formula.replace(\\\"grad\\\", \\\"grads[i]\\\").replace(\\n \\\"result\\\", \\\"result[i]\\\"\\n )\\n saved_inputs, saved_outputs = [], []\\n # note(crcrpar): This context seems necessary to call `cpp.argument_type`\\n with local.parametrize(\\n use_const_ref_for_mutable_tensors=foreach_function.use_const_ref_for_mutable_tensors,\\n use_ilistref_for_tensor_lists=foreach_function.part_of_structured_group,\\n ):\\n for ref_input in derivative.saved_inputs:\\n ref_input_jit_name = ref_input.expr.split(\\\".\\\")[0]\\n mapped_name = map_refarg2foreacharg[ref_input_jit_name]\\n if isinstance(map_name2arg[mapped_name].type, ListType):\\n mapped_expr = mapped_name + \\\"[i]\\\"\\n else:\\n mapped_expr = mapped_name\\n new_expr = ref_input.expr.replace(ref_input_jit_name, mapped_expr)\\n modified_formula = modified_formula.replace(\\n cast(str, ref_input.nctype.name), new_expr\\n )\\n\\n nctype = cpp.argument_type(map_name2arg[mapped_name], binds=mapped_name)\\n canonical_nctype = NamedCType(\\n nctype.name, nctype.type.remove_const_ref()\\n )\\n saved_inputs.append(\\n SavedAttribute(nctype=canonical_nctype, expr=mapped_name)\\n )\\n for ref_output in derivative.saved_outputs:\\n if ref_output.nctype.name == \\\"result\\\":\\n saved_outputs.append(\\n SavedAttribute(\\n nctype=NamedCType(\\n name=\\\"result\\\", type=BaseCType(tensorListT)\\n ),\\n expr=\\\"result\\\",\\n )\\n )\\n else:\\n raise RuntimeError(\\\"\\\")\\n var_names = [map_refarg2foreacharg[var] for var in derivative.var_names]\\n all_var_names.extend(var_names)\\n all_saved_inputs.extend(saved_inputs)\\n all_saved_outputs.extend(saved_outputs)\\n modified_derivative = Derivative(\\n formula=modified_formula,\\n original_formula=derivative.formula,\\n var_names=tuple(var_names),\\n saved_inputs=tuple(saved_inputs),\\n saved_outputs=tuple(saved_outputs),\\n named_gradients=set(),\\n )\\n modified_derivative_formulas.append(modified_derivative)\\n\\n with local.parametrize(\\n use_const_ref_for_mutable_tensors=foreach_function.use_const_ref_for_mutable_tensors,\\n use_ilistref_for_tensor_lists=foreach_function.part_of_structured_group,\\n ):\\n args_with_derivatives = [\\n Binding(\\n name=arg.name,\\n nctype=cpp.argument_type(arg, binds=arg.name),\\n argument=arg,\\n default=None,\\n )\\n for arg in foreach_function.func.arguments.flat_non_out\\n if arg.name in all_var_names\\n ]\\n\\n forward_derivatives: list[ForwardDerivative] = []\\n fw_derivative: ForwardDerivative\\n for fw_derivative in ref_diff_info.forward_derivatives:\\n var_names: list[str] = list(fw_derivative.var_names) # type: ignore[no-redef]\\n var_types: list[Type] = list(fw_derivative.var_types)\\n required_inputs_fw_grad: list[str] = []\\n required_inputs_primal: list[str] = []\\n if fw_derivative.required_inputs_fw_grad is not None:\\n required_inputs_fw_grad = list(fw_derivative.required_inputs_fw_grad)\\n if fw_derivative.required_inputs_primal:\\n required_inputs_primal = list(fw_derivative.required_inputs_primal)\\n modified_formula = fw_derivative.formula\\n\\n # Foreach's result is TensorList\\n if \\\"result\\\" in modified_formula:\\n modified_formula = fw_derivative.formula.replace(\\\"result\\\", \\\"result[i]\\\")\\n\\n for foreach_arg, ref_arg in zip(\\n foreach_function.func.arguments.flat_non_out,\\n ref_diff_info.func.func.arguments.flat_non_out,\\n ):\\n # Modify reference forward formula\\n if (\\n isinstance(foreach_arg.type, ListType)\\n and not foreach_arg.type.is_tensor_like()\\n ):\\n # Assuming ScalarList\\n modified_formula = modified_formula.replace(\\n ref_arg.name, foreach_arg.name + \\\"[i]\\\"\\n )\\n elif foreach_arg.type.is_tensor_like():\\n # Assuming TensorList / Tensor\\n # assert isinstance(foreach_arg.type, ListType), f\\\"{foreach_function.func.name}, {foreach_arg.type}\\\"\\n assert isinstance(foreach_arg.type, ListType) or (\\n foreach_arg.type == BaseType(BaseTy.Tensor)\\n and str(foreach_function.func.name) in _foreach_with_tensor_overload\\n ), f\\\"{foreach_function.func.name}, {foreach_arg.type}\\\"\\n for suffix in (\\\"_p\\\", \\\"_t\\\"):\\n curr_expr = ref_arg.name + suffix\\n if curr_expr in modified_formula:\\n new_expr = foreach_arg.name + suffix\\n modified_formula = modified_formula.replace(curr_expr, new_expr)\\n else:\\n # Assuming Scalar\\n if foreach_arg.name != ref_arg.name:\\n modified_formula = modified_formula.replace(\\n ref_arg.name, foreach_arg.name\\n )\\n\\n # note(crcrpar): there should exist a cooler way...\\n for i, name in enumerate(var_names):\\n if name == ref_arg.name:\\n var_names[i] = foreach_arg.name\\n var_types[i] = foreach_arg.type\\n for i, name in enumerate(required_inputs_fw_grad):\\n if name == ref_arg.name:\\n required_inputs_fw_grad[i] = foreach_arg.name\\n for i, name in enumerate(required_inputs_primal):\\n if name == ref_arg.name:\\n required_inputs_primal[i] = foreach_arg.name\\n forward_derivatives.append(\\n ForwardDerivative(\\n formula=modified_formula,\\n var_names=tuple(var_names),\\n var_types=tuple(var_types),\\n required_inputs_fw_grad=tuple(required_inputs_fw_grad),\\n required_inputs_primal=tuple(required_inputs_primal),\\n required_original_self_value=fw_derivative.required_original_self_value,\\n is_reusing_outplace_formula=fw_derivative.is_reusing_outplace_formula,\\n )\\n )\\n\\n return (\\n DifferentiabilityInfo(\\n name=foreach_function.func.name.name.base,\\n func=foreach_function,\\n op=f\\\"Foreach{ref_diff_info.op}{foreach_function.func.name.overload_name}\\\",\\n derivatives=modified_derivative_formulas,\\n forward_derivatives=forward_derivatives,\\n all_saved_inputs=tuple(set(all_saved_inputs)),\\n all_saved_outputs=tuple(set(all_saved_outputs)),\\n available_named_gradients=(),\\n used_named_gradients=set(),\\n args_with_derivatives=args_with_derivatives,\\n non_differentiable_arg_names=[],\\n output_differentiability=None,\\n output_differentiability_conditions=None,\\n ),\\n True,\\n )\\n\\n\\ndef match_differentiability_info(\\n native_functions: list[NativeFunction],\\n differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]],\\n) -> list[NativeFunctionWithDifferentiabilityInfo]:\\n \\\"\\\"\\\"Sets the \\\"derivative\\\" key on declarations to matching autograd function\\n In-place functions will use the out-of-place derivative definition if there\\n is no in-place specific derivative.\\n \\\"\\\"\\\"\\n\\n functional_info_by_signature = {\\n schema.signature(strip_default=True): info_dict\\n for schema, info_dict in differentiability_infos.items()\\n if schema.kind() == SchemaKind.functional\\n }\\n non_functional_info_by_signature = {\\n schema.signature(strip_default=True): info_dict\\n for schema, info_dict in differentiability_infos.items()\\n if schema.kind() != SchemaKind.functional\\n }\\n\\n def find_info(\\n f: NativeFunction,\\n ) -> tuple[dict[str, DifferentiabilityInfo] | None, bool]:\\n # Don't bother matching info to generated out= variants\\n if \\\"generated\\\" in f.tags and f.func.kind() == SchemaKind.out:\\n return None, False\\n\\n # (1) Check for an exact match\\n if f.func in differentiability_infos:\\n return differentiability_infos[f.func], True\\n\\n # (2) If no exact match, check if the out-of-place variant\\n # of this operator has a match.\\n # i.e mul() for mul_() or mul_out()\\n # note(crcrpar): Check foreach or not because in-place foreach functions use backward defined for the existing\\n # native functions instead of the out-place counterparts.\\n f_sig = f.func.signature(strip_default=True)\\n if f_sig in functional_info_by_signature and not is_foreach_func(f):\\n return functional_info_by_signature[f_sig], False\\n\\n # (3) Some operators have a derivative explicitly defined for the mutable\\n # variant, but get a code-generated out-of-place variant which does *not*\\n # come with a derivative formula.\\n # For the generated out-of-place variant, use the mutable variant's formula\\n # if it exists.\\n if \\\"generated\\\" in f.tags and f_sig in non_functional_info_by_signature:\\n info_dict = non_functional_info_by_signature[f_sig]\\n # See https://github.com/pytorch/pytorch/pull/76320/files#r874816389\\n assert not any(\\n any(\\\"self\\\" in str(inpt.nctype.name) for inpt in info.all_saved_inputs)\\n for info in info_dict.values()\\n ), f\\\"\\\"\\\"\\\\\\nAttempted to convert a derivative formula for a mutable operator\\n to be used by automatically by its functional variant (\\\"{str(f.func)}\\\").\\n this is not currently supported (we'd need to fix up the formula in the codegen).\\\"\\\"\\\"\\n return info_dict, False\\n\\n # (4) Generate derivative information of foreach functions if none is defined in `derivatives.yaml`\\n if is_foreach_func(f):\\n assert f.func not in differentiability_infos\\n diff_info, is_generated = gen_foreach_derivativeinfo(\\n f,\\n functional_info_by_signature,\\n non_functional_info_by_signature,\\n )\\n if diff_info is None:\\n return None, False\\n # TODO(crcrpar): Avoid hard coding \\\"Default\\\" ideally.\\n diff_info_dict = {\\\"Default\\\": diff_info}\\n if is_generated:\\n differentiability_infos[f.func] = diff_info_dict\\n functional_info_by_signature[f.func] = diff_info_dict\\n return diff_info_dict, is_generated\\n\\n return None, False\\n\\n result: list[NativeFunctionWithDifferentiabilityInfo] = []\\n for f in native_functions:\\n info_dict, is_exact_match = find_info(f)\\n\\n # Currently, the '.strides()' to 'strides_or_error' replacement does not support\\n # 'self' derivatives of an inplace function, so we must check for this case.\\n if f.func.kind() == SchemaKind.inplace and (info_dict is not None):\\n for info in info_dict.values():\\n for derivative in info.derivatives:\\n if \\\"self\\\" in derivative.var_names:\\n for saved_input in derivative.saved_inputs:\\n assert \\\"strides_or_error\\\" not in saved_input.expr, (\\n \\\"Calling '.strides()' in the 'self' derivative formula of an \\\"\\n f\\\"in-place function is not supported: {f.func}\\\"\\n )\\n\\n if not info_dict:\\n result.append(\\n NativeFunctionWithDifferentiabilityInfo(\\n func=f, info=None, fw_derivatives=None\\n )\\n )\\n continue\\n\\n fw_derivative_dict: dict[str, Sequence[ForwardDerivative]] = {}\\n for key, info in info_dict.items():\\n if not info.forward_derivatives:\\n fw_derivative_dict[key] = []\\n continue\\n\\n forward_derivatives = info.forward_derivatives\\n\\n # For functions that have a single def for out-of-place and inplace (like abs())\\n if f.func.kind() == SchemaKind.inplace:\\n # For inplace functions there is a little bit of work to do:\\n # 1) Validate the formula and make sure the input that is modified in not used:\\n # - If there is a formula for the inplace variant of the function (is_exact_match == True) then\\n # we make sure that the original value of the input that is being modified inplace (self_p) is\\n # not used in the formula. Note that the formula can use \\\"original_self_p\\\" here and that would\\n # trigger a clone of the original input.\\n # - If we are re-using the out of place formula (is_exact_match == False) then we replace every\\n # occurrence of self_p and self_t by original_self_p and original_self_t. These will be\\n # populated by cloned version of the original input (either the clone done by the backward AD\\n # logic if self is also used in a backward formula or a special clone that we add).\\n # 2) At this point, there cannot be a self_p in the formula.\\n # 3) Change \\\"result\\\" into \\\"self_p\\\" as by design, in the inplace function codegen, the result is\\n # simply called self (as it is modified inplace).\\n # 4) Update the required primals data in case it used to contain \\\"result\\\" but should now contain\\n # \\\"self\\\"\\n # 5) If it is not an exact match, the user formula is not modifying the existing forward grad\\n # inplace as it should. So add some code that makes sure that we do so if the forward grad\\n # already exists.\\n\\n assert (\\n len(info.forward_derivatives) == 1\\n ) # Only single output inplace should exist\\n fw_info = info.forward_derivatives[0]\\n formula = fw_info.formula\\n\\n def replace_self_with_original_self(formula: str, postfix: str) -> str:\\n def repl(m: re.Match[str]) -> str:\\n return f\\\"{m.group(1)}original_self{postfix}{m.group(2)}\\\"\\n\\n return re.sub(IDENT_REGEX.format(f\\\"self{postfix}\\\"), repl, formula)\\n\\n if re.search(IDENT_REGEX.format(\\\"self_p\\\"), formula):\\n if is_exact_match:\\n # For manually defined formulas, don't allow the original value to be used\\n raise RuntimeError(\\n f'The formula for \\\"{f.func.name}\\\" is using the original value of self '\\n \\\"that is being modified inplace. This would lead to wrong forward gradients. \\\"\\n 'Please use \\\"result\\\" in the formula only.'\\n )\\n else:\\n # When the original formula is out of place, we save a clone of the primal\\n # value to be able to access this value if needed\\n # replace \\\"self_p\\\"/\\\"self_t\\\" from the formula by \\\"original_self_p\\\"/\\\"original_self_t\\\"\\n formula = replace_self_with_original_self(formula, \\\"_p\\\")\\n formula = replace_self_with_original_self(formula, \\\"_t\\\")\\n\\n # replace \\\"result\\\" from the formula by \\\"self_p\\\"\\n def repl(m: re.Match[str]) -> str:\\n return f\\\"{m.group(1)}self_p{m.group(2)}\\\"\\n\\n formula = re.sub(IDENT_REGEX.format(\\\"result\\\"), repl, formula)\\n\\n required_primals = fw_info.required_inputs_primal\\n if re.search(IDENT_REGEX.format(\\\"self_p\\\"), formula):\\n required_primals = (\\n required_primals + (\\\"self\\\",) if required_primals else (\\\"self\\\",)\\n )\\n\\n if not is_exact_match:\\n # NOTE [In-place forward AD formula Optimization]\\n #\\n # This optimization transforms the formula to directly do inplace, i.e.\\n # instead of self_t.copy_(self_t.op()) we do self_t.op_() when the following are met:\\n #\\n # 1) the formula satisfies the pattern: \\\"self_t.op(*args)\\\"\\n # 2) \\\"op\\\" in (1) needs to be the same as the op the derivative is for\\n #\\n # (2) may seem too strict, but currently the only ops that satisfy (1) also satisfy (2)\\n # If there is a need, we can relax (2) to allow any op that has an in-place variant\\n is_single_method_on_self_t = False\\n directly_do_inplace = False\\n op_name: str | None = None\\n between_parens: str | None = None\\n match = re.fullmatch(r\\\"self_t.([\\\\w]*)\\\\((.*)\\\\)\\\", formula)\\n if match:\\n op_name, between_parens = match.group(1), match.group(2)\\n\\n # We want to...\\n # Match: self_t.op1(other_p.op2(arg))\\n # Avoid: self_t.op1(args) + self_t.op2(args)\\n # Avoid: self_t.op1(other_p.op2(arg)) + self_t.op2(args)\\n def check_parens_nest_level_gt_zero(s: str) -> bool:\\n level = 1\\n for ch in s:\\n if ch == \\\")\\\":\\n level -= 1\\n if level == 0:\\n return False\\n if ch == \\\"(\\\":\\n level += 1\\n return True\\n\\n is_single_method_on_self_t = check_parens_nest_level_gt_zero(\\n between_parens\\n )\\n directly_do_inplace = (\\n is_single_method_on_self_t and op_name == info.name\\n )\\n\\n if directly_do_inplace:\\n assert op_name is not None\\n assert between_parens is not None\\n formula = f\\\"self_t_raw.defined() ? self_t_raw.{op_name}_({between_parens}) : {formula}\\\"\\n else:\\n # Make sure that the forward grad is modified inplace when the original formula\\n # is out of place\\n formula = f\\\"self_t_raw.defined() ? self_t_raw.copy_({formula}) : {formula}\\\"\\n\\n required_original_self_value = bool(\\n re.search(IDENT_REGEX.format(\\\"original_self_p\\\"), formula)\\n ) or bool(re.search(IDENT_REGEX.format(\\\"original_self_t\\\"), formula))\\n\\n forward_derivatives = [\\n ForwardDerivative(\\n formula=formula,\\n var_names=(\\\"self\\\",),\\n var_types=fw_info.var_types,\\n required_inputs_fw_grad=fw_info.required_inputs_fw_grad,\\n required_inputs_primal=required_primals,\\n required_original_self_value=required_original_self_value,\\n is_reusing_outplace_formula=not is_exact_match,\\n ),\\n ]\\n\\n fw_derivative_dict[key] = forward_derivatives\\n\\n result.append(\\n NativeFunctionWithDifferentiabilityInfo(\\n func=f, info=info_dict, fw_derivatives=fw_derivative_dict\\n )\\n )\\n\\n return result\\n\\n\\ndef is_differentiable(\\n name: str, type: Type, info: DifferentiabilityInfo | None\\n) -> bool:\\n return type.is_tensor_like() and (\\n info is None or name not in info.non_differentiable_arg_names\\n )\\n\\n\\ndef gen_differentiable_outputs(\\n fn: NativeFunctionWithDifferentiabilityInfo, key: str = \\\"Default\\\"\\n) -> list[DifferentiableOutput]:\\n f = fn.func\\n info = fn.info[key] if fn.info else None\\n outputs: list[DifferentiableOutput] = [\\n DifferentiableOutput(\\n name=name,\\n type=ret.type,\\n cpp_type=cpp.return_type(ret, symint=True).cpp_type(),\\n )\\n for name, ret in zip(cpp.return_names(f), f.func.returns)\\n ]\\n output_differentiability = info.output_differentiability if info else None\\n if output_differentiability is not None:\\n if len(output_differentiability) != len(outputs):\\n raise RuntimeError(\\n f\\\"The length of output_differentiability ({len(output_differentiability)}), \\\"\\n f\\\"does not match the number of outputs ({len(outputs)}).\\\"\\n )\\n differentiable_outputs: list[DifferentiableOutput] = []\\n if False in output_differentiability and f.func.kind() == SchemaKind.inplace:\\n raise RuntimeError(\\n \\\"output_differentiability=False for inplace operation (version_counter won't get updated)\\\"\\n )\\n for differentiable, output in zip(output_differentiability, outputs):\\n if differentiable:\\n differentiable_outputs.append(output)\\n return differentiable_outputs\\n candidate_differentiable_outputs = list(\\n filter(lambda r: is_differentiable(r.name, r.type, info), outputs)\\n )\\n if uses_single_grad(info):\\n return candidate_differentiable_outputs[:1]\\n else:\\n return candidate_differentiable_outputs\\n\\n\\nfrom __future__ import annotations\\n\\nfrom typing import Any\\n\\nfrom torchgen.api.types import (\\n BaseCppType,\\n BaseCType,\\n boolT,\\n CType,\\n deviceT,\\n doubleT,\\n generatorT,\\n layoutT,\\n ListCType,\\n longT,\\n memoryFormatT,\\n NamedCType,\\n OptionalCType,\\n scalarT,\\n scalarTypeT,\\n stringT,\\n SymIntT,\\n VectorCType,\\n)\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n ListType,\\n OperatorName,\\n OptionalType,\\n Return,\\n TensorOptionsArguments,\\n Type,\\n)\\n\\n\\n_valueT: BaseCppType | None = None\\n\\n\\n# A ValueT is an IR type which represents the computation of a Tensor. In other\\n# words, a PyTorch user will do operations on lazy tensors, and each output lazy\\n# tensor internally tracks a ValueT representing the IR node that would have\\n# actually produced the value of this tensor for real.\\n#\\n# This is configurable because different lazy tensor backends (LTC vs XLA) will\\n# have different IR representations. (Though, arguably, after unification they\\n# shouldn't!)\\ndef getValueT() -> BaseCppType:\\n global _valueT\\n if not _valueT:\\n raise NotImplementedError(\\n \\\"The value type needs to be set with setValueT() in run_gen_lazy_tensor()\\\"\\n )\\n\\n return _valueT\\n\\n\\ndef setValueT(val: BaseCppType) -> None:\\n global _valueT\\n _valueT = val\\n\\n\\n# this is a bad hack. I need to refactor the data model to represent each arg in the schema as an object,\\n# making it easier to represent special properties of an arg.\\ntensorListValueT = BaseCppType(\\\"torch::lazy\\\", \\\"Value\\\")\\n\\n\\ndef process_ir_type(\\n typ: Type, properties: LazyIrProperties, *, symint: bool\\n) -> BaseCType | VectorCType | OptionalCType | ListCType:\\n \\\"\\\"\\\"\\n This function takes a type from NativeFunctions and converts it for use with\\n lazy tensor codegen.\\n\\n Type conversion for lazy currently consists of\\n (1) changing at::Tensors into lazy::Values\\n (2) wrapping everything in a BaseCType\\n (3) making cpp-reference types into cpp-value types (e.g. vector instead of IntArrayRef)\\n\\n (1) converts at::Tensors to lazy::Values (which wrap lazy::Nodes, with which Lazy IR represents tensors.)\\n There is special handling for Optional[Tensor] or List[Tensor], etc- hence 'tensor-like'\\n\\n This is incomplete- there are assertions in places that it's expected to need to add\\n more types as the codegen is used with more operators.\\n \\\"\\\"\\\"\\n if isinstance(typ, BaseType):\\n if typ.name == BaseTy.Tensor:\\n return BaseCType(getValueT())\\n elif typ.name == BaseTy.Scalar:\\n if properties.TreatScalarsAsConstants:\\n return BaseCType(scalarT)\\n # at::scalar has special handling,\\n # and is wrapped in an lazy::Value just like at::tensor\\n return BaseCType(getValueT())\\n elif typ.name == BaseTy.ScalarType:\\n return BaseCType(scalarTypeT)\\n elif typ.name == BaseTy.int:\\n return BaseCType(longT)\\n elif typ.name == BaseTy.SymInt:\\n if symint:\\n return BaseCType(getValueT())\\n else:\\n return BaseCType(longT)\\n elif typ.name == BaseTy.bool:\\n return BaseCType(boolT)\\n elif typ.name == BaseTy.float:\\n return BaseCType(doubleT)\\n elif typ.name == BaseTy.str:\\n return BaseCType(stringT)\\n elif typ.name == BaseTy.Device:\\n return BaseCType(deviceT)\\n elif typ.name == BaseTy.Generator:\\n return BaseCType(generatorT)\\n elif typ.name == BaseTy.Layout:\\n return BaseCType(layoutT)\\n elif typ.name == BaseTy.MemoryFormat:\\n return BaseCType(memoryFormatT)\\n else:\\n raise AssertionError(f\\\"TODO add support for type {repr(typ)}\\\")\\n elif isinstance(typ, OptionalType):\\n return OptionalCType(process_ir_type(typ.elem, properties, symint=symint))\\n elif isinstance(typ, ListType):\\n if str(typ.elem) == \\\"Tensor?\\\":\\n # TODO(whc) is this actually correct? or should it use a Vector like above\\n return ListCType(OptionalCType(BaseCType(getValueT())))\\n elif str(typ.elem) == \\\"Tensor\\\":\\n # this is a TensorList which comes in from GetTensorList as a Value\\n return BaseCType(tensorListValueT)\\n elif typ.elem == BaseType(BaseTy.SymInt):\\n # TODO: return a value type. The problem here is analogous to\\n # the problem with tensorListValueT: if you have SymInt[] you\\n # cannot conveniently save the list of Value directly, as nodes\\n # expect to save values as a vector for ALL arguments. So you\\n # need a separate IR node that represents all of the size nodes\\n # assembled into a list. I'm not an LTC dev so I don't want to\\n # figure it out right now. Y'all figure it out...\\n return VectorCType(BaseCType(longT))\\n\\n else:\\n return VectorCType(process_ir_type(typ.elem, properties, symint=symint))\\n else:\\n raise AssertionError(f\\\"unrecognized type {repr(typ)}\\\")\\n\\n\\n# TODO: Determining this based off of CType is bad; this should be computed\\n# from Type directly; then the same logic as process_ir_type can be used\\n#\\n# Invariant: passed typ should be an *owning* CType (e.g., we will report\\n# that ArrayRef is NOT a value type)\\ndef isValueType(typ: CType, properties: LazyIrProperties | None = None) -> bool:\\n \\\"\\\"\\\"\\n Given a type, determine if it is a Value-like type. This is equivalent to\\n being Tensor-like, but assumes the type has already been transformed.\\n \\\"\\\"\\\"\\n if isinstance(typ, BaseCType):\\n # I am regretting my naming conventions, but now we are wrapping at::scalar in\\n # lazy value, while preserving other 'scalar' types as scalars in the IR\\n treat_scalars_as_constants = properties and properties.TreatScalarsAsConstants\\n return (\\n typ.type == getValueT()\\n or (typ.type == scalarT and not treat_scalars_as_constants)\\n or typ.type == SymIntT\\n )\\n elif typ == VectorCType(BaseCType(SymIntT)):\\n # TODO: report True for this\\n return False\\n elif isinstance(typ, (OptionalCType, ListCType, VectorCType)):\\n return isValueType(typ.elem, properties)\\n return False\\n\\n\\ndef isSymIntType(typ: Type) -> bool:\\n return isinstance(typ, BaseType) and typ.name == BaseTy.SymInt\\n\\n\\ndef isWrappedScalarType(typ: Type) -> bool:\\n \\\"\\\"\\\"\\n Given a type, determine if it is a c10::scalar which we will wrap in a lazy Value.\\n Since we literally change the type from scalarT to valueT, information is lost.\\n This function helps build a list of wrapped scalars to save that information\\n \\\"\\\"\\\"\\n if isinstance(typ, BaseType):\\n # I am regretting my naming conventions, but now we are wrapping at::scalar in\\n # lazy value, while preserving other 'scalar' types as scalars in the IR\\n return typ.name == BaseTy.Scalar\\n elif isinstance(typ, (OptionalType, ListType)):\\n return isWrappedScalarType(typ.elem)\\n return False\\n\\n\\n# TODO: dedupe with Type.is_generator_like\\ndef isGeneratorType(typ: Type) -> bool:\\n if isinstance(typ, BaseType):\\n return typ.name == BaseTy.Generator\\n elif isinstance(typ, (OptionalType)):\\n return isGeneratorType(typ.elem)\\n return False\\n\\n\\n# This class caches a few derived properties computed from an Argument\\n# and LazyIrProperties\\nclass LazyArgument:\\n name: str\\n orig_type: Type\\n lazy_type_: CType | None\\n is_wrapped_scalar: bool\\n is_generator: bool\\n # TODO: this is lies, it is false for symint list\\n is_symint_or_list: bool\\n\\n # Whether or not we are treating this as symint or not\\n symint: bool\\n\\n # true if this argument is or contains a lazy IR value\\n is_lazy_value: bool\\n\\n def __init__(\\n self, arg: Argument, properties: LazyIrProperties, *, symint: bool\\n ) -> None:\\n self.name = arg.name\\n self.orig_type = arg.type\\n self.symint = symint\\n self.is_optional = isinstance(arg.type, OptionalType)\\n self.is_generator = isGeneratorType(arg.type)\\n self.lazy_type_ = process_ir_type(arg.type, properties, symint=symint)\\n self.is_wrapped_scalar = isWrappedScalarType(arg.type)\\n self.is_symint_or_list = symint and (\\n isSymIntType(arg.type)\\n or (isinstance(arg.type, OptionalType) and isSymIntType(arg.type.elem))\\n # TODO: lists of symints are not currently treated as value types\\n # or (isinstance(arg.type, ListType) and isSymIntType(arg.type.elem))\\n )\\n\\n self.is_lazy_value = isValueType(self.lazy_type, properties)\\n\\n @property\\n def lazy_type(self) -> CType:\\n assert (\\n self.lazy_type_ is not None\\n ), f\\\"Attempted to access lazy_type for invalid argument {self.name}\\\"\\n return self.lazy_type_\\n\\n\\nclass LazyIrProperties:\\n \\\"\\\"\\\"Collection of properties for an IR node\\n\\n The property groups are listed below. Each group is mutually\\n exclusive, meaning that only one property from each group can be True\\n at any one time. The properties can be accessed as if they were normal\\n attributes. The mutual exclusivity is automatically handled.\\n \\\"\\\"\\\"\\n\\n Properties: tuple[tuple[str, ...], ...] = (\\n (\\n \\\"ShapePrecompute\\\", # Assume shape has been precomputed\\n \\\"ShapeCompute\\\", # Need to compute the shape on construction\\n \\\"ShapeCache\\\", # Utilize the shape cache to defer computation\\n ),\\n (\\n \\\"Lower\\\", # Codegen full lower function\\n \\\"LowerDeclOnly\\\", # Codegen only lower function declaration\\n ),\\n (\\n \\\"CanBeReused\\\", # Codegen full reuse function\\n \\\"CanBeReusedDeclOnly\\\", # Codegen only reuse function declaration\\n ),\\n (\\n \\\"CreateFn\\\", # Codegen full create function\\n \\\"CreateFnDeclOnly\\\", # Codegen only create function declaration\\n ),\\n (\\n \\\"TreatScalarsAsConstants\\\", # Treat Scalars as constants instead of handling like values\\n ),\\n )\\n\\n def __init__(self, *default_properties: str) -> None:\\n properties: dict[tuple[str, ...], str | None] = dict.fromkeys(\\n LazyIrProperties.Properties\\n )\\n self.__dict__[\\\"properties\\\"] = properties\\n for p in default_properties:\\n setattr(self, p, True)\\n\\n def __getattr__(self, key: str) -> Any:\\n properties = self.__dict__[\\\"properties\\\"]\\n for values in LazyIrProperties.Properties:\\n if key in values:\\n return properties[values] == key\\n\\n return self.__getattribute__(key)\\n\\n def __setattr__(self, key: str, value: Any) -> Any:\\n properties = self.__dict__[\\\"properties\\\"]\\n for values in LazyIrProperties.Properties:\\n if key in values:\\n properties[values] = key if value else None\\n return value\\n\\n raise KeyError(f\\\"Invalid property: {key}\\\")\\n\\n\\n# Inspired by a FunctionSchema object, a LazyIrSchema holds the schema of a Lazy IR node.\\n# Unlike a FunctionSchema, it has no round-trippable string form (relating to the YAML),\\n# but carries type information from a native FunctionSchema modified for use with IR nodes,\\n# and preserving original argument names.\\n#\\n# TODO: This is not idiomatic with how other torchgen APIs transform on schema.\\nclass LazyIrSchema:\\n # The name of the operator this function schema describes.\\n name: OperatorName\\n\\n positional_args: tuple[LazyArgument, ...]\\n keyword_args: tuple[LazyArgument, ...]\\n\\n # TODO: Need to handle collisions with argument names at some point\\n returns: tuple[Return, ...]\\n\\n # if this schema has a Generator arg, list its orig ctype/name but don't\\n # build a LazyArgument since lazy IR doesn't support it\\n generator_arg: NamedCType | None = None\\n\\n # original function schema\\n func: FunctionSchema\\n\\n # Whether or not we are code-genning for SymInt or not\\n symint: bool\\n\\n properties: LazyIrProperties = LazyIrProperties(\\n # default properties\\n \\\"ShapePrecompute\\\",\\n \\\"Lower\\\",\\n \\\"CanBeReused\\\",\\n )\\n opkind: str | None = None\\n\\n def __init__(\\n self,\\n func: FunctionSchema,\\n properties: LazyIrProperties | None = None,\\n *,\\n symint: bool,\\n ) -> None:\\n if properties:\\n self.properties = properties\\n\\n self.func = func\\n self.symint = symint\\n positional_args: list[LazyArgument] = []\\n for arg_field in [\\\"pre_self_positional\\\", \\\"self_arg\\\", \\\"post_self_positional\\\"]:\\n if arg_field == \\\"self_arg\\\" and func.arguments.self_arg is not None:\\n arg = func.arguments.self_arg.argument\\n positional_args.append(\\n LazyArgument(arg, self.properties, symint=symint)\\n )\\n elif getattr(func.arguments, arg_field) is not None:\\n positional_args.extend(\\n LazyArgument(arg, self.properties, symint=symint)\\n for arg in getattr(func.arguments, arg_field)\\n )\\n self.positional_args = tuple(positional_args)\\n\\n keyword_args: list[LazyArgument] = []\\n for arg_field in [\\n \\\"pre_tensor_options_kwarg_only\\\",\\n \\\"tensor_options\\\",\\n \\\"post_tensor_options_kwarg_only\\\",\\n \\\"out\\\",\\n ]:\\n curr_args = getattr(func.arguments, arg_field)\\n if curr_args is not None:\\n if isinstance(curr_args, TensorOptionsArguments):\\n curr_args = curr_args.all()\\n for arg in curr_args:\\n if isGeneratorType(arg.type):\\n assert (\\n self.generator_arg is None\\n ), \\\"We expect there is only one generator arg\\\"\\n self.generator_arg = NamedCType(\\n arg.name, arg.type # type:ignore[arg-type]\\n )\\n keyword_args.extend(\\n LazyArgument(arg, self.properties, symint=symint)\\n for arg in curr_args\\n )\\n self.keyword_args = tuple(keyword_args)\\n self.name = func.name\\n self.returns = func.returns\\n\\n @property\\n def node_name(self) -> str:\\n \\\"\\\"\\\"\\n Return camel-case version of op in node.\\n\\n Note: This function also appends any `overload_name` in the operation.\\n For example, if the op is `bitwise_and.Tensor`, the returned name\\n will be `BitwiseAndTensor`.\\n \\\"\\\"\\\"\\n op_name = f\\\"{self.name.name}_{self.name.overload_name}\\\".lower()\\n return \\\"\\\".join(word.capitalize() or \\\"\\\" for word in op_name.split(\\\"_\\\"))\\n\\n @property\\n def aten_name(self) -> str:\\n return str(self.name.name)\\n\\n @property\\n def base_name(self) -> str:\\n return f\\\"{self.name.name.base}\\\"\\n\\n def filtered_args(\\n self,\\n positional: bool = True,\\n keyword: bool = True,\\n values: bool = True,\\n scalars: bool = True,\\n generator: bool = True,\\n ) -> list[LazyArgument]:\\n # This function maintains the sorted order of arguments but provides different filtered views.\\n # Some parts of the code care about kwargs vs args (TS lowerings),\\n # other parts care about whether they need to wrap the arg in a lazy value or leave it alone.\\n # Generators are special cased, as they are needed for fallback/shape-inference but not supported\\n # in TS lowerings and therefore also omitted from lazy IR.\\n args: list[LazyArgument] = []\\n if positional:\\n args.extend(self.positional_args)\\n if keyword:\\n args.extend(self.keyword_args)\\n\\n if values and scalars and generator:\\n return args\\n elif values and scalars:\\n return [a for a in args if not a.is_generator]\\n elif values:\\n return [a for a in args if a.is_lazy_value]\\n elif scalars:\\n return [\\n a\\n for a in args\\n if not a.is_lazy_value and (generator or not a.is_generator)\\n ]\\n\\n return []\\n\\n @property\\n def positional_values(self) -> list[LazyArgument]:\\n return self.filtered_args(\\n positional=True, keyword=False, values=True, scalars=False\\n )\\n\\n @property\\n def positional_scalars(self) -> list[LazyArgument]:\\n return self.filtered_args(\\n positional=True, keyword=False, values=False, scalars=True\\n )\\n\\n @property\\n def keyword_values(self) -> list[LazyArgument]:\\n return self.filtered_args(\\n positional=False, keyword=True, values=True, scalars=False\\n )\\n\\n @property\\n def keyword_scalars(self) -> list[LazyArgument]:\\n return self.filtered_args(\\n positional=False, keyword=True, values=False, scalars=True\\n )\\n\\n\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\nfrom typing import Iterator, Sequence, TYPE_CHECKING\\n\\nfrom torchgen.api.types.types_base import Binding, CType, Expr\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.model import (\\n BackendIndex,\\n FunctionSchema,\\n NativeFunction,\\n NativeFunctionsGroup,\\n NativeFunctionsViewGroup,\\n )\\n\\n\\n@dataclass(frozen=True)\\nclass CppSignature:\\n \\\"\\\"\\\"\\n A CppSignature represents a single overload in the C++ API. For\\n any given function schema, there may be multiple CppSignatures\\n corresponding to it, based on how we desugar to C++. See also\\n CppSignatureGroup.\\n \\\"\\\"\\\"\\n\\n # The schema this signature is derived from\\n func: FunctionSchema\\n\\n # Is this a C++ signature for a method, i.e. Tensor::my_op(...)?\\n method: bool\\n\\n # Is this a faithful C++ signature (i.e. following the JIT schema) or a convenience API\\n # (i.e. with a potential TensorOptions argument and out arguments in the front)\\n faithful: bool\\n\\n # Is this a symint C++ signature. For BC reasons, functions that take\\n # SymInts still present as int64_t in C++, and the SymInt variant is\\n # offered at a different overload name\\n #\\n # NB: If a function RETURNS a SymInt, this is ALWAYS false\\n symint: bool\\n\\n # The set of C++ arguments which should not have defaults applied to them\\n cpp_no_default_args: set[str]\\n\\n # Is this a fallback C++ binding? Fallback bindings are enabled by\\n # manual_cpp_binding: True and are alternate, non-public API that\\n # lets manual C++ binding implementors access the binding that would\\n # have been automatically generated\\n fallback_binding: bool = False\\n\\n # Return the unpacked argument structure of this signature,\\n # discarding information about which arguments are semantically\\n # related to each other.\\n def arguments(self) -> Sequence[Binding]:\\n return cpp.arguments(\\n self.func.arguments,\\n faithful=self.faithful,\\n symint=self.symint,\\n method=self.method,\\n cpp_no_default_args=self.cpp_no_default_args,\\n )\\n\\n def name(self, *, suppress_symint_suffix: bool = False) -> str:\\n n = cpp.name(\\n self.func,\\n faithful_name_for_out_overloads=self.faithful,\\n symint_overload=False if suppress_symint_suffix else self.symint,\\n )\\n if self.fallback_binding:\\n n = f\\\"__dispatch_{n}\\\"\\n return n\\n\\n # Render the C++ declaration for this signature\\n def decl(\\n self,\\n *,\\n name: str | None = None,\\n prefix: str = \\\"\\\",\\n is_redispatching_fn: bool = False,\\n suppress_symint_suffix: bool = False,\\n ) -> str:\\n returns_type = cpp.returns_type(\\n self.func.returns, symint=self.symint\\n ).cpp_type()\\n cpp_args = [a.decl() for a in self.arguments()]\\n if is_redispatching_fn:\\n cpp_args = [\\\"c10::DispatchKeySet dispatchKeySet\\\"] + cpp_args\\n cpp_args_str = \\\", \\\".join(cpp_args)\\n if name is None:\\n name = prefix + self.name(suppress_symint_suffix=suppress_symint_suffix)\\n return f\\\"{returns_type} {name}({cpp_args_str})\\\"\\n\\n # Render the C++ definition for this signature, not including\\n # the body (with curly braces)\\n def defn(\\n self,\\n *,\\n name: str | None = None,\\n prefix: str = \\\"\\\",\\n is_redispatching_fn: bool = False,\\n ) -> str:\\n returns_type = cpp.returns_type(\\n self.func.returns, symint=self.symint\\n ).cpp_type()\\n cpp_args = [a.defn() for a in self.arguments()]\\n if is_redispatching_fn:\\n cpp_args = [\\\"c10::DispatchKeySet dispatchKeySet\\\"] + cpp_args\\n cpp_args_str = \\\", \\\".join(cpp_args)\\n if name is None:\\n name = prefix + self.name()\\n return f\\\"{returns_type} {name}({cpp_args_str})\\\"\\n\\n def ptr_type(self) -> str:\\n args_types_str = \\\", \\\".join(a.type for a in self.arguments())\\n return f\\\"{cpp.returns_type(self.func.returns, symint=self.symint).cpp_type()} (*)({args_types_str})\\\"\\n\\n # Return the C++ function type, e.g., something like int(bool)\\n def type(self) -> str:\\n args_types_str = \\\", \\\".join(a.type for a in self.arguments())\\n return f\\\"{cpp.returns_type(self.func.returns, symint=self.symint).cpp_type()} ({args_types_str})\\\"\\n\\n\\n# Represents group of all CppSignatures associated with a\\n# FunctionSchema. Right now, that's the regular, user-visible\\n# signature, as well as a \\\"faithful\\\" signature which doesn't\\n# have grouping.\\n@dataclass(frozen=True)\\nclass CppSignatureGroup:\\n func: FunctionSchema\\n signature: CppSignature\\n faithful_signature: CppSignature | None\\n symint_signature: CppSignature | None\\n symint_faithful_signature: CppSignature | None\\n\\n def most_faithful_signature(self) -> CppSignature:\\n if self.faithful_signature:\\n return self.faithful_signature\\n else:\\n return self.signature\\n\\n def signatures(self, *, symint: bool = True) -> Iterator[CppSignature]:\\n yield self.signature\\n if self.faithful_signature:\\n yield self.faithful_signature\\n if symint:\\n if self.symint_signature:\\n yield self.symint_signature\\n if self.symint_faithful_signature:\\n yield self.symint_faithful_signature\\n\\n @staticmethod\\n def from_native_function(\\n f: NativeFunction, *, method: bool, fallback_binding: bool = False\\n ) -> CppSignatureGroup:\\n func = f.func\\n\\n def make_sig(*, faithful: bool, symint: bool) -> CppSignature:\\n return CppSignature(\\n func=func,\\n faithful=faithful,\\n symint=symint,\\n method=method,\\n fallback_binding=fallback_binding,\\n cpp_no_default_args=f.cpp_no_default_args,\\n )\\n\\n def make_sigs(*, symint: bool) -> tuple[CppSignature, CppSignature | None]:\\n faithful_signature: CppSignature | None = None\\n if func.arguments.tensor_options is not None or len(func.arguments.out) > 0:\\n faithful_signature = make_sig(faithful=True, symint=symint)\\n signature = make_sig(faithful=False, symint=symint)\\n return signature, faithful_signature\\n\\n signature, faithful_signature = make_sigs(symint=False)\\n symint_signature: CppSignature | None = None\\n symint_faithful_signature: CppSignature | None = None\\n if func.has_symint():\\n symint_signature, symint_faithful_signature = make_sigs(symint=True)\\n\\n return CppSignatureGroup(\\n func=func,\\n signature=signature,\\n faithful_signature=faithful_signature,\\n symint_signature=symint_signature,\\n symint_faithful_signature=symint_faithful_signature,\\n )\\n\\n\\n@dataclass(frozen=True)\\nclass DispatcherSignature:\\n # The schema this signature is derived from\\n func: FunctionSchema\\n\\n # Allows you to prepend an arbitrary prefix to the signature name.\\n # This is useful for parts of the codegen that generate wrappers around kernels,\\n # and need to avoid naming collisions.\\n prefix: str = \\\"\\\"\\n\\n symint: bool = True\\n\\n def arguments(self) -> list[Binding]:\\n return dispatcher.arguments(self.func, symint=self.symint)\\n\\n def name(self) -> str:\\n return self.prefix + dispatcher.name(self.func)\\n\\n def decl(self, name: str | None = None) -> str:\\n args_str = \\\", \\\".join(a.decl() for a in self.arguments())\\n if name is None:\\n name = self.name()\\n return f\\\"{self.returns_type().cpp_type()} {name}({args_str})\\\"\\n\\n def defn(\\n self, name: str | None = None, *, is_redispatching_fn: bool = False\\n ) -> str:\\n args = [a.defn() for a in self.arguments()]\\n if is_redispatching_fn:\\n args = [\\\"c10::DispatchKeySet dispatchKeySet\\\"] + args\\n args_str = \\\", \\\".join(args)\\n if name is None:\\n name = self.name()\\n return f\\\"{self.returns_type().cpp_type()} {name}({args_str})\\\"\\n\\n def exprs(self) -> list[Expr]:\\n return [Expr(a.name, a.nctype) for a in self.arguments()]\\n\\n def returns_type(self) -> CType:\\n return dispatcher.returns_type(self.func.returns, symint=self.symint)\\n\\n def ptr_type(self) -> str:\\n dispatcher_args_types_str = \\\", \\\".join(a.type for a in self.arguments())\\n return f\\\"{self.returns_type().cpp_type()} (*)({dispatcher_args_types_str})\\\"\\n\\n # Return the C++ function type, e.g., something like int(bool)\\n def type(self) -> str:\\n dispatcher_args_types_str = \\\", \\\".join(a.type for a in self.arguments())\\n return f\\\"{self.returns_type().cpp_type()} ({dispatcher_args_types_str})\\\"\\n\\n @staticmethod\\n def from_schema(\\n func: FunctionSchema, *, prefix: str = \\\"\\\", symint: bool = True\\n ) -> DispatcherSignature:\\n return DispatcherSignature(func, prefix, symint)\\n\\n\\n@dataclass(frozen=True)\\nclass NativeSignature:\\n # The schema this signature is derived from\\n func: FunctionSchema\\n\\n symint: bool\\n\\n prefix: str = \\\"\\\"\\n\\n def name(self) -> str:\\n return self.prefix + native.name(self.func)\\n\\n def decl(self, name: str | None = None) -> str:\\n args_str = \\\", \\\".join(a.decl() for a in self.arguments())\\n if name is None:\\n name = self.name()\\n return f\\\"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} {name}({args_str})\\\"\\n\\n def defn(self, name: str | None = None) -> str:\\n args_str = \\\", \\\".join(a.defn() for a in self.arguments())\\n if name is None:\\n name = self.name()\\n return f\\\"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} {name}({args_str})\\\"\\n\\n def ptr_type(self) -> str:\\n # don't include defaults in type signature!\\n args_str = \\\", \\\".join(a.defn() for a in self.arguments())\\n return f\\\"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} (*)({args_str})\\\"\\n\\n def arguments(self) -> list[Binding]:\\n return native.arguments(self.func, symint=self.symint)\\n\\n def returns_type(self) -> CType:\\n return native.returns_type(self.func.returns, symint=self.symint)\\n\\n def dispatcher_exprs(self) -> list[Expr]:\\n return translate.translate(\\n self.arguments(), dispatcher.arguments(self.func), method=False\\n )\\n\\n\\n@dataclass(frozen=True)\\nclass ViewInverseSignature:\\n g: NativeFunctionsViewGroup\\n\\n def name(self) -> str:\\n return functionalization.reverse_name(self.g.view, include_namespace=False)\\n\\n def decl(self) -> str:\\n return_type = functionalization.returns_type(self.g.view.func)\\n decls = [\\n a.decl()\\n for a in functionalization.inner_arguments(\\n self.g.view.func, is_reverse=True\\n )\\n ]\\n return f\\\"static {return_type.cpp_type()} {self.name()}({', '.join(decls)});\\\"\\n\\n\\n@dataclass(frozen=True)\\nclass FunctionalizationLambda:\\n g: NativeFunctionsViewGroup\\n\\n # are we generating the forward lambda or the reverse lambda?\\n is_reverse: bool\\n\\n def captures(self) -> list[Expr]:\\n # The lambda lives inside of a kernel following the dispatcher API, so its outer context is the dispatcher arguments\\n # We also need to read the \\\"reapply views\\\" TLS at the time that the functionalization kernel was executed,\\n # and plumb it into the lambda.\\n outer_ctx = dispatcher.arguments(self.g.view.func) + [\\n functionalization.reapply_views_binding,\\n functionalization.inverse_return_mode_binding,\\n ]\\n capture_bindings = functionalization.capture_arguments(\\n self.g.view.func, is_reverse=self.is_reverse\\n )\\n # allow_expensive_conversions is set because we want to convert\\n # some reference types (IntArrayRef) to value types (vector).\\n capture_exprs = translate.translate(\\n outer_ctx, capture_bindings, method=False, allow_expensive_conversions=True\\n )\\n return capture_exprs\\n\\n def decl(self) -> str:\\n return_type = functionalization.returns_type(self.g.view.func)\\n capture_str = \\\", \\\".join(\\n f\\\"{val.type.name} = {val.expr}\\\" for val in self.captures()\\n )\\n decls = [\\n a.decl()\\n for a in functionalization.outer_arguments(is_reverse=self.is_reverse)\\n ]\\n return f\\\"[{capture_str}]({', '.join(decls)}) -> {return_type.cpp_type()}\\\"\\n\\n def inner_call(self, *, reapply_views: bool | None = None) -> str:\\n inner_call_name = functionalization.name(\\n self.g,\\n is_reverse=self.is_reverse,\\n include_namespace=True,\\n reapply_views=reapply_views,\\n )\\n\\n arg_ctx = functionalization.outer_arguments(is_reverse=self.is_reverse)\\n capture_ctx = functionalization.capture_arguments(\\n self.g.view.func, is_reverse=self.is_reverse\\n )\\n full_ctx = arg_ctx + capture_ctx\\n\\n assert self.g.view_copy is not None\\n call_bindings = functionalization.inner_arguments(\\n self.g.view_copy.func, is_reverse=self.is_reverse\\n )\\n maybe_index = functionalization.inner_call_index(self.g.view_copy.func)\\n call_exprs = [\\n e.expr for e in translate.translate(full_ctx, call_bindings, method=False)\\n ]\\n if not self.is_reverse and maybe_index is not None:\\n return f'{inner_call_name}({\\\", \\\".join(call_exprs)})[{maybe_index.name}];'\\n else:\\n return f'{inner_call_name}({\\\", \\\".join(call_exprs)});'\\n\\n @staticmethod\\n def from_func(\\n g: NativeFunctionsViewGroup, *, is_reverse: bool\\n ) -> FunctionalizationLambda:\\n return FunctionalizationLambda(g, is_reverse)\\n\\n\\n@dataclass(frozen=True)\\nclass StructuredImplSignature:\\n g: NativeFunctionsGroup\\n name: str\\n\\n def defn(self, name: str | None = None) -> str:\\n args_str = \\\", \\\".join(a.defn() for a in self.arguments())\\n return f\\\"TORCH_IMPL_FUNC({self.name})({args_str})\\\"\\n\\n def arguments(self) -> list[Binding]:\\n return structured.impl_arguments(self.g)\\n\\n\\n# Helper functions\\n\\n\\ndef kernel_signature(\\n f: NativeFunction, backend_index: BackendIndex, *, prefix: str = \\\"\\\"\\n) -> NativeSignature | DispatcherSignature:\\n # Note [External Backends Follow Dispatcher API]\\n # Kernel signatures for in-tree backends follow the \\\"native\\\" API,\\n # while kernels for out-of-tree backends follow the dispatcher API.\\n # See the comments in `native.py` for details, but historically there have been\\n # some small differences in schema convention between them and the Dispatcher API.\\n # Any differences that require translating between the two will results in a runtime cost,\\n # so we'd like to keep the differences as small as possible.\\n # With external backends, we'd like to enforce that they write their kernels with schemas\\n # that match the Dispatcher API directly, if they can.\\n meta = backend_index.get_kernel(f)\\n symint = meta is not None and meta.supports_symint()\\n if symint:\\n assert (\\n f.func.has_symint()\\n ), f\\\"attempted to define symint kernel for {backend_index.dispatch_key} without SymInt in schema\\\"\\n if backend_index.external:\\n return DispatcherSignature.from_schema(f.func, prefix=prefix, symint=symint)\\n else:\\n return NativeSignature(f.func, prefix=prefix, symint=symint)\\n\\n\\n# Functions only, no types\\nfrom torchgen.api import (\\n cpp,\\n dispatcher,\\n functionalization,\\n native,\\n structured,\\n translate,\\n)\\n\\n\\n\\\"\\\"\\\"\\nWhere should I add a new type? `types_base.py` vs `types.py`\\n\\nThis file defines data model classes for torchgen typing system, as well as some base types such as int32_t.\\n\\n`types.py` defines ATen Tensor type and some c10 types, along with signatures that use these types.\\n\\nThe difference between these two files, is `types_base.py` should be implementation-agnostic, meaning it shouldn't\\ncontain any type definition that is tight to a specific C++ library (e.g., ATen), so that it can be easily reused\\nif we want to generate code for another C++ library.\\n\\nAdd new types to `types.py` if these types are ATen/c10 related.\\nAdd new types to `types_base.py` if they are basic and not attached to ATen/c10.\\n\\\"\\\"\\\"\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\n\\nfrom torchgen.api.types.types_base import (\\n BaseCppType,\\n BaseCType,\\n boolT,\\n byteT,\\n charT,\\n CType,\\n doubleT,\\n floatT,\\n int32T,\\n longT,\\n shortT,\\n)\\nfrom torchgen.model import BaseTy, ScalarType\\n\\n\\nTENSOR_LIST_LIKE_CTYPES = [\\n \\\"at::TensorList\\\",\\n \\\"const c10::List<::std::optional> &\\\",\\n \\\"const at::ITensorListRef &\\\",\\n]\\n\\n\\nhalfT = BaseCppType(\\\"at\\\", \\\"Half\\\")\\ncomplexHalfT = BaseCppType(\\n \\\"c10\\\", \\\"complex\\\"\\n) # stuffing template param here is an abuse\\ncomplexFloatT = BaseCppType(\\\"c10\\\", \\\"complex\\\")\\ncomplexDoubleT = BaseCppType(\\\"c10\\\", \\\"complex\\\")\\nbfloat16T = BaseCppType(\\\"at\\\", \\\"BFloat16\\\")\\nfloat8_e5m2T = BaseCppType(\\\"at\\\", \\\"Float8_e5m2\\\")\\nfloat8_e5m2fnuzT = BaseCppType(\\\"at\\\", \\\"Float8_e5m2fnuz\\\")\\nfloat8_e4m3fnT = BaseCppType(\\\"at\\\", \\\"Float8_e4m3fn\\\")\\nfloat8_e4m3fnuzT = BaseCppType(\\\"at\\\", \\\"Float8_e4m3fnuz\\\")\\nstringT = BaseCppType(\\\"c10\\\", \\\"string_view\\\")\\ngeneratorT = BaseCppType(\\\"at\\\", \\\"Generator\\\")\\nscalarTypeT = BaseCppType(\\\"at\\\", \\\"ScalarType\\\")\\ntensorT = BaseCppType(\\\"at\\\", \\\"Tensor\\\")\\noptionalTensorRefT = BaseCppType(\\\"at\\\", \\\"OptionalTensorRef\\\")\\ntensorListT = BaseCppType(\\\"at\\\", \\\"TensorList\\\")\\niTensorListRefT = BaseCppType(\\\"at\\\", \\\"ITensorListRef\\\")\\niOptTensorListRefT = BaseCppType(\\\"at\\\", \\\"IOptTensorListRef\\\")\\ndimnameT = BaseCppType(\\\"at\\\", \\\"Dimname\\\")\\ndimnameListT = BaseCppType(\\\"at\\\", \\\"DimnameList\\\")\\ndimVectorT = BaseCppType(\\\"at\\\", \\\"DimVector\\\")\\nlayoutT = BaseCppType(\\\"at\\\", \\\"Layout\\\")\\ndeviceT = BaseCppType(\\\"at\\\", \\\"Device\\\")\\ndeviceIndexT = BaseCppType(\\\"at\\\", \\\"DeviceIndex\\\")\\nscalarT = BaseCppType(\\\"at\\\", \\\"Scalar\\\")\\noptionalScalarRefT = BaseCppType(\\\"at\\\", \\\"OptionalScalarRef\\\")\\nmemoryFormatT = BaseCppType(\\\"at\\\", \\\"MemoryFormat\\\")\\nqschemeT = BaseCppType(\\\"at\\\", \\\"QScheme\\\")\\nstorageT = BaseCppType(\\\"at\\\", \\\"Storage\\\")\\nstreamT = BaseCppType(\\\"at\\\", \\\"Stream\\\")\\nintArrayRefT = BaseCppType(\\\"at\\\", \\\"IntArrayRef\\\")\\noptionalIntArrayRefT = BaseCppType(\\\"at\\\", \\\"OptionalIntArrayRef\\\")\\noptionalSymIntArrayRefT = BaseCppType(\\\"at\\\", \\\"OptionalSymIntArrayRef\\\")\\ntensorOptionsT = BaseCppType(\\\"at\\\", \\\"TensorOptions\\\")\\ntypeAndSizeT = BaseCppType(\\\"torch::autograd::generated\\\", \\\"TypeAndSize\\\")\\ntensorGeometryT = BaseCppType(\\\"at\\\", \\\"TensorGeometry\\\")\\nSymIntT = BaseCppType(\\\"c10\\\", \\\"SymInt\\\")\\nsymIntArrayRefT = BaseCppType(\\\"c10\\\", \\\"SymIntArrayRef\\\")\\n\\n# Types representing template parameters. Technically, we probably shouldn't\\n# represent them this way in codegen, but it was pretty convenient.\\nscalar_t = BaseCppType(\\\"\\\", \\\"scalar_t\\\")\\nopmath_t = BaseCppType(\\\"\\\", \\\"opmath_t\\\")\\n\\nScalarTypeToCppMapping: dict[ScalarType, BaseCppType] = {\\n ScalarType.Byte: byteT,\\n ScalarType.Char: charT,\\n ScalarType.Short: shortT,\\n ScalarType.Int: int32T,\\n ScalarType.Long: longT,\\n ScalarType.Half: halfT,\\n ScalarType.Float: floatT,\\n ScalarType.Double: doubleT,\\n ScalarType.ComplexHalf: complexHalfT,\\n ScalarType.ComplexFloat: complexFloatT,\\n ScalarType.ComplexDouble: complexDoubleT,\\n ScalarType.Bool: boolT,\\n ScalarType.Float8_e5m2: float8_e5m2T,\\n ScalarType.Float8_e5m2fnuz: float8_e5m2fnuzT,\\n ScalarType.Float8_e4m3fn: float8_e4m3fnT,\\n ScalarType.Float8_e4m3fnuz: float8_e4m3fnuzT,\\n}\\n\\nBaseTypeToCppMapping: dict[BaseTy, BaseCppType] = {\\n BaseTy.int: longT,\\n BaseTy.float: doubleT,\\n BaseTy.bool: boolT,\\n BaseTy.str: stringT,\\n BaseTy.Generator: generatorT,\\n BaseTy.ScalarType: scalarTypeT,\\n BaseTy.Tensor: tensorT,\\n BaseTy.Dimname: dimnameT,\\n BaseTy.DimVector: dimVectorT,\\n BaseTy.Layout: layoutT,\\n BaseTy.Device: deviceT,\\n BaseTy.DeviceIndex: deviceIndexT,\\n BaseTy.Scalar: scalarT,\\n BaseTy.MemoryFormat: memoryFormatT,\\n BaseTy.QScheme: qschemeT,\\n BaseTy.Storage: storageT,\\n BaseTy.Stream: streamT,\\n BaseTy.SymInt: SymIntT,\\n}\\n\\n# CTypes encode C++ type structure as needed for translation.\\n\\n\\n@dataclass(frozen=True)\\nclass OptionalCType(CType):\\n elem: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n # Do not pass `strip_ref` recursively.\\n return f\\\"::std::optional<{self.elem.cpp_type()}>\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"::std::optional<{self.elem.cpp_type_registration_declarations()}>\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return OptionalCType(self.elem.remove_const_ref())\\n\\n\\n@dataclass(frozen=True)\\nclass ListCType(CType):\\n elem: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n # Do not pass `strip_ref` recursively.\\n return f\\\"c10::List<{self.elem.cpp_type()}>\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"c10::List<{self.elem.cpp_type_registration_declarations()}>\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return ListCType(self.elem.remove_const_ref())\\n\\n\\n@dataclass(frozen=True)\\nclass ArrayRefCType(CType):\\n elem: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n # Do not pass `strip_ref` recursively.\\n return f\\\"at::ArrayRef<{self.elem.cpp_type()}>\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"ArrayRef<{self.elem.cpp_type_registration_declarations()}>\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return ArrayRefCType(self.elem.remove_const_ref())\\n\\n\\n@dataclass(frozen=True)\\nclass VectorizedCType(CType):\\n # This template is explicitly specialized, so the only valid\\n # elems are those we have specializations for (e.g., float, double, ...)\\n # scalar_t is also a common argument here (when we are codegen in\\n # a templated context)\\n elem: BaseCType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n return f\\\"at::vec::Vectorized<{self.elem.cpp_type()}>\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n raise NotImplementedError\\n\\n def remove_const_ref(self) -> CType:\\n return self\\n\\n\\n\\\"\\\"\\\"\\nWhere should I add a new type? `types_base.py` vs `types.py`\\n\\nThis file defines data model classes for torchgen typing system, as well as some base types such as int32_t.\\n\\n`types.py` defines ATen Tensor type and some c10 types, along with signatures that use these types.\\n\\nThe difference between these two files, is `types_base.py` should be implementation-agnostic, meaning it shouldn't\\ncontain any type definition that is tight to a specific C++ library (e.g., ATen), so that it can be easily reused\\nif we want to generate code for another C++ library.\\n\\nAdd new types to `types.py` if these types are ATen/c10 related.\\nAdd new types to `types_base.py` if they are basic and not attached to ATen/c10.\\n\\\"\\\"\\\"\\n\\nfrom __future__ import annotations\\n\\nfrom abc import ABC, abstractmethod\\nfrom dataclasses import dataclass\\nfrom enum import auto, Enum\\nfrom typing import TYPE_CHECKING, Union\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.model import Argument, SelfArgument, TensorOptionsArguments\\n\\n\\n# An ArgName is just the str name of the argument in schema;\\n# but in some special circumstances, we may add a little extra\\n# context. The Enum SpecialArgName covers all of these cases;\\n# grep for their construction sites to see when they can occur.\\n\\n\\nclass SpecialArgName(Enum):\\n possibly_redundant_memory_format = auto()\\n\\n\\nArgName = Union[str, SpecialArgName]\\n\\n\\n# This class shouldn't be created directly; instead, use/create one of the singletons below.\\n@dataclass(frozen=True)\\nclass BaseCppType:\\n ns: str | None\\n name: str\\n\\n def __str__(self) -> str:\\n if self.ns is None or self.ns == \\\"\\\":\\n return self.name\\n return f\\\"{self.ns}::{self.name}\\\"\\n\\n\\n# The set of all non-templated, valid, fully-qualified names of C++ types that are used in the codegen.\\n# Templated types get their own dataclass, mainly to make namespace parsing easier.\\nbyteT = BaseCppType(\\\"\\\", \\\"uint8_t\\\")\\ncharT = BaseCppType(\\\"\\\", \\\"int8_t\\\")\\nshortT = BaseCppType(\\\"\\\", \\\"int16_t\\\")\\n# It would be more symmetric for this to be called intT, but it easy to mix\\n# this up with JIT int (which is int64_t in C++), so we intentionally don't\\n# define intT to make it obvious when you've stuffed it up\\nint32T = BaseCppType(\\\"\\\", \\\"int32_t\\\")\\nlongT = BaseCppType(\\\"\\\", \\\"int64_t\\\")\\ndoubleT = BaseCppType(\\\"\\\", \\\"double\\\")\\nfloatT = BaseCppType(\\\"\\\", \\\"float\\\")\\nboolT = BaseCppType(\\\"\\\", \\\"bool\\\")\\nvoidT = BaseCppType(\\\"\\\", \\\"void\\\")\\n\\n\\nclass CType(ABC):\\n @abstractmethod\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n raise NotImplementedError\\n\\n @abstractmethod\\n def cpp_type_registration_declarations(self) -> str:\\n raise NotImplementedError\\n\\n @abstractmethod\\n def remove_const_ref(self) -> CType:\\n return self\\n\\n\\n@dataclass(frozen=True)\\nclass BaseCType(CType):\\n type: BaseCppType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n return str(self.type)\\n\\n # For BC reasons, we don't want to introduce at:: namespaces to RegistrationDeclarations.yaml\\n # TODO: Kill this when we eventually remove it!\\n def cpp_type_registration_declarations(self) -> str:\\n return str(self.type).replace(\\\"at::\\\", \\\"\\\")\\n\\n def remove_const_ref(self) -> CType:\\n return self\\n\\n\\n@dataclass(frozen=True)\\nclass ConstRefCType(CType):\\n elem: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n if strip_ref:\\n return self.elem.cpp_type(strip_ref=strip_ref)\\n return f\\\"const {self.elem.cpp_type()} &\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"const {self.elem.cpp_type_registration_declarations()} &\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return self.elem.remove_const_ref()\\n\\n\\n@dataclass(frozen=True)\\nclass VectorCType(CType):\\n elem: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n # Do not pass `strip_ref` recursively.\\n return f\\\"::std::vector<{self.elem.cpp_type()}>\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"::std::vector<{self.elem.cpp_type_registration_declarations()}>\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return VectorCType(self.elem.remove_const_ref())\\n\\n\\n@dataclass(frozen=True)\\nclass ArrayCType(CType):\\n elem: CType\\n size: int\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n # Do not pass `strip_ref` recursively.\\n return f\\\"::std::array<{self.elem.cpp_type()},{self.size}>\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"::std::array<{self.elem.cpp_type_registration_declarations()},{self.size}>\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return ArrayCType(self.elem.remove_const_ref(), self.size)\\n\\n\\n@dataclass(frozen=True)\\nclass TupleCType(CType):\\n elems: list[CType]\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n # Do not pass `strip_ref` recursively.\\n return f'::std::tuple<{\\\",\\\".join([e.cpp_type() for e in self.elems])}>'\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f'::std::tuple<{\\\",\\\".join([e.cpp_type_registration_declarations() for e in self.elems])}>'\\n\\n def remove_const_ref(self) -> CType:\\n return TupleCType([e.remove_const_ref() for e in self.elems])\\n\\n\\n@dataclass(frozen=True)\\nclass MutRefCType(CType):\\n elem: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n if strip_ref:\\n return self.elem.cpp_type(strip_ref=strip_ref)\\n return f\\\"{self.elem.cpp_type()} &\\\"\\n\\n def cpp_type_registration_declarations(self) -> str:\\n return f\\\"{self.elem.cpp_type_registration_declarations()} &\\\"\\n\\n def remove_const_ref(self) -> CType:\\n return self.elem.remove_const_ref()\\n\\n\\n# A NamedCType is short for Named C++ semantic type. A NamedCType represents a C++ type, plus\\n# semantic information about what it represents. For example, consider the\\n# argument \\\"bool pin_memory\\\"; its normal C++ type is \\\"bool\\\", but its C++\\n# semantic type also keeps track that this represents a \\\"pin_memory\\\"; you can't\\n# just use a random other boolean in a context where you need a \\\"pin_memory\\\"!\\n#\\n\\n\\n@dataclass(frozen=True)\\nclass NamedCType:\\n name: ArgName\\n type: CType\\n\\n def cpp_type(self, *, strip_ref: bool = False) -> str:\\n return self.type.cpp_type(strip_ref=strip_ref)\\n\\n # For BC reasons, we don't want to introduce at:: namespaces to RegistrationDeclarations.yaml\\n # TODO: Kill this when we eventually remove it!\\n def cpp_type_registration_declarations(self) -> str:\\n return self.type.cpp_type_registration_declarations()\\n\\n def remove_const_ref(self) -> NamedCType:\\n return NamedCType(self.name, self.type.remove_const_ref())\\n\\n def with_name(self, name: str) -> NamedCType:\\n return NamedCType(name, self.type)\\n\\n\\n# A binding represents any C++ binding site for a formal parameter.\\n# We don't distinguish between binding sites for different APIs;\\n# instead, all of the important distinctions are encoded in CType,\\n# which you can use to figure out if a given Binding is appropriate\\n# for use in another context. (See torchgen.api.translate)\\n\\n\\n@dataclass(frozen=True)\\nclass Binding:\\n name: str\\n nctype: NamedCType\\n argument: Argument | TensorOptionsArguments | SelfArgument\\n # TODO: maybe don't represent default here\\n default: str | None = None\\n\\n def rename(self, name: str) -> Binding:\\n return Binding(\\n name=name,\\n nctype=self.nctype,\\n argument=self.argument,\\n default=self.default,\\n )\\n\\n @property\\n def type(self) -> str:\\n return self.nctype.cpp_type()\\n\\n def no_default(self) -> Binding:\\n return Binding(\\n name=self.name,\\n nctype=self.nctype,\\n default=None,\\n argument=self.argument,\\n )\\n\\n def decl(self, *, func_ptr_cast: bool = False) -> str:\\n mb_default = \\\"\\\"\\n if self.default is not None:\\n mb_default = f\\\"={self.default}\\\"\\n\\n # casting only needs to know the type\\n if func_ptr_cast:\\n return f\\\"{self.type}\\\"\\n else:\\n return f\\\"{self.type} {self.name}{mb_default}\\\"\\n\\n # For BC reasons, we don't want to introduce at:: namespaces to RegistrationDeclarations.yaml\\n # TODO: Kill this when we eventually remove it!\\n def decl_registration_declarations(self) -> str:\\n type_s = self.nctype.cpp_type_registration_declarations()\\n mb_default = \\\"\\\"\\n if self.default is not None:\\n mb_default = f\\\"={self.default}\\\"\\n return f\\\"{type_s} {self.name}{mb_default}\\\"\\n\\n def defn(self) -> str:\\n return f\\\"{self.type} {self.name}\\\"\\n\\n def with_name(self, name: str) -> Binding:\\n return Binding(\\n name=name, nctype=self.nctype, argument=self.argument, default=self.default\\n )\\n\\n\\n# An Expr is a C++ expression. It has a C++ string representing its syntax,\\n# as well as a CType saying what it provides.\\n\\n\\n@dataclass(frozen=True)\\nclass Expr:\\n expr: str\\n type: NamedCType\\n\\n\\nfrom torchgen.api.types.types import *\\nfrom torchgen.api.types.types_base import *\\n\\n\\nfrom torchgen.api.types.signatures import * # usort: skip\\n\\n\\nfrom __future__ import annotations\\n\\nimport itertools\\nimport textwrap\\nfrom dataclasses import dataclass\\nfrom typing import Literal, TYPE_CHECKING\\n\\nimport torchgen.api.cpp as cpp\\nimport torchgen.api.meta as meta\\nimport torchgen.api.structured as structured\\nfrom torchgen.api.translate import translate\\nfrom torchgen.api.types import (\\n BaseCType,\\n Binding,\\n ConstRefCType,\\n CppSignature,\\n CppSignatureGroup,\\n DispatcherSignature,\\n Expr,\\n kernel_signature,\\n MutRefCType,\\n NamedCType,\\n NativeSignature,\\n tensorT,\\n)\\nfrom torchgen.context import method_with_native_function, native_function_manager\\nfrom torchgen.model import (\\n Argument,\\n BackendIndex,\\n DeviceCheckType,\\n DispatchKey,\\n gets_generated_out_inplace_wrapper,\\n is_cuda_dispatch_key,\\n NativeFunction,\\n NativeFunctionsGroup,\\n SchemaKind,\\n TensorOptionsArguments,\\n)\\nfrom torchgen.utils import assert_never, mapMaybe, Target\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.selective_build.selector import SelectiveBuilder\\n\\n\\ndef gen_registration_headers(\\n backend_index: BackendIndex,\\n per_operator_headers: bool,\\n rocm: bool,\\n) -> list[str]:\\n if per_operator_headers:\\n headers = [\\\"#include \\\"]\\n else:\\n headers = [\\\"#include \\\"]\\n\\n if backend_index.dispatch_key in (DispatchKey.CPU, DispatchKey.Meta):\\n headers.append(\\\"#include \\\")\\n elif backend_index.dispatch_key == DispatchKey.CUDA:\\n if rocm:\\n headers.append(\\\"#include \\\")\\n else:\\n headers.append(\\\"#include \\\")\\n elif backend_index.dispatch_key == DispatchKey.MPS:\\n headers.append(\\\"#include \\\")\\n elif backend_index.dispatch_key == DispatchKey.XPU:\\n # XPU specific, this header resides in third_party/torch-xpu-ops\\n headers.append(\\\"#include \\\")\\n elif per_operator_headers:\\n headers += [\\n \\\"#include \\\",\\n \\\"#include \\\",\\n \\\"#include \\\",\\n \\\"#include \\\",\\n ]\\n else:\\n headers.append(\\\"#include \\\")\\n\\n headers.append(\\\"#include \\\")\\n return headers\\n\\n\\ndef gen_empty_impl_names(\\n backend_index: BackendIndex,\\n) -> tuple[str | None, str | None]:\\n empty_impl = None\\n empty_strided_impl = None\\n\\n if backend_index.dispatch_key in (\\n DispatchKey.Meta,\\n DispatchKey.CPU,\\n DispatchKey.CUDA,\\n DispatchKey.MPS,\\n DispatchKey.XPU,\\n ):\\n dispatch = str(backend_index.dispatch_key).lower()\\n empty_impl = f\\\"at::detail::empty_{dispatch}\\\"\\n empty_strided_impl = f\\\"at::detail::empty_strided_{dispatch}\\\"\\n elif backend_index.dispatch_key in (\\n DispatchKey.CompositeExplicitAutogradNonFunctional,\\n DispatchKey.QuantizedCPU,\\n DispatchKey.QuantizedCUDA,\\n DispatchKey.XPU,\\n ):\\n empty_impl = \\\"at::empty\\\"\\n empty_strided_impl = \\\"at::empty_strided\\\"\\n\\n return empty_impl, empty_strided_impl\\n\\n\\ndef gen_create_out_helper(backend_index: BackendIndex) -> list[str]:\\n if backend_index.dispatch_key == DispatchKey.Meta:\\n empty_options = \\\"options.device(at::kMeta)\\\"\\n else:\\n empty_options = \\\"options\\\"\\n\\n empty_impl, empty_strided_impl = gen_empty_impl_names(backend_index)\\n if empty_impl is None:\\n return []\\n\\n return [\\n f\\\"\\\"\\\"\\nTensor create_out(IntArrayRef sizes, IntArrayRef strides, const TensorOptions &options) {{\\n if (strides.empty()) {{\\n return {empty_impl}(sizes, {empty_options});\\n }} else {{\\n return {empty_strided_impl}(sizes, strides, {empty_options});\\n }}\\n}}\\n\\\"\\\"\\\"\\n ]\\n\\n\\ndef gen_maybe_create_proxy_helper(backend_index: BackendIndex) -> list[str]:\\n _, empty_strided_impl = gen_empty_impl_names(backend_index)\\n return (\\n []\\n if empty_strided_impl is None\\n else [\\n f\\\"\\\"\\\"\\nstd::optional maybe_create_proxy(const Tensor &out, IntArrayRef sizes, IntArrayRef strides, const TensorOptions &options) {{\\n if (out.strides() != strides) {{\\n return {empty_strided_impl}(sizes, strides, options);\\n }}\\n return std::nullopt;\\n}}\\n\\\"\\\"\\\"\\n ]\\n )\\n\\n\\ndef gen_resize_out_helper(backend_index: BackendIndex) -> list[str]:\\n if backend_index.dispatch_key == DispatchKey.CompositeExplicitAutogradNonFunctional:\\n # The function isn't used by this key (since only functional ops have a kernel for this key),\\n # so we need to not include it to avoid a defined-but-not-used error.\\n return []\\n return [\\n \\\"\\\"\\\"\\nvoid resize_out(const Tensor &out, IntArrayRef sizes, IntArrayRef strides, const TensorOptions &options) {\\n TORCH_CHECK(options.dtype() == out.dtype(),\\n \\\"Expected out tensor to have dtype \\\", options.dtype(), \\\", but got \\\", out.dtype(), \\\" instead\\\");\\n TORCH_CHECK(options.device() == out.device(),\\n \\\"Expected out tensor to have device \\\", options.device(), \\\", but got \\\", out.device(), \\\" instead\\\");\\n const bool resized = at::native::resize_output(out, sizes);\\n // Only restride if a resize occurred; otherwise we ignore the (advisory)\\n // strides from the meta function and directly use the output tensor's\\n // preexisting strides\\n if (resized) {\\n if (!strides.empty()) {\\n TORCH_INTERNAL_ASSERT(!options.memory_format_opt().has_value());\\n // TODO: avoid the redispatch here\\n out.as_strided_(sizes, strides);\\n } else if (options.memory_format_opt().has_value()) {\\n out.unsafeGetTensorImpl()->empty_tensor_restride(*options.memory_format_opt());\\n }\\n }\\n}\\n\\\"\\\"\\\"\\n ]\\n\\n\\ndef gen_check_inplace_helper(backend_index: BackendIndex) -> list[str]:\\n return [\\n \\\"\\\"\\\"\\nvoid check_inplace(const Tensor &self, IntArrayRef sizes, const TensorOptions &options) {\\n // These checks are needed on those operators that:\\n // 1) don't use 'TensorIterator' (e.g. 'addmm' and 'baddbmm')\\n // 2) have particular typing rules (e.g. 'cumsum' and 'cumprod')\\n // For other operators (e.g. 'add'), 'TensorIterator' already checks\\n // these things separately.\\n TORCH_CHECK(options.dtype() == self.dtype(),\\n \\\"Bad in-place call: \\\",\\n \\\"input tensor dtype \\\", self.dtype(), \\\" and output tensor dtype \\\", options.dtype(), \\\" should match\\\");\\n TORCH_CHECK(options.device() == self.device(),\\n \\\"Bad in-place call: \\\",\\n \\\"input tensor device \\\", self.device(), \\\" and output tensor device \\\", options.device(), \\\" should match\\\");\\n TORCH_CHECK(sizes == self.sizes(),\\n \\\"Bad in-place call: \\\",\\n \\\"input tensor size \\\", self.sizes(), \\\" and output tensor size \\\", sizes, \\\" should match\\\");\\n}\\n\\\"\\\"\\\"\\n ]\\n\\n\\ndef gen_registration_helpers(backend_index: BackendIndex) -> list[str]:\\n return [\\n 'C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED(\\\"-Wunused-function\\\")',\\n *gen_create_out_helper(backend_index),\\n *gen_resize_out_helper(backend_index),\\n *gen_check_inplace_helper(backend_index),\\n *gen_maybe_create_proxy_helper(backend_index),\\n \\\"C10_DIAGNOSTIC_POP()\\\",\\n ]\\n\\n\\n# Generates Register{dispatch}.cpp (e.g., RegisterCPU.cpp).\\n#\\n# - The primary function of this file is to register all of the\\n# implementations for the given dispatch key to the dispatcher,\\n# so they are available for use in PyTorch. If dispatch is\\n# None, we generate schema (def) registrations and catchall\\n# registrations.\\n# - The secondary function of this file is to generate a wrapper\\n# around functions. In CPUType these wrappers do nothing\\n# (and should be removed), but in other cases they handle\\n# DeviceGuard. A small extra benefit of wrappers is they\\n# are not overloaded, so they can be used in the registration\\n# API without having to disambiguate which overload you want\\n# (as would be the case if you directly registered native::\\n# functions).\\n# - The tertiary function of this file is to generate *static*\\n# cpp API bindings which can be used to bypass dispatcher\\n# directly to kernels, but with user-friendly cpp-style API\\n@dataclass(frozen=True)\\nclass RegisterDispatchKey:\\n backend_index: BackendIndex\\n\\n target: Literal[\\n Target.ANONYMOUS_DEFINITION,\\n Target.NAMESPACED_DEFINITION,\\n Target.NAMESPACED_DECLARATION,\\n Target.REGISTRATION,\\n ]\\n\\n # Selector object to determine which operators to generate\\n # registration code for.\\n selector: SelectiveBuilder\\n\\n # Whether or not we are actually code-genning for ROCm\\n rocm: bool\\n\\n # Whether or not to generate symint registrations or not. External users\\n # of codegen who don't care about symints can set this to false to get\\n # non-SymInt codegen\\n symint: bool\\n\\n # The class that all unstructured native functions live under. This is used to improve\\n # compiler error messages when a kernel writer adds a native function with the wrong signature.\\n # This is only used in unstructured kernels, since structured kernels already live in a class.\\n # Finally, this field is currently Optional because it is only used by external backends.\\n # It would be nice if we can add the same logic to in-tree kernels too, but that requires updating\\n # all of the existing kernel signatures scattered across aten/src/ATen/native.\\n class_method_name: str | None\\n\\n # Only set to true in lightweight dispatch. If lightweight dispatch is enabled we are registering\\n # operators into JIT op registry, thus we need to avoid generating code to register into the dispatcher.\\n skip_dispatcher_op_registration: bool\\n\\n @staticmethod\\n def gen_device_check(\\n type: DeviceCheckType, args: list[Argument], method_name: str\\n ) -> str:\\n if type == DeviceCheckType.NoCheck:\\n return \\\" // No device check\\\\n\\\"\\n\\n device_check = \\\"std::optional common_device = std::nullopt;\\\\n\\\"\\n device_check += \\\"(void)common_device; // Suppress unused variable warning\\\\n\\\"\\n for arg in args:\\n # Only tensor like arguments are eligible\\n if arg.type.is_tensor_like():\\n device_check += f\\\"\\\"\\\"\\n c10::impl::check_and_update_common_device(common_device, {arg.name}, \\\"{method_name}\\\", \\\"{arg.name}\\\");\\\"\\\"\\\"\\n return device_check\\n\\n @method_with_native_function\\n def __call__(self, f: NativeFunctionsGroup | NativeFunction) -> list[str]:\\n if isinstance(f, NativeFunctionsGroup):\\n g: NativeFunctionsGroup = f\\n # Note: We call gen_structured() if the operator is marked structured, regardless of the backend.\\n # gen_structured() has special logic to handle auto-generated kernels.\\n if g.structured:\\n return self.gen_structured(g)\\n else:\\n return list(\\n mapMaybe(lambda f: self.gen_unstructured(f, g), g.functions())\\n )\\n elif isinstance(f, NativeFunction):\\n r = self.gen_unstructured(f)\\n return [] if r is None else [r]\\n else:\\n assert_never(f)\\n\\n def wrapper_kernel_sig(\\n self, f: NativeFunction\\n ) -> NativeSignature | DispatcherSignature:\\n # The prefix is just to ensure uniqueness. The Dispatcher API doesn't guarantee unique kernel names.\\n return DispatcherSignature.from_schema(\\n f.func,\\n prefix=f\\\"wrapper_{self.backend_index.dispatch_key}_{f.func.name.overload_name}_\\\",\\n symint=self.symint,\\n )\\n\\n def gen_out_inplace_wrapper(\\n self, f: NativeFunction, g: NativeFunctionsGroup | None\\n ) -> str | None:\\n if g is None:\\n return None\\n k = f.func.kind()\\n if k is SchemaKind.inplace:\\n copy_op = \\\"at::_copy_from\\\"\\n elif k is SchemaKind.out:\\n copy_op = \\\"at::_copy_from_and_resize\\\"\\n else:\\n raise AssertionError(\\\"gen_out_inplace_wrapper called on a functional op\\\")\\n\\n sig = self.wrapper_kernel_sig(f)\\n name = sig.name()\\n\\n func_res = f\\\"{name}_tmp\\\"\\n return_names = cpp.return_names(f)\\n if len(return_names) > 1:\\n updates = \\\"\\\\n \\\".join(\\n f\\\"{copy_op}(std::get<{i}>({func_res}), {ret_name});\\\"\\n for i, ret_name in enumerate(return_names)\\n )\\n returns = f'{sig.returns_type().cpp_type()}({\\\", \\\".join(return_names)})'\\n elif len(return_names) == 1:\\n ret_name = return_names[0]\\n updates = f\\\"{copy_op}({func_res}, {ret_name});\\\"\\n returns = ret_name\\n else:\\n assert len(f.func.arguments.out) == 1\\n returns = \\\"\\\"\\n out_arg = f.func.arguments.out[0]\\n if out_arg.type.is_list_like():\\n updates = f\\\"\\\"\\\"\\\\\\n for (int64_t i = 0; i < {func_res}.size(); ++i) {{\\n {copy_op}({func_res}[i], {out_arg.name}[i]);\\n }}\\\"\\\"\\\"\\n else:\\n updates = f\\\"{copy_op}({func_res}, {out_arg.name});\\\"\\n\\n functional_sig = self.wrapper_kernel_sig(g.functional)\\n wrapper_name = sig.name()\\n\\n return f\\\"\\\"\\\"\\\\\\n{sig.defn(name=wrapper_name)} {{\\n auto {func_res} = {functional_sig.name()}({\\\", \\\".join(e.expr for e in translate(sig.arguments(), functional_sig.arguments()))});\\n {updates}\\n return {returns};\\n}}\\n\\\"\\\"\\\"\\n\\n def gen_structured(self, g: NativeFunctionsGroup) -> list[str]:\\n metadata = self.backend_index.get_kernel(g)\\n if self.backend_index.dispatch_key == DispatchKey.Meta:\\n assert not self.backend_index.has_kernel(g.out), (\\n \\\"Do not explicitly specify Meta dispatch key on structured \\\"\\n \\\"functions, they will be automatically generated for you\\\"\\n )\\n elif (\\n self.backend_index.dispatch_key\\n == DispatchKey.CompositeExplicitAutogradNonFunctional\\n ):\\n assert not self.backend_index.has_kernel(g.out), (\\n \\\"Do not explicitly specify CompositeExplicitAutograd dispatch key on structured \\\"\\n \\\"functions, they will be automatically generated for you\\\"\\n )\\n elif metadata is None or not metadata.structured:\\n return list(mapMaybe(lambda f: self.gen_unstructured(f, g), g.functions()))\\n structured_gen = StructuredRegisterDispatchKey(\\n self.backend_index,\\n self.target,\\n self.selector,\\n self.rocm,\\n self.symint,\\n self.class_method_name,\\n self.skip_dispatcher_op_registration,\\n g,\\n )\\n return list(mapMaybe(structured_gen.gen_one, g.functions()))\\n\\n def gen_unstructured(\\n self, f: NativeFunction, g: NativeFunctionsGroup | None = None\\n ) -> str | None:\\n with native_function_manager(f):\\n inplace_meta = False\\n gets_out_inplace_wrapper = False\\n if not self.backend_index.has_kernel(f):\\n if (\\n self.backend_index.dispatch_key == DispatchKey.Meta\\n and f.func.kind() is SchemaKind.inplace\\n and\\n # Defer to composites for meta implementation\\n not f.has_composite_kernel\\n and\\n # Inplace list operations are not supported\\n len(f.func.returns) == 1\\n ):\\n inplace_meta = True\\n elif (\\n not self.backend_index.use_out_as_primary\\n and g is not None\\n and gets_generated_out_inplace_wrapper(f, g, self.backend_index)\\n ):\\n # We want to generate inplace/out wrappers, that don't have a kernel for the backend.\\n gets_out_inplace_wrapper = True\\n else:\\n return None\\n if f.manual_kernel_registration:\\n return None\\n\\n if (\\n self.target is Target.REGISTRATION\\n and not self.selector.is_native_function_selected(f)\\n ):\\n return None\\n\\n sig = self.wrapper_kernel_sig(f)\\n\\n name = sig.name()\\n returns_type = sig.returns_type().cpp_type()\\n args = sig.arguments()\\n args_str = \\\", \\\".join(a.defn() for a in args)\\n\\n # See Note [Direct dispatch bindings]\\n cpp_sig_group = CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=False\\n )\\n\\n # TODO: dedupe this with the structured codegen\\n if self.target is Target.NAMESPACED_DECLARATION:\\n result = \\\"\\\"\\n for cpp_sig in cpp_sig_group.signatures(symint=self.symint):\\n result += f\\\"TORCH_API {cpp_sig.decl()};\\\\n\\\"\\n return result\\n elif self.target is Target.NAMESPACED_DEFINITION:\\n\\n def generate_defn(cpp_sig: CppSignature) -> str:\\n return f\\\"\\\"\\\"\\n{cpp_sig.defn()} {{\\nreturn {sig.name()}({', '.join(e.expr for e in translate(cpp_sig.arguments(), sig.arguments()))});\\n}}\\n\\\"\\\"\\\"\\n\\n result = \\\"\\\"\\n for cpp_sig in cpp_sig_group.signatures(symint=self.symint):\\n result += generate_defn(cpp_sig)\\n return result\\n\\n elif self.target is Target.ANONYMOUS_DEFINITION:\\n # short circuit for inplace_meta\\n if inplace_meta:\\n assert f.func.arguments.self_arg is not None\\n self_arg_name = f.func.arguments.self_arg.argument.name\\n # TODO: handle in place on tensor list\\n return f\\\"\\\"\\\"\\n{returns_type} {name}({args_str}) {{\\n TORCH_CHECK_NOT_IMPLEMENTED({self_arg_name}.is_meta(),\\n \\\"Cannot inplace into non-meta tensor with meta tensor argument\\\");\\n return {self_arg_name};\\n}}\\n\\\"\\\"\\\"\\n\\n # short circuit for generated inplace/out wrappers\\n if gets_out_inplace_wrapper:\\n return self.gen_out_inplace_wrapper(f, g)\\n\\n metadata = self.backend_index.get_kernel(f)\\n if metadata is None:\\n return None\\n if self.class_method_name is None:\\n impl_name = f\\\"{metadata.cpp_namespace}::{metadata.kernel}\\\"\\n else:\\n impl_name = f\\\"{metadata.cpp_namespace}::{self.class_method_name}::{metadata.kernel}\\\"\\n\\n kernel_sig = kernel_signature(f, self.backend_index)\\n\\n args_exprs_str = \\\", \\\".join(\\n e.expr\\n for e in translate(\\n sig.arguments(), kernel_sig.arguments(), method=False\\n )\\n )\\n\\n device_check = \\\" // No device check\\\\n\\\"\\n # Backends that require device guards presumably also require device checks.\\n if self.backend_index.device_guard:\\n device_check_args = itertools.chain(\\n f.func.arguments.out, f.func.arguments.flat_positional\\n )\\n device_check = RegisterDispatchKey.gen_device_check(\\n f.device_check, list(device_check_args), name\\n )\\n\\n device_guard = \\\"// DeviceGuard omitted\\\" # default\\n if f.device_guard and self.backend_index.device_guard:\\n has_tensor_options = any(\\n isinstance(a, TensorOptionsArguments)\\n for a in f.func.arguments.non_out\\n )\\n if has_tensor_options:\\n # kernel is creating a tensor\\n device_guard = \\\"\\\"\\\"\\n const DeviceGuard device_guard(device_or_default(device));\\\"\\\"\\\"\\n\\n # CUDA requires special handling\\n if is_cuda_dispatch_key(self.backend_index.dispatch_key):\\n device_guard = (\\n f\\\"globalContext().lazyInitCUDA();\\\\n{device_guard}\\\"\\n )\\n else:\\n # kernel is operating on existing tensors\\n\\n # There is precedence for which argument we use to do\\n # device guard. This describes the precedence order.\\n self_arg = (\\n [f.func.arguments.self_arg.argument]\\n if f.func.arguments.self_arg is not None\\n else []\\n )\\n candidate_args = itertools.chain(\\n self_arg,\\n f.func.arguments.out,\\n f.func.arguments.flat_positional,\\n )\\n\\n # Only tensor like arguments are eligible\\n device_of = next(\\n (\\n f\\\"{a.name}\\\"\\n for a in candidate_args\\n if a.type.is_tensor_like()\\n ),\\n None,\\n )\\n if device_of is not None:\\n device_guard = f\\\"const OptionalDeviceGuard device_guard(device_of({device_of}));\\\"\\n\\n return f\\\"\\\"\\\"\\\\\\nnamespace {{\\n\\n{returns_type} {name}({args_str}) {{\\n {device_check}\\n\\n {device_guard}\\n return {impl_name}({args_exprs_str});\\n}}\\n\\n}} // anonymous namespace\\n\\\"\\\"\\\"\\n\\n elif self.target is Target.REGISTRATION:\\n if f.manual_kernel_registration or self.skip_dispatcher_op_registration:\\n return None\\n else:\\n payload = f\\\"TORCH_FN({name})\\\"\\n return f'm.impl(\\\"{f.func.name}\\\",\\\\n{payload});\\\\n'\\n else:\\n assert_never(self.target)\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# STRUCTURED\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\n@dataclass(frozen=True)\\nclass StructuredRegisterDispatchKey(RegisterDispatchKey):\\n g: NativeFunctionsGroup\\n\\n def gen_class_set_output_functions(\\n self, k: SchemaKind, parent_class: str, generate_super: bool\\n ) -> str:\\n if generate_super:\\n set_output_super = f\\\"{parent_class}::set_output_raw_strided(output_idx, sizes, strides, options, names);\\\"\\n else:\\n set_output_super = \\\"\\\"\\n\\n def gen_set_output_function(name: str, maybe_create_proxy: bool) -> str:\\n return f\\\"\\\"\\\"\\nvoid set_output_{name}(\\n int64_t output_idx, IntArrayRef sizes, IntArrayRef strides,\\n TensorOptions options, DimnameList names\\n) override {{\\n{textwrap.indent(self.gen_class_set_output_body(k, maybe_create_proxy), \\\" \\\")}\\n if (!names.empty()) {{\\n namedinference::propagate_names(outputs_[output_idx], names);\\n }}\\n // super must happen after, so that downstream can use maybe_get_output\\n // to retrieve the output\\n{textwrap.indent(set_output_super, \\\" \\\")}\\n}}\\n\\\"\\\"\\\"\\n\\n return f\\\"\\\"\\\"\\n{gen_set_output_function(\\\"strided\\\", maybe_create_proxy=True)}\\n{gen_set_output_function(\\\"raw_strided\\\", maybe_create_proxy=False)}\\n\\\"\\\"\\\"\\n\\n def gen_class_set_output_body(self, k: SchemaKind, maybe_create_proxy: bool) -> str:\\n if self.backend_index.dispatch_key in [\\n DispatchKey.CUDA,\\n DispatchKey.MPS,\\n DispatchKey.CompositeExplicitAutogradNonFunctional,\\n ]:\\n maybe_set_guard = \\\"\\\"\\\"\\nauto current_device = guard_.current_device();\\nif (C10_UNLIKELY(current_device.has_value())) {\\n TORCH_INTERNAL_ASSERT(*current_device == options.device(),\\n \\\"structured kernels don't support multi-device outputs\\\");\\n} else {\\n guard_.reset_device(options.device());\\n}\\n\\\"\\\"\\\"\\n maybe_set_guard_line = maybe_set_guard + \\\"\\\\n\\\"\\n else:\\n maybe_set_guard_line = maybe_set_guard = \\\"\\\"\\n\\n if maybe_create_proxy:\\n create_proxy = \\\"\\\"\\\"\\nauto maybe_proxy = maybe_create_proxy(out, sizes, strides, options);\\nif (C10_UNLIKELY(maybe_proxy.has_value())) {\\n proxy_outputs_[output_idx] = std::move(maybe_proxy).value();\\n}\\n\\\"\\\"\\\"\\n else:\\n create_proxy = \\\"\\\"\\n\\n if k is SchemaKind.functional:\\n assert self.backend_index.dispatch_key in (\\n DispatchKey.Meta,\\n DispatchKey.CPU,\\n DispatchKey.CUDA,\\n DispatchKey.MPS,\\n DispatchKey.XPU,\\n DispatchKey.CompositeExplicitAutogradNonFunctional,\\n )\\n return f\\\"\\\"\\\"{maybe_set_guard_line}\\noutputs_[output_idx] = create_out(sizes, strides, options);\\\"\\\"\\\"\\n elif k is SchemaKind.inplace:\\n return f\\\"\\\"\\\"{maybe_set_guard_line}\\nconst auto& out = outputs_[output_idx].get();\\ncheck_inplace(out, sizes, options);\\n{create_proxy}\\\"\\\"\\\"\\n elif k is SchemaKind.out:\\n return f\\\"\\\"\\\"{maybe_set_guard_line}\\nconst auto& out = outputs_[output_idx].get();\\nresize_out(out, sizes, strides, options);\\n{create_proxy}\\\"\\\"\\\"\\n elif k is SchemaKind.mutable or k is SchemaKind.scratch:\\n raise AssertionError(\\n f\\\"{k} structured operators are currently not supported\\\"\\n )\\n else:\\n assert_never(k)\\n\\n # returns the definition of a ctor, as well as how to construct\\n # this class to a variable named op\\n def gen_class_ctor(self, k: SchemaKind, class_name: str, returns: int) -> str:\\n if k is SchemaKind.functional:\\n return \\\"\\\"\\n elif k is SchemaKind.inplace:\\n # TODO: Make sure out argument is guaranteed to be self\\n return f\\\"{class_name}(Tensor& self) : outputs_{{std::ref(self)}} {{}}\\\"\\n elif k is SchemaKind.out:\\n out_args = \\\", \\\".join(f\\\"Tensor& out{i}\\\" for i in range(returns))\\n out_refs = \\\", \\\".join(f\\\"std::ref(out{i})\\\" for i in range(returns))\\n return f\\\"{class_name}({out_args}) : outputs_{{ {out_refs} }} {{}}\\\"\\n elif k is SchemaKind.mutable or k is SchemaKind.scratch:\\n raise AssertionError(\\n f\\\"{k} structured operators are currently not supported\\\"\\n )\\n else:\\n assert_never(k)\\n\\n def gen_class(\\n self,\\n f: NativeFunction,\\n k: SchemaKind,\\n *,\\n class_name: str,\\n parent_class: str,\\n generate_super: bool,\\n ) -> str:\\n if k is SchemaKind.functional:\\n output_type = \\\"Tensor\\\"\\n output_value = \\\"outputs_[output_idx]\\\"\\n proxy_field = \\\"\\\"\\n elif k is SchemaKind.inplace:\\n output_type = \\\"std::reference_wrapper\\\"\\n output_value = \\\"proxy_outputs_[output_idx].has_value() ? *proxy_outputs_[output_idx] : outputs_[output_idx].get()\\\"\\n proxy_field = f\\\"std::array<::std::optional, {len(f.func.returns)}> proxy_outputs_;\\\"\\n elif k is SchemaKind.out:\\n output_type = \\\"std::reference_wrapper\\\"\\n output_value = \\\"proxy_outputs_[output_idx].has_value() ? *proxy_outputs_[output_idx] : outputs_[output_idx].get()\\\"\\n proxy_field = f\\\"std::array<::std::optional, {len(f.func.returns)}> proxy_outputs_;\\\"\\n else:\\n raise RuntimeError(f\\\"Unsupported SchemaKind {k}\\\")\\n\\n if self.backend_index.dispatch_key == DispatchKey.CUDA:\\n if self.rocm:\\n guard_field = \\\"c10::hip::OptionalHIPGuardMasqueradingAsCUDA guard_;\\\"\\n else:\\n guard_field = \\\"c10::cuda::OptionalCUDAGuard guard_;\\\"\\n elif (\\n self.backend_index.dispatch_key\\n == DispatchKey.CompositeExplicitAutogradNonFunctional\\n ):\\n guard_field = \\\"c10::OptionalDeviceGuard guard_;\\\"\\n elif self.backend_index.dispatch_key == DispatchKey.MPS:\\n # TODO: Move to OptionalMPSGuard.\\n guard_field = \\\"c10::OptionalDeviceGuard guard_;\\\"\\n else:\\n guard_field = \\\"\\\"\\n\\n indent = \\\" \\\" * 4\\n class_ctor_str = self.gen_class_ctor(k, class_name, len(f.func.returns))\\n lines = (\\n f\\\"struct {class_name} final : public {parent_class} {{\\\",\\n f\\\"{textwrap.indent(class_ctor_str, indent)}\\\",\\n f\\\"{textwrap.indent(self.gen_class_set_output_functions(k, parent_class, generate_super), indent)}\\\",\\n \\\" const Tensor& maybe_get_output(int64_t output_idx) override {\\\",\\n f\\\" return {output_value};\\\\n\\\", # type: ignore[possibly-undefined] # TODO: audit\\n \\\" }\\\",\\n # type: ignore[possibly-undefined] # TODO: audit\\n f\\\" std::array<{output_type}, {len(f.func.returns)}> outputs_;\\\",\\n f\\\"{textwrap.indent(proxy_field, indent)}\\\", # type: ignore[possibly-undefined] # TODO: audit\\n f\\\"{textwrap.indent(guard_field, indent)}\\\",\\n \\\"};\\\",\\n )\\n return \\\"\\\\n\\\".join(line for line in lines if line)\\n\\n @method_with_native_function\\n def gen_one(self, f: NativeFunction) -> str | None:\\n assert not f.manual_kernel_registration\\n\\n if (\\n self.target is Target.REGISTRATION\\n and not self.selector.is_native_function_selected(f)\\n ):\\n return None\\n\\n # TODO: Now, there is something interesting going on here. In the code below,\\n # we generate CompositeExplicitAutogradNonFunctional implementations of functional and inplace\\n # based on the out implementation. But in fact, out is definable by\\n # functional too (just not very efficiently), and this is honestly the\\n # MORE likely situation for a backend implementor. How do we pick?\\n # Well, taking a page from Haskell type classes and default methods,\\n # we could conceivably register a circular definition (out in terms\\n # of functional, and functional in terms of out) and just require\\n # someone to implement one or the other. We'd have to do a little bit\\n # of work to not register one of these \\\"weak\\\" definitions unless there\\n # is a strong definition somewhere in the DAG! So it's not implemented yet.\\n if (\\n self.backend_index.dispatch_key\\n == DispatchKey.CompositeExplicitAutogradNonFunctional\\n and f.func.kind() is SchemaKind.out\\n ):\\n # Never generate a default implementation for out, that's what you\\n # have to define as a backend implementor\\n return None\\n\\n # Note [Direct dispatch bindings]\\n # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n # Signature of the non-dispatched function we'll expose in a header\\n # (e.g., at::cpu::add). We don't generate methods (TODO: do this\\n # when CPUTensor class is a thing); nor do we generate fallback\\n # bindings for manual_cpp_binding functions.\\n cpp_sig_group = CppSignatureGroup.from_native_function(\\n f, method=False, fallback_binding=False\\n )\\n\\n # Signature of the wrapper function we'll register to the dispatcher\\n kern = self.backend_index.get_kernel(f)\\n sig = NativeSignature(\\n f.func,\\n prefix=f\\\"wrapper_{self.backend_index.dispatch_key}_\\\",\\n symint=kern is not None and kern.supports_symint(),\\n )\\n\\n if self.target is Target.NAMESPACED_DECLARATION:\\n result = \\\"\\\"\\n for cpp_sig in cpp_sig_group.signatures(symint=self.symint):\\n result += f\\\"TORCH_API {cpp_sig.decl()};\\\\n\\\"\\n return result\\n\\n elif self.target is Target.NAMESPACED_DEFINITION:\\n\\n def generate_defn(cpp_sig: CppSignature) -> str:\\n return f\\\"\\\"\\\"\\n{cpp_sig.defn()} {{\\nreturn {sig.name()}({', '.join(e.expr for e in translate(cpp_sig.arguments(), sig.arguments()))});\\n}}\\n\\\"\\\"\\\"\\n\\n result = \\\"\\\"\\n for cpp_sig in cpp_sig_group.signatures(symint=self.symint):\\n result += generate_defn(cpp_sig)\\n return result\\n\\n elif self.target is Target.ANONYMOUS_DEFINITION:\\n k = f.func.kind()\\n\\n # Construct the body of the wrapper function with signature sig\\n sig_body = []\\n # We'll use context to keep track of any variables we've brought\\n # into scope while generating code\\n context: list[Binding | Expr] = list(sig.arguments())\\n\\n # Initialize the class corresponding to this structured\\n # operator; feeding it the output argument(s) if it is known\\n if self.backend_index.dispatch_key is DispatchKey.Meta:\\n class_name = f\\\"structured_{meta.name(self.g)}_meta_{k.name}\\\"\\n parent_class = f\\\"at::meta::structured_{meta.name(self.g)}\\\"\\n elif (\\n self.backend_index.dispatch_key\\n is DispatchKey.CompositeExplicitAutogradNonFunctional\\n ):\\n # TODO: dedup this branch\\n class_name = f\\\"structured_{meta.name(self.g)}_default_backend_{k.name}\\\"\\n parent_class = f\\\"at::meta::structured_{meta.name(self.g)}\\\"\\n else:\\n metadata = self.backend_index.get_kernel(self.g)\\n assert metadata is not None\\n class_name = f\\\"structured_{metadata.kernel}_{k.name}\\\"\\n parent_class = f\\\"{metadata.cpp_namespace}::structured_{metadata.kernel}\\\"\\n\\n if self.backend_index.device_guard:\\n device_check_args = itertools.chain(\\n f.func.arguments.out, f.func.arguments.flat_positional\\n )\\n sig_body.append(\\n RegisterDispatchKey.gen_device_check(\\n f.device_check, list(device_check_args), sig.name()\\n )\\n )\\n\\n if k is SchemaKind.functional:\\n sig_body.append(f\\\"{class_name} op;\\\")\\n elif k is SchemaKind.inplace:\\n sig_body.append(f\\\"{class_name} op(self);\\\")\\n elif k is SchemaKind.out:\\n out_args_str = \\\", \\\".join(a.name for a in f.func.arguments.out)\\n sig_body.append(f\\\"{class_name} op({out_args_str});\\\")\\n\\n # Translate the input native arguments into structured\\n # arguments for the meta call\\n meta_exprs = \\\", \\\".join(\\n e.expr\\n for e in translate(\\n context, structured.meta_arguments(self.g), method=False\\n )\\n )\\n\\n if self.g.out.precomputed:\\n # If this function group has precomputed elements, the meta function\\n # returns a struct containing them which must be saved so that it\\n # can be unpacked when generating code to call the impl.\\n sig_body.append(f\\\"auto precompute = op.meta({meta_exprs});\\\")\\n\\n # Put all of the contents of the precompute struct into the context\\n # so that translate will be able to return the correct args for the\\n # call to the impl.\\n precomputed_values = [\\n *self.g.out.precomputed.replace.values(),\\n self.g.out.precomputed.add,\\n ]\\n for precomputed_elems in precomputed_values:\\n for arg in precomputed_elems:\\n context.append(\\n Expr(\\n expr=f\\\"precompute.{arg.name}\\\",\\n type=structured.argument_type(arg, binds=arg.name),\\n )\\n )\\n\\n # Add a use of the precompute struct so FB internal compilers don't\\n # complain that there is an unused variable.\\n sig_body.append(\\\"(void)precompute;\\\")\\n else:\\n sig_body.append(f\\\"op.meta({meta_exprs});\\\")\\n\\n # After running meta, op.outputs_ is guaranteed to be valid;\\n # add it to the context\\n out_args = structured.out_arguments(self.g)\\n for i, out_arg in enumerate(out_args):\\n assert ConstRefCType(BaseCType(tensorT)) == out_arg.nctype.type\\n\\n if k is SchemaKind.out:\\n expr = f\\\"op.maybe_get_output({i})\\\"\\n else:\\n expr = f\\\"op.outputs_[{i}]\\\"\\n\\n context.append(\\n Expr(\\n expr=expr,\\n # TODO: Stop hardcoding that the output type is a Tensor. Note\\n # that for the codegen here this is fine because outputs_ is\\n # hardcoded to be tensor already\\n type=NamedCType(\\n out_arg.nctype.name, MutRefCType(BaseCType(tensorT))\\n ),\\n )\\n )\\n\\n # With the expanded context, do the impl call (if not a meta\\n # function)\\n if (\\n self.backend_index.dispatch_key\\n == DispatchKey.CompositeExplicitAutogradNonFunctional\\n ):\\n # TODO: https://github.com/pytorch/pytorch/issues/53023\\n out_sig_group = CppSignatureGroup.from_native_function(\\n self.g.out, method=False, fallback_binding=f.manual_cpp_binding\\n )\\n out_sig = out_sig_group.most_faithful_signature()\\n api_name = out_sig.name()\\n out_exprs = \\\", \\\".join(\\n e.expr\\n for e in translate(context, out_sig.arguments(), method=False)\\n )\\n # TODO: I think this means structured won't work with method\\n # only functions (but maybe you're saved by faithful? iunno.)\\n # NB: Originally I wrote this as an at::redispatch call, but\\n # I got in trouble because that meant I needed a DispatchKeySet\\n # in the wrapper function, which meant I needed a DispatchKeySet\\n # in the DispatchKeyFunctions declarations, but the defined API\\n # there does NOT permit a dispatch key set. I think you can\\n # probably unwind this by calling some function to do the TLS\\n # fetch and get the DispatchKeySet when you don't have it, but\\n # I didn't do it for this version\\n sig_body.append(f\\\"at::{api_name}({out_exprs});\\\")\\n elif self.backend_index.dispatch_key != DispatchKey.Meta:\\n impl_exprs = \\\", \\\".join(\\n e.expr\\n for e in translate(\\n context, structured.impl_arguments(self.g), method=False\\n )\\n )\\n sig_body.append(f\\\"op.impl({impl_exprs});\\\")\\n\\n # Go over each output, and check if there is a proxy created for it.\\n # If so, copy it over to the original output.\\n if k is SchemaKind.out or k is SchemaKind.inplace:\\n for i in range(len(f.func.returns)):\\n sig_body.append(\\n f\\\"if (op.proxy_outputs_[{i}].has_value()) op.outputs_[{i}].get().copy_(*op.proxy_outputs_[{i}]);\\\"\\n )\\n\\n # Destructively return the final tensors\\n # TODO: Do this in translate instead\\n if k is SchemaKind.functional:\\n if len(f.func.returns) == 1:\\n ret_expr = \\\"std::move(op.outputs_[0])\\\" # small optimization\\n else:\\n moved = \\\", \\\".join(\\n f\\\"std::move(op.outputs_[{i}])\\\"\\n for i in range(len(f.func.returns))\\n )\\n ret_expr = f\\\"std::make_tuple({moved})\\\"\\n elif k is SchemaKind.inplace:\\n ret_expr = \\\"self\\\"\\n elif k is SchemaKind.out:\\n if len(f.func.returns) == 1:\\n ret_expr = f.func.arguments.out[0].name\\n else:\\n refs = \\\", \\\".join(a.name for a in f.func.arguments.out)\\n ret_expr = f\\\"std::forward_as_tuple({refs})\\\"\\n sig_body.append(f\\\"return {ret_expr};\\\") # type: ignore[possibly-undefined] # TODO: audit\\n\\n sig_body_str = \\\"\\\\n\\\".join(sig_body)\\n\\n # For an overview of what this template code looks like, see\\n # https://github.com/pytorch/rfcs/pull/9\\n return f\\\"\\\"\\\"\\\\\\n{self.gen_class(\\nf, k,\\nclass_name=class_name,\\nparent_class=parent_class,\\ngenerate_super=self.g.out.structured_inherits is not None\\n)}\\n\\n{sig.defn()} {{\\n{sig_body_str}\\n}}\\n\\\"\\\"\\\"\\n\\n elif self.target is Target.REGISTRATION:\\n return f'm.impl(\\\"{f.func.name}\\\", TORCH_FN({sig.name()}));'\\n else:\\n assert_never(self.target)\\n # Silence mypy's \\\"Missing return statement\\\" error\\n return None\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\nfrom typing import Sequence, TYPE_CHECKING\\n\\nimport torchgen.api.ufunc as ufunc\\nfrom torchgen.api.translate import translate\\nfrom torchgen.api.types import (\\n BaseCType,\\n Binding,\\n CType,\\n Expr,\\n NamedCType,\\n opmath_t,\\n scalar_t,\\n StructuredImplSignature,\\n VectorizedCType,\\n)\\nfrom torchgen.context import with_native_function\\nfrom torchgen.model import (\\n Argument,\\n BaseTy,\\n BaseType,\\n DispatchKey,\\n NativeFunctionsGroup,\\n ScalarType,\\n UfuncKey,\\n)\\nfrom torchgen.utils import OrderedSet\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.api.ufunc import UfunctorBindings\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# CUDA STUFF\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n# NB: not bothering to generate dispatch stub forward declaration in header,\\n# we can just paste it whereever necessary\\n\\n# TODO: use BackendIndex\\n# dispatch_key: DispatchKey # only CPU/CUDA right now\\n\\n\\n# Represents functors for implementing CUDA ufuncs.\\n# Functors are templated by scalar_t because when USERS instantiate functors\\n# they are templated. A functor looks something like this:\\n#\\n# template \\n# struct CUDAFunctorOnSelf_add {\\n# using opmath_t = at::opmath_type;\\n# opmath_t other_;\\n# opmath_t alpha_;\\n# CUDAFunctorOnSelf_add(opmath_t other, opmath_t alpha)\\n# : other_(other), alpha_(alpha) {}\\n# __device__ scalar_t operator()(scalar_t self) {\\n# return ufunc::add(static_cast(self), other_, alpha_);\\n# }\\n# };\\n#\\n@dataclass(frozen=True)\\nclass UfunctorSignature:\\n g: NativeFunctionsGroup\\n scalar_tensor_idx: int | None\\n name: str\\n\\n def arguments(self) -> UfunctorBindings:\\n return ufunc.ufunctor_arguments(\\n self.g, scalar_tensor_idx=self.scalar_tensor_idx, scalar_t=scalar_t\\n )\\n\\n def fields(self) -> list[Binding]:\\n # fields are renamed to have a trailing underscore, as is conventional\\n return [b.rename(f\\\"{b.name}_\\\") for b in self.arguments().ctor]\\n\\n def returns_type(self) -> CType:\\n # TODO: don't hardcode; return type will be inferred based on tags on\\n # the native function\\n return BaseCType(scalar_t)\\n\\n def decl_fields(self) -> str:\\n return \\\"\\\\n\\\".join(f\\\"{f.type} {f.name};\\\" for f in self.fields())\\n\\n def inline_defn_ctor(self) -> str:\\n args_str = \\\", \\\".join(a.decl() for a in self.arguments().ctor)\\n # NB: hypothetically could do this with translate but the\\n # transition here is very regular\\n init_str = \\\", \\\".join(f\\\"{a.name}_({a.name})\\\" for a in self.arguments().ctor)\\n return f\\\"{self.name}({args_str}) : {init_str} {{}}\\\"\\n\\n def decl_apply(self) -> str:\\n args_str = \\\", \\\".join(a.decl() for a in self.arguments().apply)\\n return f\\\"{self.returns_type().cpp_type()} operator()({args_str}) const\\\"\\n\\n\\n@dataclass(frozen=True)\\nclass UfuncSignature:\\n g: NativeFunctionsGroup\\n name: str\\n compute_t: CType\\n\\n def arguments(self) -> list[Binding]:\\n return ufunc.ufunc_arguments(self.g, compute_t=self.compute_t)\\n\\n def call(self, ctx: Sequence[Binding | Expr]) -> str:\\n return f\\\"{self.name}({', '.join(a.expr for a in translate(ctx, self.arguments()))})\\\"\\n\\n\\n# steps:\\n# 1. take the functional signature\\n# 2. use api.ufunc to convert it to template signature. this establishes\\n# the type of the template function\\n# 3. use api.ufunc (II) to generate a split struct / operator() signature.\\n# this establish context in which we call the template signature\\n#\\n# StructuredImplSignature context\\n# ~> functor constructor sig\\n#\\n# Functor constructor context\\n# ~> functor fields sig\\n#\\n# Functor apply context (functor fields + functor apply sig)\\n# ~> template sig\\n#\\n\\n\\ndef eligible_for_binary_scalar_specialization(g: NativeFunctionsGroup) -> bool:\\n num_tensors = sum(\\n 1 for a in g.functional.func.arguments.flat_non_out if a.type.is_tensor_like()\\n )\\n return num_tensors == 2\\n\\n\\ndef compute_ufunc_cuda_functors(\\n g: NativeFunctionsGroup,\\n) -> tuple[dict[ScalarType, dict[UfuncKey, UfunctorSignature]], str]:\\n # First, build the functors.\\n ufunctor_sigs: dict[ScalarType, dict[UfuncKey, UfunctorSignature]] = {}\\n ufunctors: list[str] = []\\n loops = g.out.ufunc_inner_loop\\n scalar_tensor_idx_lookup = {\\n UfuncKey.CUDAFunctorOnSelf: 1,\\n UfuncKey.CUDAFunctorOnOther: 0,\\n UfuncKey.CUDAFunctor: None,\\n }\\n if eligible_for_binary_scalar_specialization(g):\\n keys = [\\n UfuncKey.CUDAFunctorOnSelf,\\n UfuncKey.CUDAFunctorOnOther,\\n UfuncKey.CUDAFunctor,\\n ]\\n else:\\n keys = [UfuncKey.CUDAFunctor]\\n for k in [UfuncKey.CUDAFunctorOnSelf, UfuncKey.CUDAFunctorOnOther]:\\n assert k not in loops, f\\\"cannot use {k} on non-binary function\\\"\\n for k in keys:\\n # If the key was directly defined, skip functor codegen; we assume the\\n # user already done it for us\\n if k in loops:\\n ufunctor_sig = UfunctorSignature(\\n g, scalar_tensor_idx=scalar_tensor_idx_lookup[k], name=loops[k].name\\n )\\n for dtype in loops[k].supported_dtypes:\\n ufunctor_sigs.setdefault(dtype, {})[k] = ufunctor_sig\\n continue\\n\\n # Note [ScalarOnly and Generic must match names for CUDA]\\n # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\\n # Otherwise, look in ANY of the generic entries. For simplicity of\\n # codegen, both ScalarOnly and Generic are defined, the ufunc name\\n # must match (if they didn't match, we'd have to generate distinct\\n # functors per dtype, which is awful, so we're not going to do it unless\\n # someone really forces us to)\\n ufunc_name = None\\n supported_dtypes: OrderedSet[ScalarType] = OrderedSet()\\n for lk in [UfuncKey.ScalarOnly, UfuncKey.Generic]:\\n if lk not in loops:\\n continue\\n if ufunc_name is None:\\n ufunc_name = loops[lk].name\\n else:\\n # See Note [ScalarOnly and Generic must match names for CUDA]\\n assert (\\n ufunc_name == loops[lk].name\\n ), \\\"ScalarOnly and Generic must have same ufunc name\\\"\\n supported_dtypes |= loops[lk].supported_dtypes\\n assert ufunc_name is not None\\n\\n name = f\\\"{k}_{ufunc_name}\\\"\\n ufunctor_sig = UfunctorSignature(\\n g, scalar_tensor_idx=scalar_tensor_idx_lookup[k], name=name\\n )\\n for dtype in supported_dtypes:\\n ufunctor_sigs.setdefault(dtype, {})[k] = ufunctor_sig\\n\\n ufunc_sig = UfuncSignature(\\n g, name=f\\\"ufunc::{ufunc_name}\\\", compute_t=BaseCType(opmath_t)\\n )\\n apply_ctx = ufunctor_sig.fields() + ufunctor_sig.arguments().apply\\n ufunctors.append(\\n f\\\"\\\"\\\"\\ntemplate \\nstruct {ufunctor_sig.name} {{\\n using opmath_t = at::opmath_type;\\n {ufunctor_sig.decl_fields()}\\n {ufunctor_sig.inline_defn_ctor()}\\n __device__ {ufunctor_sig.decl_apply()} {{\\n return {ufunc_sig.call(apply_ctx)};\\n }}\\n}};\\n\\\"\\\"\\\"\\n )\\n\\n return ufunctor_sigs, \\\"\\\\n\\\".join(ufunctors)\\n\\n\\n@dataclass(frozen=True)\\nclass BinaryScalarSpecializationConfig:\\n scalar_idx: int\\n ctor_tensor: str\\n ufunc_key: UfuncKey\\n\\n\\nBinaryScalarSpecializationConfigs = [\\n BinaryScalarSpecializationConfig(\\n scalar_idx=0,\\n ctor_tensor=\\\"self\\\",\\n ufunc_key=UfuncKey.CUDAFunctorOnOther,\\n ),\\n BinaryScalarSpecializationConfig(\\n scalar_idx=1,\\n ctor_tensor=\\\"other\\\",\\n ufunc_key=UfuncKey.CUDAFunctorOnSelf,\\n ),\\n]\\n\\n\\ndef compute_ufunc_cuda_dtype_body(\\n g: NativeFunctionsGroup,\\n dtype: ScalarType,\\n inner_loops: dict[UfuncKey, UfunctorSignature],\\n parent_ctx: Sequence[Binding],\\n) -> str:\\n body = \\\"using opmath_t = at::opmath_type;\\\"\\n body += \\\"if (false) {}\\\\n\\\" # for ease of codegen\\n for config in BinaryScalarSpecializationConfigs:\\n if config.ufunc_key not in inner_loops:\\n continue\\n ufunctor_sig = inner_loops[config.ufunc_key]\\n scalar_idx = config.scalar_idx + 1\\n # Make a copy and at the same time widen the type (not permissible\\n # without copy; we don't want to mutate the input argument anyway)\\n ctx: list[Expr | Binding] = list(parent_ctx)\\n ctx.append(\\n Expr(\\n expr=f\\\"iter.scalar_value({scalar_idx})\\\",\\n type=NamedCType(config.ctor_tensor, BaseCType(opmath_t)),\\n )\\n )\\n ufunctor_ctor_exprs_str = \\\", \\\".join(\\n a.expr for a in translate(ctx, ufunctor_sig.arguments().ctor)\\n )\\n\\n # NB: ufunctor must be allocated before iter.remove_operand is called,\\n # as it relies on iter\\n body += f\\\"\\\"\\\"\\\\\\nelse if (iter.is_cpu_scalar({scalar_idx})) {{\\n {ufunctor_sig.name} ufunctor({ufunctor_ctor_exprs_str});\\n iter.remove_operand({scalar_idx});\\n gpu_kernel(iter, ufunctor);\\n}}\\\"\\\"\\\"\\n\\n ufunctor_sig = inner_loops[UfuncKey.CUDAFunctor]\\n ufunctor_ctor_exprs_str = \\\", \\\".join(\\n a.expr for a in translate(parent_ctx, ufunctor_sig.arguments().ctor)\\n )\\n body += f\\\"\\\"\\\"\\nelse {{\\n gpu_kernel(iter, {ufunctor_sig.name}({ufunctor_ctor_exprs_str}));\\n}}\\n \\\"\\\"\\\"\\n return body\\n\\n\\n@with_native_function\\ndef compute_ufunc_cuda(g: NativeFunctionsGroup) -> str:\\n # First, build the functors, indexing them by dtype\\n ufunctor_sigs, ufunctors = compute_ufunc_cuda_functors(g)\\n\\n # Next, build the conditionals\\n sig = StructuredImplSignature(g, ufunc.kernel_name(g, DispatchKey.CUDA))\\n dtype_cases = []\\n for dtype, inner_ufunc_sigs in ufunctor_sigs.items():\\n dtype_cases.append(\\n f\\\"\\\"\\\"\\nAT_DISPATCH_CASE(at::ScalarType::{dtype},\\n [&]() {{\\n {compute_ufunc_cuda_dtype_body(g, dtype, inner_ufunc_sigs, sig.arguments())}\\n }}\\n)\\n\\\"\\\"\\\"\\n )\\n\\n dtype_cases_str = \\\"\\\\n\\\".join(dtype_cases)\\n\\n stub_sig = StubSignature(g)\\n\\n return f\\\"\\\"\\\"\\n{ufunctors}\\n\\n{stub_sig.type_defn()};\\n{stub_sig.dispatch_decl()};\\n\\n{stub_sig.kernel_defn()} {{\\n AT_DISPATCH_SWITCH(iter.common_dtype(), \\\"{sig.name}\\\",\\n {dtype_cases_str}\\n );\\n}}\\nREGISTER_DISPATCH({stub_sig.name}, &{stub_sig.kernel_name});\\n\\n{sig.defn()} {{\\n {stub_sig.direct_call(sig.arguments())};\\n}}\\n\\\"\\\"\\\"\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# CPU STUFF\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\n@dataclass(frozen=True)\\nclass StubSignature:\\n g: NativeFunctionsGroup\\n\\n @property\\n def name(self) -> str:\\n return f\\\"{str(self.g.functional.func.name.name)}_stub\\\"\\n\\n @property\\n def kernel_name(self) -> str:\\n return f\\\"{str(self.g.functional.func.name.name)}_kernel\\\"\\n\\n @property\\n def type_name(self) -> str:\\n return f\\\"{str(self.g.functional.func.name.name)}_fn\\\"\\n\\n def arguments(self) -> list[Binding]:\\n return ufunc.stub_arguments(self.g)\\n\\n def type(self) -> str:\\n cpp_args = self.arguments()\\n return f\\\"void(*)(TensorIteratorBase&, {', '.join(a.type for a in cpp_args)})\\\"\\n\\n def dispatch_decl(self) -> str:\\n return f\\\"DECLARE_DISPATCH({self.type_name}, {self.name})\\\"\\n\\n def dispatch_defn(self) -> str:\\n return f\\\"DEFINE_DISPATCH({self.name})\\\"\\n\\n def kernel_defn(self) -> str:\\n return f\\\"void {self.kernel_name}(TensorIteratorBase& iter, {', '.join(a.defn() for a in self.arguments())})\\\"\\n\\n def type_defn(self) -> str:\\n return f\\\"using {self.type_name} = {self.type()}\\\"\\n\\n # must be called from context where this is TensorIteratorBase*\\n def call(self, ctx: Sequence[Binding]) -> str:\\n return f\\\"{self.name}(device_type(), *this, {', '.join(a.expr for a in translate(ctx, self.arguments()))})\\\"\\n\\n # used in CUDA to skip the unnecessary dynamic dispatch\\n def direct_call(self, ctx: Sequence[Binding]) -> str:\\n return f\\\"{self.kernel_name}(*this, {', '.join(a.expr for a in translate(ctx, self.arguments()))})\\\"\\n\\n\\n@with_native_function\\ndef compute_ufunc_cpu(g: NativeFunctionsGroup) -> str:\\n stub_sig = StubSignature(g)\\n sig = StructuredImplSignature(g, ufunc.kernel_name(g, DispatchKey.CPU))\\n\\n return f\\\"\\\"\\\"\\n{stub_sig.type_defn()};\\n{stub_sig.dispatch_decl()};\\n{stub_sig.dispatch_defn()};\\n\\n{sig.defn()} {{\\n {stub_sig.call(sig.arguments())};\\n}}\\n\\\"\\\"\\\"\\n\\n\\ndef compute_ufunc_cpu_dtype_body(\\n g: NativeFunctionsGroup,\\n dtype: ScalarType,\\n inner_loops: dict[UfuncKey, UfuncSignature],\\n parent_ctx: Sequence[Binding],\\n) -> str:\\n assert UfuncKey.CPUScalar in inner_loops, f\\\"{dtype}, {inner_loops.keys()}\\\"\\n assert inner_loops.keys() <= {UfuncKey.CPUScalar, UfuncKey.CPUVector}\\n scalar_loop = inner_loops[UfuncKey.CPUScalar]\\n vec_loop = None\\n if UfuncKey.CPUVector in inner_loops:\\n vec_loop = inner_loops[UfuncKey.CPUVector]\\n\\n # NB: We DON'T use translate here, because translate is\\n # incapable of CSE'ing the scalar accesses in case it is also\\n # used by Vectorized; also, the unpacking here is very simple\\n # and only affects Scalar; everything else is implicitly captured\\n # by the lambda\\n\\n # Setup scalar in scope\\n body = []\\n ctx = []\\n for b in parent_ctx:\\n if isinstance(b.argument, Argument) and b.argument.type != BaseType(\\n BaseTy.Scalar\\n ):\\n continue\\n body.append(f\\\"auto _s_{b.name} = {b.name}.to();\\\")\\n ctx.append(Expr(f\\\"_s_{b.name}\\\", NamedCType(b.nctype.name, BaseCType(scalar_t))))\\n if vec_loop is not None:\\n for b in parent_ctx:\\n if isinstance(b.argument, Argument) and b.argument.type != BaseType(\\n BaseTy.Scalar\\n ):\\n continue\\n body.append(\\n f\\\"auto _v_{b.name} = at::vec::Vectorized(_s_{b.name});\\\"\\n )\\n ctx.append(\\n Expr(\\n f\\\"_v_{b.name}\\\",\\n NamedCType(b.nctype.name, VectorizedCType(BaseCType(scalar_t))),\\n )\\n )\\n\\n # Setup lambda signature\\n # NB: simplified version of ufunctor_arguments\\n scalar_bindings = []\\n vec_bindings = []\\n for a in g.functional.func.arguments.flat_non_out:\\n if not a.type.is_tensor_like():\\n continue\\n assert a.type == BaseType(BaseTy.Tensor)\\n scalar_bindings.append(\\n Binding(\\n name=a.name,\\n nctype=NamedCType(a.name, BaseCType(scalar_t)),\\n argument=a,\\n )\\n )\\n if vec_loop is not None:\\n vec_bindings.append(\\n Binding(\\n name=a.name,\\n nctype=NamedCType(a.name, VectorizedCType(BaseCType(scalar_t))),\\n argument=a,\\n )\\n )\\n\\n def with_ctx(b: Sequence[Binding]) -> list[Expr | Binding]:\\n r: list[Expr | Binding] = []\\n r.extend(ctx)\\n r.extend(b)\\n return r\\n\\n body_str = \\\"\\\\n\\\".join(body)\\n if vec_loop is not None:\\n return f\\\"\\\"\\\"\\n{body_str}\\ncpu_kernel_vec(iter,\\n [=]({', '.join(b.decl() for b in scalar_bindings)}) {{ return {scalar_loop.call(with_ctx(scalar_bindings))}; }},\\n [=]({', '.join(b.decl() for b in vec_bindings)}) {{ return {vec_loop.call(with_ctx(vec_bindings))}; }}\\n);\\n\\\"\\\"\\\"\\n else:\\n return f\\\"\\\"\\\"\\n{body_str}\\ncpu_kernel(iter,\\n [=]({', '.join(b.decl() for b in scalar_bindings)}) {{ return {scalar_loop.call(with_ctx(scalar_bindings))}; }}\\n);\\n\\\"\\\"\\\"\\n\\n\\n@with_native_function\\ndef compute_ufunc_cpu_kernel(g: NativeFunctionsGroup) -> str:\\n stub_sig = StubSignature(g)\\n\\n # Reindex the ufunc by dtypes; processing generic/scalaronly as well\\n loops = g.out.ufunc_inner_loop\\n ufunc_sigs: dict[ScalarType, dict[UfuncKey, UfuncSignature]] = {}\\n for k in [UfuncKey.CPUScalar, UfuncKey.CPUVector]:\\n lks = []\\n # ORDER MATTERS: this specifies overriding precedence\\n if k in loops: # should happen rarely\\n lks.append(k)\\n if UfuncKey.ScalarOnly in loops and k is UfuncKey.CPUScalar:\\n lks.append(UfuncKey.ScalarOnly)\\n if UfuncKey.Generic in loops:\\n lks.append(UfuncKey.Generic)\\n # TODO: don't hardcode ufunc:: namespace here, should be centralized smh\\n for lk in lks:\\n for dtype in loops[lk].supported_dtypes:\\n compute_t: CType\\n if k is UfuncKey.CPUScalar:\\n compute_t = BaseCType(scalar_t)\\n elif k is UfuncKey.CPUVector:\\n compute_t = VectorizedCType(BaseCType(scalar_t))\\n else:\\n raise AssertionError\\n inner_ufunc_sigs = ufunc_sigs.setdefault(dtype, {})\\n if k not in inner_ufunc_sigs:\\n inner_ufunc_sigs[k] = UfuncSignature(\\n g, name=f\\\"ufunc::{loops[lk].name}\\\", compute_t=compute_t\\n )\\n\\n # Build the conditionals\\n dtype_cases = []\\n for dtype, inner_ufunc_sigs in ufunc_sigs.items():\\n dtype_cases.append(\\n f\\\"\\\"\\\"\\nAT_DISPATCH_CASE(at::ScalarType::{dtype},\\n [&]() {{\\n {compute_ufunc_cpu_dtype_body(g, dtype, inner_ufunc_sigs, stub_sig.arguments())}\\n }}\\n)\\n\\\"\\\"\\\"\\n )\\n\\n dtype_cases_str = \\\"\\\\n\\\".join(dtype_cases)\\n return f\\\"\\\"\\\"\\nnamespace {{\\n\\n{stub_sig.kernel_defn()} {{\\n AT_DISPATCH_SWITCH(iter.common_dtype(), \\\"{stub_sig.name}\\\",\\n {dtype_cases_str}\\n );\\n}}\\n\\n}} // anonymous namespace\\n\\n{stub_sig.type_defn()};\\n{stub_sig.dispatch_decl()};\\nREGISTER_DISPATCH({stub_sig.name}, &{stub_sig.kernel_name});\\n\\\"\\\"\\\"\\n\\n\\nfrom __future__ import annotations\\n\\nimport torchgen.api.meta as meta\\nimport torchgen.api.structured as structured\\nfrom torchgen.api.types import kernel_signature\\nfrom torchgen.context import with_native_function_and_index\\nfrom torchgen.model import BackendIndex, NativeFunction, NativeFunctionsGroup\\nfrom torchgen.utils import mapMaybe\\n\\n\\n@with_native_function_and_index\\ndef gen_unstructured(f: NativeFunction, backend_index: BackendIndex) -> str | None:\\n sig = kernel_signature(f, backend_index)\\n metadata = backend_index.get_kernel(f)\\n if metadata is None:\\n return None\\n if \\\"legacy::\\\" in metadata.kernel:\\n return None\\n else:\\n prefix = \\\"static\\\" if backend_index.external else \\\"TORCH_API\\\"\\n return f\\\"{prefix} {sig.decl(name=metadata.kernel)};\\\"\\n\\n\\n@with_native_function_and_index\\ndef gen_structured(g: NativeFunctionsGroup, backend_index: BackendIndex) -> list[str]:\\n meta_name = meta.name(g)\\n out_args = structured.impl_arguments(g)\\n metadata = backend_index.get_kernel(g)\\n if metadata is None:\\n return []\\n prefix = \\\"\\\" if backend_index.external else \\\"TORCH_API \\\"\\n return [\\n f\\\"\\\"\\\"\\\\\\nstruct {prefix}structured_{metadata.kernel} : public at::meta::structured_{meta_name} {{\\nvoid impl({', '.join(a.decl() for a in out_args)});\\n}};\\n\\\"\\\"\\\"\\n ]\\n\\n\\n# Generates NativeFunctions.h, a list of forward declarations of all\\n# actual kernel definitions we keep in aten/src/ATen/native/\\n@with_native_function_and_index\\ndef compute_native_function_declaration(\\n g: NativeFunctionsGroup | NativeFunction, backend_index: BackendIndex\\n) -> list[str]:\\n metadata = backend_index.get_kernel(g)\\n if isinstance(g, NativeFunctionsGroup):\\n if metadata is not None and metadata.structured:\\n if backend_index.external:\\n # Structured hasn't been tested with external backends yet.\\n raise AssertionError(\\n \\\"Structured external backend functions are not implemented yet.\\\"\\n )\\n else:\\n return gen_structured(g, backend_index)\\n else:\\n return list(\\n mapMaybe(lambda f: gen_unstructured(f, backend_index), g.functions())\\n )\\n else:\\n x = gen_unstructured(g, backend_index)\\n return [] if x is None else [x]\\n\\n\\nfrom __future__ import annotations\\n\\nimport itertools\\nfrom abc import ABC\\nfrom dataclasses import dataclass\\nfrom typing import Any\\n\\nimport torchgen.api.dispatcher as dispatcher\\nfrom torchgen.api.lazy import (\\n getValueT,\\n isValueType,\\n LazyArgument,\\n LazyIrProperties,\\n LazyIrSchema,\\n tensorListValueT,\\n)\\nfrom torchgen.api.translate import translate\\nfrom torchgen.api.types import (\\n BaseCType,\\n Binding,\\n deviceT,\\n DispatcherSignature,\\n kernel_signature,\\n NativeSignature,\\n OptionalCType,\\n VectorCType,\\n)\\nfrom torchgen.context import method_with_native_function\\nfrom torchgen.dest.lazy_ts_lowering import ts_lowering_body\\nfrom torchgen.model import (\\n Argument,\\n BackendIndex,\\n BackendMetadata,\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n ListType,\\n NativeFunction,\\n NativeFunctionsGroup,\\n)\\n\\n\\ndef node_ctor_arg_rvalue_string(arg: LazyArgument) -> str:\\n \\\"\\\"\\\"\\n Given a LazyArgument,\\n generate a c++ string for materializing an rvalue of that arg for passing into\\n a lazy Node constructor.\\n \\\"\\\"\\\"\\n\\n # TODO: Matching on CType seems wrong; should be matching on Type\\n if isValueType(arg.lazy_type):\\n if isinstance(arg.lazy_type, BaseCType):\\n if arg.is_wrapped_scalar:\\n return f\\\"node_{arg.name}\\\"\\n elif arg.lazy_type.type is tensorListValueT:\\n return f\\\"lazy_{arg.name}_tensorlist\\\"\\n elif arg.is_symint_or_list:\\n return f\\\"GetSymIntValue({arg.name})\\\"\\n return f\\\"lazy_{arg.name}->GetIrValue()\\\"\\n elif isinstance(arg.lazy_type, OptionalCType):\\n if arg.is_symint_or_list:\\n # TODO: I don't understand when you should put lazy_ in the name\\n # or not\\n return f\\\"{arg.name} ? std::make_optional(GetSymIntValue(*{arg.name})) : ::std::nullopt\\\"\\n elif arg.is_wrapped_scalar:\\n return f\\\"node_{arg.name}\\\"\\n return (\\n f\\\"lazy_{arg.name} ? \\\"\\n f\\\"std::make_optional(lazy_{arg.name}->GetIrValue()) : \\\"\\n \\\"::std::nullopt\\\"\\n )\\n else:\\n raise AssertionError(\\n f\\\"TODO not sure if there are other valid types to handle here ({arg.lazy_type})\\\"\\n )\\n else:\\n # NB: this is here because right now we aren't treating SymInt[] as a\\n # value type; when we do this needs to move above\\n # NB: we cannot test arg.lazy_type as we've already specified it is an\\n # int64_t and so we cannot distinguish between SymInt and int64_t\\n if isinstance(arg.orig_type, ListType) and arg.orig_type.elem == BaseType(\\n BaseTy.SymInt\\n ):\\n if arg.symint:\\n return f\\\"GetSymIntArrayRefValue({arg.name})\\\"\\n else:\\n return f\\\"std::vector({arg.name}.begin(), {arg.name}.end())\\\"\\n elif isinstance(arg.lazy_type, VectorCType) and isinstance(\\n arg.lazy_type.elem, BaseCType\\n ):\\n return f\\\"std::vector<{arg.lazy_type.elem.type}>({arg.name}.begin(), {arg.name}.end())\\\"\\n elif (\\n isinstance(arg.lazy_type, OptionalCType)\\n and isinstance(arg.lazy_type.elem, VectorCType)\\n and isinstance(arg.lazy_type.elem.elem, BaseCType)\\n ):\\n return f\\\"torch::lazy::ToOptionalVector<{arg.lazy_type.elem.elem.type}>({arg.name})\\\"\\n else:\\n return f\\\"{arg.name}\\\"\\n\\n\\ndef node_ctor_inputs(schema: LazyIrSchema) -> str:\\n \\\"\\\"\\\"\\n Produce a formatted string with the arguments as passed into the constructor of a node class.\\n \\\"\\\"\\\"\\n node_ctor_values = [\\n node_ctor_arg_rvalue_string(arg) for arg in schema.filtered_args()\\n ]\\n return \\\", \\\".join(node_ctor_values)\\n\\n\\ndef gen_fallback_code(\\n schema: LazyIrSchema,\\n sig: DispatcherSignature | NativeSignature,\\n overload_name: str,\\n) -> str:\\n \\\"\\\"\\\"\\n Generate code that falls back to eager conditioned on a predicate\\n \\\"\\\"\\\"\\n dispatcher_sig = DispatcherSignature.from_schema(schema.func)\\n exprs = translate(sig.arguments(), dispatcher_sig.arguments())\\n fallback_args = \\\",\\\\n \\\".join([a.expr for a in exprs])\\n if len(overload_name):\\n aten_op_str = f\\\"ATEN_OP2({schema.aten_name}, {overload_name})\\\"\\n else:\\n aten_op_str = f\\\"ATEN_OP({schema.aten_name})\\\"\\n return f\\\"\\\"\\\"\\n if (force_eager_fallback({aten_symbol(schema)})) {{\\n return at::native::call_fallback_fn_symint<<c_eager_fallback, {aten_op_str}>::call(\\n {fallback_args}\\n );\\n }}\\n\\\"\\\"\\\"\\n\\n\\ndef aten_symbol(schema: LazyIrSchema) -> str:\\n missing_interned_strings = {\\n \\\"sigmoid_backward\\\",\\n }\\n if schema.aten_name in missing_interned_strings:\\n return f'c10::Symbol::fromQualString(\\\"aten::{schema.aten_name}\\\")'\\n\\n if not schema.aten_name.startswith(\\\"at::\\\"):\\n return f\\\"at::aten::{schema.aten_name}\\\"\\n else:\\n return schema.aten_name\\n\\n\\n# converts all tensor-like arguments to meta tensors. Returns:\\n# (1) a string containing all of the logic that does the conversions.\\n# (2) a context, to be used by translate(), with all of the relevant bindings.\\ndef convert_to_meta_tensors(sig: DispatcherSignature) -> tuple[str, list[Binding]]:\\n context: list[Binding] = []\\n unwrapped_tensor_args: list[str] = []\\n for arg in sig.arguments():\\n if isinstance(arg.argument, Argument) and arg.argument.type.is_tensor_like():\\n unwrapped_name = f\\\"{arg.name}_meta\\\"\\n unwrapped_tensor_args.append(\\n f\\\"auto {unwrapped_name} = to_meta({arg.name});\\\"\\n )\\n context.append(arg.with_name(unwrapped_name))\\n else:\\n context.append(arg)\\n unwrap_tensor_args_str = \\\"\\\\n \\\".join(unwrapped_tensor_args)\\n return unwrap_tensor_args_str, context\\n\\n\\n@dataclass(frozen=True)\\nclass GenLazyIR(ABC):\\n backend_index: BackendIndex\\n backend_name: str\\n node_base: str\\n use_lazy_shape: bool\\n\\n @method_with_native_function\\n def __call__(self, f: NativeFunctionsGroup | NativeFunction) -> list[str]:\\n func = f.functional.func if isinstance(f, NativeFunctionsGroup) else f.func\\n metadata = self.backend_index.get_kernel(\\n f.functional if isinstance(f, NativeFunctionsGroup) else f\\n )\\n schema = LazyIrSchema(\\n func, symint=metadata is not None and metadata.supports_symint()\\n )\\n return self.gen(schema)\\n\\n # there is no lowering functionality generated unless this IR base class is subclassed and\\n # implemented as a backend-specific node\\n def lowering_function(self, schema: LazyIrSchema) -> str:\\n return \\\"\\\"\\n\\n def create_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str:\\n return \\\"\\\"\\n\\n def can_be_reused_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str:\\n return f\\\"\\\"\\\"bool CanBeReused({node_ctor_args}) const {{\\n return false;\\n }}\\\"\\\"\\\"\\n\\n def node_base_ctor_call(self, schema: LazyIrSchema) -> str:\\n value_args = schema.filtered_args(values=True, scalars=False)\\n # backends can customize the way the node base class constructor is called,\\n # as long as all of its arguments can be generated from information available from the schema\\n base_ctor_value_args_list = []\\n for arg in value_args:\\n if isinstance(arg.lazy_type, (BaseCType, VectorCType)):\\n base_ctor_value_args_list.append(f\\\"{arg.name}\\\")\\n elif isinstance(arg.lazy_type, OptionalCType):\\n base_ctor_value_args_list.append(f\\\"{arg.name}.value_or(kNullValue)\\\")\\n else:\\n raise AssertionError(\\n f\\\"Unsupported type ({arg.lazy_type}) - add support if necessary\\\"\\n )\\n base_ctor_value_args = \\\", \\\".join(base_ctor_value_args_list)\\n\\n scalar_args = schema.filtered_args(values=False, scalars=True)\\n\\n # Shape construction.\\n # Conditionally build shape depending on specified shape property\\n if schema.properties.ShapePrecompute:\\n shape_ctor_arg = \\\"std::move(shapes),\\\"\\n elif schema.properties.ShapeCompute:\\n shape_args = [a.name for a in value_args]\\n shape_args.extend(a.name for a in scalar_args)\\n shape_ctor_arg = f\\\"compute_shape_{schema.name}({', '.join(shape_args)}),\\\"\\n elif schema.properties.ShapeCache:\\n shape_args = [f\\\"operand({i})\\\" for i in range(len(value_args))]\\n shape_args.extend(a.name for a in scalar_args)\\n shape_ctor_arg = f\\\"[&](){{ return compute_shape_{schema.name}({', '.join(shape_args)})[0]; }},\\\"\\n else:\\n shape_ctor_arg = \\\"\\\"\\n\\n scalar_hashes = \\\", \\\".join(f\\\"{a.name}\\\" for a in scalar_args)\\n\\n return f\\\"\\\"\\\"{self.node_base}(\\n {schema.node_name}::ClassOpKind(),\\n OpList{{{base_ctor_value_args}}},\\n {shape_ctor_arg}\\n /* num_outputs */ {len(schema.returns)},\\n torch::lazy::MHash({scalar_hashes}))\\\"\\\"\\\"\\n\\n def gen(self, schema: LazyIrSchema) -> list[str]:\\n opkind = schema.opkind or aten_symbol(schema)\\n\\n # for now, we just want one IR class decl and soon after also the method defs\\n # and we use the functional version not out/inplace.\\n all_args = schema.filtered_args()\\n scalar_args = schema.filtered_args(values=False, scalars=True)\\n\\n ctor_args = [f\\\"const {i.lazy_type.cpp_type()}& {i.name}\\\" for i in all_args]\\n reuse_ctor_args = \\\", \\\".join(ctor_args)\\n if self.use_lazy_shape and schema.properties.ShapePrecompute:\\n ctor_args.append(\\\"std::vector&& shapes\\\")\\n node_ctor_args = \\\", \\\".join(ctor_args)\\n\\n scalar_initializers = \\\",\\\\n \\\".join(\\n [\\n # This code is just special casing the mapping from string_view -> strings\\n f\\\"{a.name}({a.name}.has_value() ? ::std::make_optional(std::string(*{a.name})) : ::std::nullopt)\\\"\\n if a.lazy_type.cpp_type() == \\\"::std::optional\\\"\\n else f\\\"{a.name}({a.name})\\\"\\n for a in scalar_args\\n ]\\n )\\n if len(scalar_initializers):\\n scalar_initializers = f\\\",\\\\n {scalar_initializers}\\\"\\n scalar_decls = \\\"\\\\n \\\".join(\\n [\\n f\\\"std::string {a.name};\\\"\\n if a.lazy_type.cpp_type() == \\\"c10::string_view\\\"\\n else f\\\"::std::optional {a.name};\\\"\\n if a.lazy_type.cpp_type() == \\\"::std::optional\\\"\\n else f\\\"{a.lazy_type.cpp_type()} {a.name};\\\"\\n for a in scalar_args\\n ]\\n )\\n optional_values = [\\n arg.name\\n for arg in schema.filtered_args(values=True, scalars=False)\\n if isinstance(arg.lazy_type, OptionalCType)\\n ]\\n has_optional_decls = \\\"\\\\n \\\".join(\\n [f\\\"bool has_{value}: 1;\\\" for value in optional_values]\\n )\\n has_optional_defs = \\\"\\\\n \\\".join(\\n [f\\\"has_{value} = !!{value};\\\" for value in optional_values]\\n )\\n members_to_string = []\\n for arg in scalar_args:\\n if isinstance(arg.lazy_type, OptionalCType):\\n value = f\\\"{arg.name}.value()\\\"\\n if arg.is_generator:\\n value = '\\\"torch.Generator()\\\"'\\n members_to_string.append(\\n f\\\"\\\"\\\"if ({arg.name}.has_value()) {{\\n ss << \\\", {arg.name}=\\\" << {value};\\n }} else {{\\n ss << \\\", {arg.name}=null\\\";\\n }}\\\"\\\"\\\"\\n )\\n else:\\n members_to_string.append(f'ss << \\\", {arg.name}=\\\" << {arg.name};')\\n members_to_string_str = \\\"\\\\n \\\".join(members_to_string)\\n\\n return [\\n f\\\"\\\"\\\"\\\\\\nclass {schema.node_name} : public {self.node_base} {{\\n public:\\n static torch::lazy::OpKind ClassOpKind() {{\\n return torch::lazy::OpKind({opkind});\\n }}\\n\\n {schema.node_name}({node_ctor_args})\\n : {self.node_base_ctor_call(schema)}{scalar_initializers}\\n {{\\n {has_optional_defs}\\n }}\\n\\n std::string ToString() const override {{\\n std::stringstream ss;\\n ss << {self.node_base}::ToString();\\n {members_to_string_str}\\n return ss.str();\\n }}\\n\\n {self.create_function(schema, reuse_ctor_args)}\\n\\n {self.can_be_reused_function(schema, reuse_ctor_args)}\\n\\n {self.lowering_function(schema)}\\n\\n {scalar_decls}\\n {has_optional_decls}\\n\\n}};\\n\\n\\\"\\\"\\\",\\n ]\\n\\n\\n@dataclass(frozen=True)\\nclass GenTSLazyIR(GenLazyIR):\\n def lowering_function(self, schema: LazyIrSchema) -> str:\\n signature = \\\"\\\"\\\"\\n torch::lazy::TSOpVector Lower(\\n std::shared_ptr function,\\n torch::lazy::TSLoweringContext* loctx) const override\\\"\\\"\\\"\\n\\n if schema.properties.LowerDeclOnly:\\n return f\\\"{signature};\\\"\\n elif schema.properties.Lower:\\n return f\\\"\\\"\\\"{signature} {{\\n {ts_lowering_body(schema)}\\n }}\\n \\\"\\\"\\\"\\n else:\\n return \\\"\\\"\\n\\n def create_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str:\\n signature = f\\\"static NodePtr Create({node_ctor_args})\\\"\\n if schema.properties.CreateFnDeclOnly:\\n return f\\\"{signature};\\\"\\n elif not schema.properties.CreateFn:\\n return \\\"\\\"\\n return f\\\"\\\"\\\"{signature} {{\\n return ReuseOrMakeNode<{schema.node_name}>(data);\\n }}\\\"\\\"\\\"\\n\\n def can_be_reused_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str:\\n signature = f\\\"bool CanBeReused({node_ctor_args}) const\\\"\\n if schema.properties.CanBeReusedDeclOnly:\\n return f\\\"{signature};\\\"\\n elif not schema.properties.CanBeReused:\\n return \\\"\\\"\\n value_comparison = []\\n for arg in itertools.chain(schema.positional_values, schema.keyword_values):\\n if isinstance(arg.lazy_type, OptionalCType):\\n value_comparison.append(\\n f\\\"nullable_operand(i++) == {arg.name}.value_or(kNullValue)\\\"\\n )\\n else:\\n value_comparison.append(f\\\"operand(i++) == {arg.name}\\\")\\n for arg in itertools.chain(schema.positional_scalars, schema.keyword_scalars):\\n if isinstance(arg.lazy_type, OptionalCType):\\n value_comparison.append(\\n f\\\"((!this->{arg.name}&&!{arg.name}) || (this->{arg.name}&&{arg.name} && *(this->{arg.name}) == *{arg.name}))\\\"\\n )\\n else:\\n value_comparison.append(f\\\"this->{arg.name} == {arg.name}\\\")\\n value_comparison_str = \\\" &&\\\\n \\\".join(value_comparison)\\n\\n return f\\\"\\\"\\\"{signature} {{\\n size_t i = 0;\\n return ({value_comparison_str});\\n }}\\\"\\\"\\\"\\n\\n\\n@dataclass(frozen=True)\\nclass GenLazyNativeFuncDefinition:\\n class_method_name: str\\n backend_index: BackendIndex\\n tensor_class: str\\n gen_forced_fallback_code: bool\\n backend_namespace: str\\n get_tensorlist: str\\n get_tensor_or_wrap_number: str\\n try_get_tensor: str\\n metrics_counter: str\\n create_tensor: str\\n create_from_first_tensor: bool\\n create_aten_from_ltc_tensor: str\\n tuple_aten_from_ltc_tensors: str\\n lazy_tensor_ptr: str\\n get_device_fn: str\\n\\n def lazy_tensor_decls(self, func: NativeFunction, schema: LazyIrSchema) -> str:\\n value_args = schema.filtered_args(values=True, scalars=False)\\n # Generates lazy_{name} variables for LazyTensors wrapping input tensors\\n lazy_tensor_decls: list[str] = []\\n for arg in value_args:\\n if arg.is_wrapped_scalar:\\n if isinstance(arg.lazy_type, OptionalCType):\\n lazy_tensor_decls.append(\\n f\\\"\\\"\\\"auto node_{arg.name} = {arg.name} ?\\n std::make_optional(torch::lazy::LazyGraphExecutor::Get()->\\n GetIrValueForScalarFromCodegen(*{arg.name}, *common_device)):\\n ::std::nullopt;\\\"\\\"\\\"\\n )\\n else:\\n lazy_tensor_decls.append(\\n f\\\"\\\"\\\"auto node_{arg.name} = torch::lazy::LazyGraphExecutor::Get()->\\n GetIrValueForScalarFromCodegen({arg.name}, *common_device);\\\"\\\"\\\"\\n )\\n elif arg.is_symint_or_list:\\n continue # values are extracted in isValueType\\n elif isinstance(arg.lazy_type, BaseCType):\\n if arg.lazy_type.type is tensorListValueT:\\n lazy_tensor_decls.append(\\n f\\\"auto lazy_{arg.name}_tensorlist = \\\"\\n f\\\"{self.backend_namespace}::{self.get_tensorlist}({arg.name});\\\"\\n )\\n else:\\n lazy_tensor_decls.append(\\n f\\\"{self.lazy_tensor_ptr} lazy_{arg.name} = \\\"\\n f\\\"{self.backend_namespace}::{self.get_tensor_or_wrap_number}({arg.name}, *common_device);\\\"\\n )\\n elif isinstance(arg.lazy_type, OptionalCType):\\n assert arg.lazy_type.elem == BaseCType(getValueT()), arg.lazy_type.elem\\n # TODO(alanwaketan): Maybe we want to apply GetLtcTensorOrCreateForWrappedNumber here, but hold it\\n # until we encounter a real world example.\\n lazy_tensor_decls.append(\\n f\\\"{self.lazy_tensor_ptr} lazy_{arg.name} = \\\"\\n f\\\"{self.backend_namespace}::{self.try_get_tensor}({arg.name}.value_or(at::Tensor()));\\\"\\n )\\n else:\\n raise AssertionError(\\n f\\\"TODO not sure if there are other valid types to handle here ({arg.lazy_type})\\\"\\n )\\n return (\\\"\\\\n \\\").join(lazy_tensor_decls)\\n\\n def force_eager_fallback(\\n self,\\n func: NativeFunction,\\n schema: LazyIrSchema,\\n metadata: BackendMetadata,\\n sig: DispatcherSignature | NativeSignature,\\n ) -> str:\\n if self.gen_forced_fallback_code:\\n return gen_fallback_code(\\n schema, sig, overload_name=func.func.name.overload_name\\n )\\n return \\\"\\\"\\n\\n def metrics(self, func: NativeFunction, schema: LazyIrSchema) -> str:\\n return f\\\"{self.metrics_counter};\\\"\\n\\n def get_device(self, func: NativeFunction, schema: LazyIrSchema) -> str:\\n value_args = schema.filtered_args(values=True, scalars=False)\\n scalar_args = schema.filtered_args(values=False, scalars=True)\\n value_types_names = [f\\\"{a.name}\\\" for a in value_args if not a.is_wrapped_scalar]\\n optional_device = OptionalCType(BaseCType(deviceT))\\n optional_devices = [\\n a.name for a in scalar_args if a.lazy_type == optional_device\\n ]\\n assert (\\n len(value_types_names) > 0 or len(optional_devices) > 0\\n ), \\\"Expected at least one Value or Device type\\\"\\n get_device_str = (\\n f\\\"{self.get_device_fn}({', '.join(value_types_names + optional_devices)})\\\"\\n )\\n return f\\\"\\\"\\\"auto common_device = {get_device_str};\\n TORCH_INTERNAL_ASSERT(common_device);\\n \\\"\\\"\\\"\\n\\n def shape_inference(self, func: NativeFunction, schema: LazyIrSchema) -> str:\\n metadata = self.backend_index.get_kernel(func)\\n assert metadata is not None\\n all_args = schema.filtered_args()\\n returns_length = len(schema.returns)\\n # call the meta kernel if it exists, to compute output shape/dtype for our IR\\n # Note [Generated LTC Shape Functions]\\n # LTC uses meta tensors from core to do shape inference when possible, and otherwise\\n # we generate a shape function declaration that needs to be manually implemented.\\n # How do we detect which ops are eligible to use meta tensors?\\n # In general we should be able to use meta tensors not just on structured operators,\\n # but also on composite operators that are implemented in terms of structured kernels.\\n # We don't currently have a way of knowing at codegen time which ops are implemented that way.\\n # This is the case for all view and view_copy operators however, so we're going to\\n # use them specifically for all of the view_copy ops (instead of manually writing shape rules for all of them).\\n is_view_copy_op = \\\"view_copy\\\" in func.tags\\n is_structured = func.structured or func.structured_delegate is not None\\n if is_structured or is_view_copy_op:\\n meta_out = \\\"\\\"\\\"\\nstd::vector shapes{torch::lazy::Shape(out_meta.scalar_type(), out_meta.sizes().vec())};\\\"\\\"\\\"\\n if returns_length > 1:\\n\\n def this_shape(i: int) -> str:\\n return f\\\"torch::lazy::Shape(std::get<{i}>(out_meta).scalar_type(), std::get<{i}>(out_meta).sizes().vec())\\\"\\n\\n shapes_str = \\\",\\\".join([this_shape(i) for i in range(returns_length)])\\n meta_out = \\\"std::vector shapes{\\\" + shapes_str + \\\"};\\\"\\n\\n # Convert tensor args to the meta device and call it.\\n # (We can't pass in the input tensors directly, because they are \\\"functional wrappers\\\".\\n # If any of the meta kernels call a tensor op and redispatch, we don't want to hit the functionalize kernels.)\\n # Even at::meta:: functions might redispatch, e.g. if they call into view ops.\\n dispatcher_sig = DispatcherSignature.from_schema(func.func)\\n meta_conversion_str, meta_call_ctx = convert_to_meta_tensors(dispatcher_sig)\\n meta_call_args = [\\n e.expr\\n for e in translate(\\n meta_call_ctx, dispatcher_sig.arguments(), method=False\\n )\\n ]\\n if is_view_copy_op:\\n # view_copy ops always have a CompositeExplicitAutogradNonFunctional kernel\\n assert func.has_composite_explicit_autograd_non_functional_kernel\\n dispatch_ns = \\\"compositeexplicitautogradnonfunctional\\\"\\n else:\\n dispatch_ns = \\\"meta\\\"\\n aten_name = schema.aten_name\\n # TODO: this is trolling\\n if func.func.has_symint() and metadata.supports_symint():\\n aten_name += \\\"_symint\\\"\\n shape_str = f\\\"\\\"\\\"\\\\\\n {meta_conversion_str}\\n auto out_meta = at::{dispatch_ns}::{aten_name}({', '.join(meta_call_args)});\\n {meta_out}\\\"\\\"\\\"\\n else:\\n shape_sig = ComputeShapeSignature(\\n metadata.kernel, func, symint=metadata.supports_symint()\\n )\\n shape_str = f\\\"\\\"\\\"\\n auto shapes = {shape_sig.shape_call};\\\"\\\"\\\"\\n\\n shape_str += f\\\"\\\"\\\"\\n TORCH_INTERNAL_ASSERT(shapes.size() == {returns_length});\\\"\\\"\\\"\\n\\n # Calculating which dimensions are symbolic\\n func_schema_str = \\\"aten::\\\" + str(func.func)\\n shape_str += f\\\"\\\"\\\"\\n if(torch::lazy::symbolicShapeEnabled()){{\\n std::vector inputs = {{ {', '.join(str(a.name) for a in all_args)} }};\\n const char* schema_str = \\\"{func_schema_str}\\\";\\n applySymbolicShapesOnLT(schema_str, inputs, shapes);\\n }}\\n \\\"\\\"\\\"\\n return shape_str\\n\\n def build_ir_node(self, func: NativeFunction, schema: LazyIrSchema) -> str:\\n node_ctor_input_str = node_ctor_inputs(schema)\\n return f\\\"\\\"\\\"torch::lazy::NodePtr node = torch::lazy::ReuseNode<{schema.node_name}>({node_ctor_input_str});\\n if (!node) {{\\n {self.shape_inference(func, schema)}\\n node = torch::lazy::MakeNode<{schema.node_name}>({node_ctor_input_str}, std::move(shapes));\\n CacheNode(node);\\n }}\\n \\\"\\\"\\\"\\n\\n def create_lazy_tensor(self, first_tensor_name: str | None = None) -> str:\\n # xla uses an instance method for tensor creation, for the time being\\n if self.create_from_first_tensor:\\n # TODO(whc) remove this if XLA switches to using static method for creation\\n assert (\\n first_tensor_name is not None\\n ), \\\"Requires first tensor to create lazy tensor\\\"\\n return f\\\"{first_tensor_name}.{self.create_tensor}\\\"\\n return f\\\"{self.backend_namespace}::{self.create_tensor}\\\"\\n\\n def return_aten_tensor(self, func: NativeFunction, schema: LazyIrSchema) -> str:\\n returns_length = len(schema.returns)\\n value_args = schema.filtered_args(values=True, scalars=False)\\n value_types_names = [f\\\"{a.name}\\\" for a in value_args if not a.is_wrapped_scalar]\\n first_tensor_name = value_types_names[0] if len(value_types_names) > 0 else None\\n bridge_str = f\\\"\\\"\\\"auto result = {self.create_aten_from_ltc_tensor}(\\n {self.create_lazy_tensor(first_tensor_name)}(std::move(node), *common_device));\\\"\\\"\\\"\\n\\n if returns_length > 1:\\n assert (\\n len(value_types_names) > 0\\n ), \\\"Code below assumes there is at least one tensor arg\\\"\\n bridge_str = f\\\"\\\"\\\"std::vector<{self.lazy_tensor_ptr}> lazy_tensors;\\n for (int i = 0; i < {returns_length}; i++) {{\\n lazy_tensors.push_back({self.create_lazy_tensor(first_tensor_name)}({getValueT()}(node, i), *common_device));\\n }}\\n auto result = {self.tuple_aten_from_ltc_tensors}<{returns_length}>(lazy_tensors);\\\"\\\"\\\"\\n\\n if schema.name.name.inplace or func.func.is_out_fn():\\n assert returns_length == 1, (\\n \\\"We assumed there was no such case where an op is an in-place variant \\\"\\n f\\\"and has tuple outputs, but got tuple of len {returns_length}.\\\"\\n )\\n bridge_str = f\\\"\\\"\\\"lazy_{first_tensor_name}->SetInPlaceIrValue(node);\\n auto& result = {first_tensor_name};\\\"\\\"\\\"\\n\\n bridge_str += \\\"\\\"\\\"\\n return result;\\\"\\\"\\\"\\n return bridge_str\\n\\n @method_with_native_function\\n def __call__(self, func: NativeFunction) -> list[str]:\\n sig = kernel_signature(func, self.backend_index)\\n metadata = self.backend_index.get_kernel(func)\\n assert metadata is not None\\n schema = LazyIrSchema(func.func, symint=metadata.supports_symint())\\n return [\\n f\\\"\\\"\\\"\\\\\\n {sig.decl(name=f\\\"{self.class_method_name}::{metadata.kernel}\\\")} {{\\n {self.force_eager_fallback(func, schema, metadata, sig)}\\n {self.metrics(func, schema)}\\n {self.get_device(func, schema)}\\n {self.lazy_tensor_decls(func, schema)}\\n {self.build_ir_node(func, schema)}\\n {self.return_aten_tensor(func, schema)}\\n }}\\\\n\\n \\\"\\\"\\\"\\n ]\\n\\n\\nclass ComputeShapeSignature:\\n \\\"\\\"\\\"\\n Here we use the base name as the suffix of the signature to avoid generating for in-place variants.\\n \\\"\\\"\\\"\\n\\n def __init__(self, kernel_name: str, f: NativeFunction, *, symint: bool) -> None:\\n self.__schema = LazyIrSchema(f.func, symint=symint)\\n self.__dispatch_args = \\\", \\\".join(\\n [a.decl() for a in dispatcher.arguments(f.func, symint=symint)]\\n )\\n self.__call_args = \\\", \\\".join(\\n [f\\\"{arg.name}\\\" for arg in self.__schema.filtered_args(generator=True)]\\n )\\n self.__kernel_name = kernel_name\\n\\n def __decl_suffix(self) -> str:\\n return f\\\"{self.__kernel_name}({self.__dispatch_args})\\\"\\n\\n def __call_suffix(self) -> str:\\n return f\\\"{self.__kernel_name}({self.__call_args})\\\"\\n\\n @property\\n def shape_decl(self) -> str:\\n return f\\\"TORCH_API std::vector compute_shape_{self.__decl_suffix()}\\\"\\n\\n @property\\n def shape_call(self) -> str:\\n return f\\\"torch::lazy::compute_shape_{self.__call_suffix()}\\\"\\n\\n\\n@dataclass(frozen=True)\\nclass GenLazyShapeInferenceDefinition:\\n backend_index: BackendIndex\\n tensor_class: str\\n\\n @method_with_native_function\\n def __call__(self, f: NativeFunction) -> list[str]:\\n metadata = self.backend_index.get_kernel(f)\\n assert metadata is not None\\n\\n # See Note [Generated LTC Shape Functions]\\n is_view_copy_op = \\\"view_copy\\\" in f.tags\\n is_structured = f.structured or f.structured_delegate is not None\\n if is_structured or is_view_copy_op:\\n return []\\n else:\\n shape_sig = ComputeShapeSignature(\\n metadata.kernel, f, symint=metadata.supports_symint()\\n )\\n return [\\\"\\\\n\\\".join([f\\\"{shape_sig.shape_decl};\\\"])]\\n\\n\\ndef generate_non_native_lazy_ir_nodes(\\n non_native: list[dict[str, Any]], gen_lazy_ir: GenLazyIR\\n) -> list[str]:\\n \\\"\\\"\\\"Generate the non-native lazy IR node classes\\\"\\\"\\\"\\n nodes = []\\n for op in non_native:\\n # Set default properties for Non-Native IRs\\n properties = LazyIrProperties(\\\"ShapeCache\\\", \\\"CanBeReused\\\", \\\"LowerDeclOnly\\\")\\n for p in op.get(\\\"properties\\\", []):\\n setattr(properties, p, True)\\n\\n # non-native is assumed to want symint bindings if you wrote symint\\n schema = LazyIrSchema(FunctionSchema.parse(op[\\\"func\\\"]), properties, symint=True)\\n schema.opkind = op.get(\\\"opkind\\\")\\n nodes.append(gen_lazy_ir.gen(schema)[0])\\n\\n return nodes\\n\\n\\nfrom torchgen.api.lazy import LazyArgument, LazyIrSchema\\nfrom torchgen.api.types import OptionalCType\\n\\n\\ndef ts_lowering_body(schema: LazyIrSchema) -> str:\\n # for now, we just want one IR class decl and soon after also the method defs\\n # and we use the functional version not out/inplace.\\n emplace_arguments = []\\n\\n def get_value(arg: LazyArgument) -> str:\\n if isinstance(arg.lazy_type, OptionalCType):\\n return f\\\"has_{arg.name} ? loctx->GetOutputOp(operand(i++)) : nullptr\\\"\\n return \\\"loctx->GetOutputOp(operand(i++))\\\"\\n\\n for arg in schema.positional_args:\\n if arg.is_lazy_value:\\n emplace_arguments.append(get_value(arg))\\n continue\\n emplace_arguments.append(f'\\\"{arg.name}\\\", {arg.name}')\\n\\n emplace_arguments_str = \\\"\\\\n \\\".join(\\n [f\\\"arguments.emplace_back({a});\\\" for a in emplace_arguments]\\n )\\n emplace_kwarg_values = [\\n f'\\\"{arg.name}\\\", {get_value(arg)}' for arg in schema.keyword_values\\n ]\\n emplace_kwarg_scalars = [\\n f'\\\"{arg.name}\\\", {arg.name}' for arg in schema.keyword_scalars\\n ]\\n emplace_kwarguments = \\\"\\\\n \\\".join(\\n [\\n f\\\"kwarguments.emplace_back({a});\\\"\\n for a in emplace_kwarg_values + emplace_kwarg_scalars\\n ]\\n )\\n return f\\\"\\\"\\\"\\\\\\n std::vector arguments;\\n std::vector kwarguments;\\n arguments.reserve({len(emplace_arguments)});\\n kwarguments.reserve({len(emplace_kwarg_values + emplace_kwarg_scalars)});\\n size_t i = 0;\\n {emplace_arguments_str}\\n {emplace_kwarguments}\\n torch::lazy::TSOpVector {schema.aten_name}_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments);\\n TORCH_CHECK_EQ({schema.aten_name}_out.size(), {len(schema.returns)});\\n\\n return {schema.aten_name}_out;\\n\\\"\\\"\\\"\\n\\n\\nfrom torchgen.dest.lazy_ir import (\\n generate_non_native_lazy_ir_nodes as generate_non_native_lazy_ir_nodes,\\n GenLazyIR as GenLazyIR,\\n GenLazyNativeFuncDefinition as GenLazyNativeFuncDefinition,\\n GenLazyShapeInferenceDefinition as GenLazyShapeInferenceDefinition,\\n)\\nfrom torchgen.dest.native_functions import (\\n compute_native_function_declaration as compute_native_function_declaration,\\n)\\nfrom torchgen.dest.register_dispatch_key import (\\n gen_registration_headers as gen_registration_headers,\\n gen_registration_helpers as gen_registration_helpers,\\n RegisterDispatchKey as RegisterDispatchKey,\\n)\\nfrom torchgen.dest.ufunc import (\\n compute_ufunc_cpu as compute_ufunc_cpu,\\n compute_ufunc_cpu_kernel as compute_ufunc_cpu_kernel,\\n compute_ufunc_cuda as compute_ufunc_cuda,\\n)\\n\\n\\nfrom __future__ import annotations\\n\\nfrom collections import defaultdict\\nfrom collections.abc import Iterable\\nfrom dataclasses import dataclass\\nfrom typing import TYPE_CHECKING\\n\\nimport yaml\\n\\nfrom torchgen.selective_build.operator import (\\n merge_debug_info,\\n merge_operator_dicts,\\n SelectiveBuildOperator,\\n strip_operator_overload_name,\\n)\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.model import NativeFunction\\n\\n\\n# A SelectiveBuilder holds information extracted from the selective build\\n# YAML specification.\\n#\\n# It includes information about the build's selectivity, the debug_info\\n# associated with this selective build (opaque string), and the set of\\n# operators that should be included in the build.\\n#\\n@dataclass(frozen=True)\\nclass SelectiveBuilder:\\n # If true, then the build is not selective, and includes all\\n # operators.\\n include_all_operators: bool\\n\\n # Debug Information at the selective/custom build level.\\n _debug_info: tuple[str, ...] | None\\n\\n # A dictionary of operator -> operator metadata.\\n operators: dict[str, SelectiveBuildOperator]\\n\\n # A dictionary of selected kernel tags and dtypes. Typically a\\n # PyTorch Operator Kernel (function) may have many code paths\\n # that are specialized for many many Tensor dtypes, so it's not\\n # one per kernel function, but there could be many per kernel\\n # function. The tag isn't a kernel function name, but some fragment\\n # of the kernel function implementation itself.\\n kernel_metadata: dict[str, list[str]]\\n\\n # ExecuTorch only. A dictionary of kernel tag -> list of (list of input\\n # dtypes for tensor-like input args).\\n # This is from selective.yaml\\n et_kernel_metadata: dict[str, list[str]]\\n\\n # A set of all the custom torch bind classes used by the selected models\\n # Stored as a set internally to remove duplicates proactively, but written\\n # as a list to yamls\\n custom_classes: set[str]\\n\\n # A set of all the build features used by the selected models\\n # Stored as a set internally to remove duplicates proactively, but written\\n # as a list to yamls\\n build_features: set[str]\\n\\n # If true, then fragments for all dtypes for all kernel functions\\n # are included as well as all custom classes. This is typically set when any one of the\\n # operator lists is generated from a mechanism other than\\n # tracing based selective build.\\n include_all_non_op_selectives: bool\\n\\n @staticmethod\\n def get_nop_selector() -> SelectiveBuilder:\\n return SelectiveBuilder.from_yaml_dict({\\\"include_all_operators\\\": True})\\n\\n @staticmethod\\n def from_yaml_dict(data: dict[str, object]) -> SelectiveBuilder:\\n valid_top_level_keys = {\\n \\\"include_all_non_op_selectives\\\",\\n \\\"include_all_operators\\\",\\n \\\"debug_info\\\",\\n \\\"operators\\\",\\n \\\"kernel_metadata\\\",\\n \\\"et_kernel_metadata\\\",\\n \\\"custom_classes\\\",\\n \\\"build_features\\\",\\n }\\n top_level_keys = set(data.keys())\\n if len(top_level_keys - valid_top_level_keys) > 0:\\n raise Exception( # noqa: TRY002\\n \\\"Got unexpected top level keys: {}\\\".format(\\n \\\",\\\".join(top_level_keys - valid_top_level_keys),\\n )\\n )\\n include_all_operators = data.get(\\\"include_all_operators\\\", False)\\n assert isinstance(include_all_operators, bool)\\n\\n debug_info = None\\n if \\\"debug_info\\\" in data:\\n di_list = data[\\\"debug_info\\\"]\\n assert isinstance(di_list, list)\\n\\n debug_info = tuple(str(x) for x in di_list)\\n\\n operators = {}\\n operators_dict = data.get(\\\"operators\\\", {})\\n assert isinstance(operators_dict, dict)\\n\\n for k, v in operators_dict.items():\\n operators[k] = SelectiveBuildOperator.from_yaml_dict(k, v)\\n\\n kernel_metadata = {}\\n kernel_metadata_dict = data.get(\\\"kernel_metadata\\\", {})\\n assert isinstance(kernel_metadata_dict, dict)\\n\\n for k, v in kernel_metadata_dict.items():\\n kernel_metadata[str(k)] = [str(dtype) for dtype in v]\\n\\n et_kernel_metadata = data.get(\\\"et_kernel_metadata\\\", {})\\n assert isinstance(et_kernel_metadata, dict)\\n\\n custom_classes = data.get(\\\"custom_classes\\\", [])\\n assert isinstance(custom_classes, Iterable)\\n custom_classes = set(custom_classes)\\n\\n build_features = data.get(\\\"build_features\\\", [])\\n assert isinstance(build_features, Iterable)\\n build_features = set(build_features)\\n\\n include_all_non_op_selectives = data.get(\\\"include_all_non_op_selectives\\\", False)\\n assert isinstance(include_all_non_op_selectives, bool)\\n\\n return SelectiveBuilder(\\n include_all_operators,\\n debug_info,\\n operators,\\n kernel_metadata,\\n et_kernel_metadata,\\n custom_classes, # type: ignore[arg-type]\\n build_features, # type: ignore[arg-type]\\n include_all_non_op_selectives,\\n )\\n\\n @staticmethod\\n def from_yaml_str(config_contents: str) -> SelectiveBuilder:\\n contents = yaml.safe_load(config_contents)\\n return SelectiveBuilder.from_yaml_dict(contents)\\n\\n @staticmethod\\n def from_yaml_path(config_path: str) -> SelectiveBuilder:\\n with open(config_path) as f:\\n contents = yaml.safe_load(f)\\n return SelectiveBuilder.from_yaml_dict(contents)\\n\\n @staticmethod\\n def from_legacy_op_registration_allow_list(\\n allow_list: set[str], is_root_operator: bool, is_used_for_training: bool\\n ) -> SelectiveBuilder:\\n operators = {}\\n for op in allow_list:\\n operators[op] = {\\n \\\"name\\\": op,\\n \\\"is_root_operator\\\": is_root_operator,\\n \\\"is_used_for_training\\\": is_used_for_training,\\n \\\"include_all_overloads\\\": True,\\n }\\n return SelectiveBuilder.from_yaml_dict(\\n {\\n \\\"operators\\\": operators,\\n \\\"include_all_non_op_selectives\\\": True,\\n }\\n )\\n\\n def is_operator_selected(self, name: str) -> bool:\\n if self.include_all_operators:\\n return True\\n\\n if name in self.operators:\\n return True\\n name = strip_operator_overload_name(name)\\n return name in self.operators and self.operators[name].include_all_overloads\\n\\n def is_native_function_selected(self, func: NativeFunction) -> bool:\\n op_name = op_name_from_native_function(func)\\n return self.is_operator_selected(op_name)\\n\\n def is_operator_selected_for_training(self, name: str) -> bool:\\n if not self.is_operator_selected(name):\\n return False\\n if self.include_all_operators:\\n return True\\n\\n not_training_op = SelectiveBuildOperator(\\n name=\\\"\\\",\\n is_root_operator=False,\\n is_used_for_training=False,\\n include_all_overloads=False,\\n _debug_info=None,\\n )\\n op = not_training_op\\n if name in self.operators:\\n op = self.operators[name]\\n\\n name = strip_operator_overload_name(name)\\n base_op = not_training_op\\n if name in self.operators:\\n base_op = self.operators[name]\\n\\n return op.is_used_for_training or (\\n base_op.include_all_overloads and base_op.is_used_for_training\\n )\\n\\n def is_native_function_selected_for_training(self, func: NativeFunction) -> bool:\\n op_name = op_name_from_native_function(func)\\n return self.is_operator_selected_for_training(op_name)\\n\\n def is_root_operator(self, name: str) -> bool:\\n if not self.is_operator_selected(name):\\n return False\\n if self.include_all_operators:\\n return True\\n\\n if name in self.operators:\\n op: SelectiveBuildOperator = self.operators[name]\\n return op.is_root_operator\\n name = strip_operator_overload_name(name)\\n if name not in self.operators:\\n return False\\n base_op: SelectiveBuildOperator = self.operators[name]\\n return base_op.include_all_overloads and base_op.is_root_operator\\n\\n def is_kernel_dtype_selected(self, kernel_tag: str, dtype: str) -> bool:\\n if self.include_all_operators or self.include_all_non_op_selectives:\\n return True\\n\\n return (\\n kernel_tag in self.kernel_metadata\\n and dtype in self.kernel_metadata[kernel_tag]\\n )\\n\\n def et_get_selected_kernels(self, op_name: str, kernel_key: list[str]) -> list[str]:\\n \\\"\\\"\\\"\\n Return a list of kernel keys that cover the used ops\\n \\\"\\\"\\\"\\n # If no kernel metadata, either it's implied by include_all_operators=True or the op is not used.\\n if op_name not in self.et_kernel_metadata:\\n return kernel_key if self.include_all_operators else []\\n # Otherwise, only return the specific kernel keys.\\n\\n result_set = set()\\n\\n for model_kernel_keys in self.et_kernel_metadata[op_name]:\\n key_found = False\\n for key in kernel_key:\\n # Don't compare the version for now\\n if (\\n key != \\\"default\\\"\\n and key.split(\\\"/\\\")[1] == model_kernel_keys.split(\\\"/\\\")[1]\\n ):\\n result_set.add(key)\\n key_found = True\\n break\\n if not key_found:\\n if \\\"default\\\" not in kernel_key:\\n raise Exception(\\\"Missing kernel for the model\\\") # noqa: TRY002\\n else:\\n result_set.add(\\\"default\\\")\\n\\n return list(result_set)\\n\\n def to_dict(self) -> dict[str, object]:\\n ret: dict[str, object] = {\\n \\\"include_all_non_op_selectives\\\": self.include_all_non_op_selectives,\\n \\\"include_all_operators\\\": self.include_all_operators,\\n }\\n operators = {}\\n for op_name, op in self.operators.items():\\n operators[op_name] = op.to_dict()\\n ret[\\\"operators\\\"] = operators\\n\\n if self._debug_info is not None:\\n ret[\\\"debug_info\\\"] = sorted(self._debug_info)\\n\\n ret[\\\"kernel_metadata\\\"] = {\\n k: sorted(v) for (k, v) in self.kernel_metadata.items()\\n }\\n\\n ret[\\\"et_kernel_metadata\\\"] = self.et_kernel_metadata\\n\\n ret[\\\"custom_classes\\\"] = sorted(self.custom_classes)\\n\\n ret[\\\"build_features\\\"] = sorted(self.build_features)\\n\\n return ret\\n\\n\\ndef merge_kernel_metadata(\\n lhs: dict[str, list[str]],\\n rhs: dict[str, list[str]],\\n) -> dict[str, list[str]]:\\n kernel_metadata: dict[str, list[str]] = {}\\n for tag_name, dtypes in list(lhs.items()) + list(rhs.items()):\\n dtypes_copy = set(dtypes)\\n if tag_name in kernel_metadata:\\n dtypes_copy |= set(kernel_metadata[tag_name])\\n\\n kernel_metadata[tag_name] = list(dtypes_copy)\\n\\n return kernel_metadata\\n\\n\\ndef merge_et_kernel_metadata(\\n lhs: dict[str, list[str]],\\n rhs: dict[str, list[str]],\\n) -> dict[str, list[str]]:\\n merge_et_kernel_metadata: dict[str, set[str]] = defaultdict(set)\\n for op in list(lhs.keys()) + list(rhs.keys()):\\n merge_et_kernel_metadata[op].update(lhs.get(op, []))\\n merge_et_kernel_metadata[op].update(rhs.get(op, []))\\n\\n return {op: sorted(val) for op, val in merge_et_kernel_metadata.items()}\\n\\n\\ndef combine_selective_builders(\\n lhs: SelectiveBuilder, rhs: SelectiveBuilder\\n) -> SelectiveBuilder:\\n include_all_operators = lhs.include_all_operators or rhs.include_all_operators\\n debug_info = merge_debug_info(lhs._debug_info, rhs._debug_info)\\n operators = merge_operator_dicts(lhs.operators, rhs.operators)\\n kernel_metadata = merge_kernel_metadata(lhs.kernel_metadata, rhs.kernel_metadata)\\n et_kernel_metadata = merge_et_kernel_metadata(\\n lhs.et_kernel_metadata, rhs.et_kernel_metadata\\n )\\n include_all_non_op_selectives = (\\n lhs.include_all_non_op_selectives or rhs.include_all_non_op_selectives\\n )\\n custom_classes = lhs.custom_classes.union(rhs.custom_classes)\\n build_features = lhs.build_features.union(rhs.build_features)\\n return SelectiveBuilder(\\n include_all_operators,\\n debug_info,\\n operators,\\n kernel_metadata,\\n et_kernel_metadata,\\n custom_classes,\\n build_features,\\n include_all_non_op_selectives,\\n )\\n\\n\\ndef op_name_from_native_function(f: NativeFunction) -> str:\\n # This was originally read from the 'operator_name_with_overload' field in the\\n # declaration dict, which was the part before the first '(' in 'schema_string'.\\n return f\\\"{f.namespace}::{f.func.name}\\\"\\n\\n\\nfrom __future__ import annotations\\n\\nfrom dataclasses import dataclass\\n\\n\\n# This class holds information about a single operator used to determine\\n# the outcome of a selective/custom PyTorch build that doesn't include\\n# registration code for all the supported operators. This is done to\\n# reduce the size of the generated binary so that it can be deployed in\\n# situations where binary size comes at a premium.\\n#\\n@dataclass(frozen=True)\\nclass SelectiveBuildOperator:\\n # The name of the operator. This includes the aten::, etc... prefix\\n # The operator name may or may not have the overload name. If this\\n # operator name does not specify an overload name, the way to determine\\n # if this entry refers to the family of operators with this base name\\n # or just the operator with this name is to look at the value of the\\n # 'include_all_overloads' flag in this class.\\n name: str\\n\\n # True if this is a root operator (i.e. called directly from a\\n # TorchScript model, etc...). An operator is considered to be a\\n # root operator if it is called directly from any one of the models\\n # that this instance of the pytorch library was built for. Hence, it\\n # may not be a root operator in all of the models that are used in\\n # this instance of the pytorch library.\\n is_root_operator: bool\\n\\n # Is this operator used for on-device training? If True, then we need to\\n # use the information to generate code in VariableType_N.cpp for registration\\n # of training related operators. Again, this is True if this operator\\n # is used for training in one or more models used by this instance of the\\n # pytorch library.\\n is_used_for_training: bool\\n\\n # If True, it indicates that this operator instance (object) refers to an\\n # operator without the overload name and should apply to all overloads\\n # which have this operator name as the base name. This flag is applicable\\n # only for objects that have operator names without a DOT (period) character\\n # in them.\\n #\\n # Note: This flag is a temporary workaround to grandfather in the current\\n # static selective (custom) build mechanism, which largely ignores overload\\n # names when determining whether to select operators for registration\\n # purposes.\\n include_all_overloads: bool\\n\\n # Debug Information at the operator level\\n _debug_info: tuple[str, ...] | None\\n\\n @staticmethod\\n def from_yaml_dict(\\n op_name: str, op_info: dict[str, object]\\n ) -> SelectiveBuildOperator:\\n allowed_keys = {\\n \\\"name\\\",\\n \\\"is_root_operator\\\",\\n \\\"is_used_for_training\\\",\\n \\\"include_all_overloads\\\",\\n \\\"debug_info\\\",\\n }\\n\\n if len(set(op_info.keys()) - allowed_keys) > 0:\\n raise Exception( # noqa: TRY002\\n \\\"Got unexpected top level keys: {}\\\".format(\\n \\\",\\\".join(set(op_info.keys()) - allowed_keys),\\n )\\n )\\n\\n if \\\"name\\\" in op_info:\\n assert op_name == op_info[\\\"name\\\"]\\n\\n is_root_operator = op_info.get(\\\"is_root_operator\\\", True)\\n assert isinstance(is_root_operator, bool)\\n\\n is_used_for_training = op_info.get(\\\"is_used_for_training\\\", True)\\n assert isinstance(is_used_for_training, bool)\\n\\n include_all_overloads = op_info.get(\\\"include_all_overloads\\\", True)\\n assert isinstance(include_all_overloads, bool)\\n\\n debug_info: tuple[str, ...] | None = None\\n if \\\"debug_info\\\" in op_info:\\n di_list = op_info[\\\"debug_info\\\"]\\n assert isinstance(di_list, list)\\n debug_info = tuple(str(x) for x in di_list)\\n\\n return SelectiveBuildOperator(\\n name=op_name,\\n is_root_operator=is_root_operator,\\n is_used_for_training=is_used_for_training,\\n include_all_overloads=include_all_overloads,\\n _debug_info=debug_info,\\n )\\n\\n @staticmethod\\n def from_legacy_operator_name_without_overload(\\n name: str,\\n ) -> SelectiveBuildOperator:\\n return SelectiveBuildOperator(\\n name=name,\\n is_root_operator=True,\\n is_used_for_training=True,\\n include_all_overloads=True,\\n _debug_info=None,\\n )\\n\\n def to_dict(self) -> dict[str, object]:\\n ret: dict[str, object] = {\\n \\\"is_root_operator\\\": self.is_root_operator,\\n \\\"is_used_for_training\\\": self.is_used_for_training,\\n \\\"include_all_overloads\\\": self.include_all_overloads,\\n }\\n if self._debug_info is not None:\\n ret[\\\"debug_info\\\"] = self._debug_info\\n\\n return ret\\n\\n\\ndef merge_debug_info(\\n lhs: tuple[str, ...] | None,\\n rhs: tuple[str, ...] | None,\\n) -> tuple[str, ...] | None:\\n # Ensure that when merging, each entry shows up just once.\\n if lhs is None and rhs is None:\\n return None\\n\\n return tuple(set((lhs or ()) + (rhs or ())))\\n\\n\\ndef combine_operators(\\n lhs: SelectiveBuildOperator, rhs: SelectiveBuildOperator\\n) -> SelectiveBuildOperator:\\n if str(lhs.name) != str(rhs.name):\\n raise Exception( # noqa: TRY002\\n f\\\"Expected both arguments to have the same name, but got '{str(lhs.name)}' and '{str(rhs.name)}' instead\\\"\\n )\\n\\n return SelectiveBuildOperator(\\n name=lhs.name,\\n # Consider this operator to be a root operator if it is a\\n # root operator in any of the models used in this instance of\\n # the pytorch library.\\n is_root_operator=lhs.is_root_operator or rhs.is_root_operator,\\n # Consider this operator to be a training operator if it is\\n # an operator used for training in any of the models used\\n # in this instance of the pytorch library.\\n is_used_for_training=lhs.is_used_for_training or rhs.is_used_for_training,\\n include_all_overloads=lhs.include_all_overloads or rhs.include_all_overloads,\\n _debug_info=merge_debug_info(lhs._debug_info, rhs._debug_info),\\n )\\n\\n\\ndef merge_operator_dicts(\\n lhs: dict[str, SelectiveBuildOperator],\\n rhs: dict[str, SelectiveBuildOperator],\\n) -> dict[str, SelectiveBuildOperator]:\\n operators: dict[str, SelectiveBuildOperator] = {}\\n for op_name, op in list(lhs.items()) + list(rhs.items()):\\n new_op = op\\n if op_name in operators:\\n new_op = combine_operators(operators[op_name], op)\\n\\n operators[op_name] = new_op\\n\\n return operators\\n\\n\\ndef strip_operator_overload_name(op_name: str) -> str:\\n return op_name.split(\\\".\\\")[0]\\n\\n\\n\\n\\n# Be extra careful when you edit this file, because it affects AOTInductor ABI compatbility. See\\n# https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436\\n# for details.\\n#\\n# The inductor_fallback_ops list is based on the fallback ops from torch/_inductor/lowering.py.\\n# Generally speaking, it is ok to add a new op to the list, but you need to run\\n# `python torchgen/gen.py --update-aoti-c-shim` in order to regenerate C shim header files.\\n# But it is NOT ok to remove an existing fallback op from the list, since that will break\\n# some existing AOTInductor-compiled models.\\ninductor_fallback_ops = {\\n \\\"aten._adaptive_avg_pool2d_backward.default\\\",\\n \\\"aten._adaptive_avg_pool2d.default\\\",\\n \\\"aten._adaptive_avg_pool3d.default\\\",\\n \\\"aten._adaptive_avg_pool3d_backward.default\\\",\\n \\\"aten.adaptive_max_pool2d_backward.default\\\",\\n \\\"aten.adaptive_max_pool2d.default\\\",\\n \\\"aten.adaptive_max_pool3d.default\\\",\\n \\\"aten.adaptive_max_pool3d_backward.default\\\",\\n \\\"aten.addbmm.default\\\",\\n \\\"aten._addmm_activation.default\\\",\\n \\\"aten.addmm.out\\\",\\n \\\"aten.addmv.default\\\",\\n \\\"aten.angle.default\\\",\\n \\\"aten.avg_pool2d_backward.default\\\",\\n \\\"aten.avg_pool2d.default\\\",\\n \\\"aten.avg_pool3d_backward.default\\\",\\n \\\"aten.avg_pool3d.default\\\",\\n \\\"aten.bernoulli_.float\\\",\\n \\\"aten.bernoulli_.Tensor\\\",\\n \\\"aten.bmm.out\\\",\\n \\\"aten.bucketize.Tensor\\\",\\n \\\"aten.cat.default\\\",\\n \\\"aten._cdist_backward.default\\\",\\n \\\"aten._cdist_forward.default\\\",\\n \\\"aten.cholesky_inverse.default\\\",\\n \\\"aten.cholesky_solve.default\\\",\\n \\\"aten.convolution_backward.default\\\",\\n \\\"aten._cudnn_rnn.default\\\",\\n \\\"aten._cudnn_rnn_backward.default\\\",\\n \\\"aten.convolution.default\\\",\\n \\\"aten.cummax.default\\\",\\n \\\"aten.cummin.default\\\",\\n \\\"aten.cumprod.default\\\",\\n \\\"aten.cumsum.default\\\",\\n \\\"aten._efficient_attention_backward.default\\\",\\n \\\"aten._efficient_attention_forward.default\\\",\\n \\\"aten._efficientzerotensor.default\\\",\\n \\\"aten._embedding_bag.default\\\",\\n \\\"aten._embedding_bag_dense_backward.default\\\",\\n \\\"aten._embedding_bag_forward_only.default\\\",\\n \\\"aten._embedding_bag_per_sample_weights_backward.default\\\",\\n \\\"aten.exponential.default\\\",\\n \\\"aten._fft_c2c.default\\\",\\n \\\"aten._fft_r2c.default\\\",\\n \\\"aten._flash_attention_backward.default\\\",\\n \\\"aten._flash_attention_forward.default\\\",\\n \\\"aten.fractional_max_pool2d_backward.default\\\",\\n \\\"aten.fractional_max_pool2d.default\\\",\\n \\\"aten.fractional_max_pool3d.default\\\",\\n \\\"aten.fractional_max_pool3d_backward.default\\\",\\n \\\"aten._fused_moving_avg_obs_fq_helper.default\\\",\\n \\\"aten._fused_moving_avg_obs_fq_helper_functional.default\\\",\\n \\\"aten.gcd.default\\\",\\n \\\"aten.geqrf.default\\\",\\n \\\"aten.grid_sampler_2d_backward.default\\\",\\n \\\"aten.histc.default\\\",\\n \\\"aten.histogram.bin_ct\\\",\\n \\\"aten._histogramdd_bin_edges.default\\\",\\n \\\"aten._histogramdd_from_bin_cts.default\\\",\\n \\\"aten.index_put.default\\\",\\n \\\"aten.index_reduce.default\\\",\\n \\\"aten.index.Tensor\\\",\\n \\\"aten.kthvalue.default\\\",\\n \\\"aten.logcumsumexp.default\\\",\\n \\\"aten.lu_unpack.default\\\",\\n \\\"aten.masked_scatter.default\\\",\\n \\\"aten.masked_scatter_backward.default\\\",\\n \\\"aten.max_pool2d_with_indices_backward.default\\\",\\n \\\"aten.max_pool2d_with_indices.default\\\",\\n \\\"aten.max_pool3d_with_indices.default\\\",\\n \\\"aten.max_pool3d_with_indices_backward.default\\\",\\n \\\"aten.max_unpool2d.default\\\",\\n \\\"aten.max_unpool3d.default\\\",\\n \\\"aten.median.default\\\",\\n \\\"aten.mm.out\\\",\\n \\\"aten.mode.default\\\",\\n \\\"aten.mul.Scalar\\\",\\n \\\"aten.mul.Tensor\\\",\\n \\\"aten.nanmedian.default\\\",\\n \\\"aten.native_dropout.default\\\",\\n \\\"aten.normal_functional.default\\\",\\n \\\"aten.nonzero.default\\\",\\n \\\"aten.ormqr.default\\\",\\n \\\"aten._pdist_backward.default\\\",\\n \\\"aten._pdist_forward.default\\\",\\n \\\"aten.polar.default\\\",\\n \\\"aten.pow.Scalar\\\",\\n \\\"aten.pow.Tensor_Scalar\\\",\\n \\\"aten.pow.Tensor_Tensor\\\",\\n \\\"aten.rand.default\\\",\\n \\\"aten.rand.generator\\\",\\n \\\"aten.randint.default\\\",\\n \\\"aten.randint.generator\\\",\\n \\\"aten.randint.low\\\",\\n \\\"aten.randint.low_out\\\",\\n \\\"aten.randn.default\\\",\\n \\\"aten.randn.generator\\\",\\n \\\"aten.randperm.default\\\",\\n \\\"aten.repeat_interleave.Tensor\\\",\\n \\\"aten.replication_pad1d_backward.default\\\",\\n \\\"aten.replication_pad2d_backward.default\\\",\\n \\\"aten.reshape.default\\\",\\n \\\"aten.resize_.default\\\",\\n \\\"aten.resize_as_.default\\\",\\n \\\"aten._scaled_dot_product_efficient_attention_backward.default\\\",\\n \\\"aten._scaled_dot_product_efficient_attention.default\\\",\\n \\\"aten._scaled_dot_product_flash_attention_backward.default\\\",\\n \\\"aten._scaled_dot_product_flash_attention.default\\\",\\n \\\"aten._scaled_dot_product_cudnn_attention_backward.default\\\",\\n \\\"aten._scaled_dot_product_cudnn_attention.default\\\",\\n \\\"aten._scaled_dot_product_flash_attention_for_cpu_backward.default\\\",\\n \\\"aten._scaled_dot_product_flash_attention_for_cpu.default\\\",\\n \\\"aten._scaled_mm.default\\\",\\n \\\"aten.scatter_reduce.two_out\\\",\\n \\\"aten.scatter.src_out\\\",\\n \\\"aten.scatter.value_out\\\",\\n \\\"aten.searchsorted.default\\\",\\n \\\"aten._segment_reduce_backward.default\\\",\\n \\\"aten.segment_reduce.default\\\",\\n \\\"aten.slice.Tensor\\\",\\n \\\"aten.soft_margin_loss_backward.default\\\",\\n \\\"aten.sort.default\\\",\\n \\\"aten.sort.stable\\\",\\n \\\"aten._sparse_coo_tensor_with_dims_and_tensors.default\\\",\\n \\\"aten._thnn_fused_lstm_cell.default\\\",\\n \\\"aten.topk.default\\\",\\n \\\"aten._to_sparse.default\\\",\\n \\\"aten.to_sparse.default\\\",\\n \\\"aten.triangular_solve.default\\\",\\n \\\"aten._trilinear.default\\\",\\n \\\"aten.uniform.default\\\",\\n \\\"aten.upsample_bicubic2d_backward.default\\\",\\n \\\"aten.upsample_linear1d_backward.default\\\",\\n \\\"aten.upsample_trilinear3d_backward.default\\\",\\n \\\"aten.view_as_complex.default\\\",\\n \\\"aten.view_as_real.default\\\",\\n \\\"aten.view.dtype\\\",\\n \\\"aten.zeros.names\\\",\\n}\\n\\n\\n\\n\\nfrom __future__ import annotations\\n\\nimport itertools\\nfrom typing import Sequence\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.types import DispatcherSignature\\nfrom torchgen.code_template import CodeTemplate\\nfrom torchgen.context import with_native_function\\nfrom torchgen.model import Argument, NativeFunction, SchemaKind, TensorOptionsArguments\\nfrom torchgen.utils import FileManager\\n\\n\\n# Note [Manual Backend kernels]\\n# For these ops, we want to manually register to dispatch key Backend and\\n# skip codegen-ed registeration to all keys before Backend.\\n# For codegen this means:\\n# - op set below must match ops with manual_kernel_registration=True in native_functions.yaml\\n# where we skip codegen backend kernels\\n# - all ops below are part of MANUAL_AUTOGRAD to skip codegen Autograd kernel registration\\n# - all ops below are part of MANUAL_TRACER to skip codegen Tracer kernel registration\\n# Note: we still register to dispatch key Profiler for these ops, keeping it untouched for now.\\n# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp\\nMANUAL_BACKEND = {\\n \\\"options\\\",\\n \\\"data\\\",\\n \\\"set_data\\\",\\n \\\"is_leaf\\\",\\n \\\"output_nr\\\",\\n \\\"_version\\\",\\n \\\"retain_grad\\\",\\n \\\"_backward\\\",\\n \\\"requires_grad_\\\",\\n}\\n\\n# For these ops we want to skip the codegen-ed registration to both Autograd and Tracer keys.\\n# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp\\nMANUAL_AUTOGRAD_AND_TRACER = {\\n \\\"resize_\\\",\\n \\\"resize_as_\\\",\\n \\\"detach\\\",\\n \\\"detach_\\\",\\n \\\"copy_\\\",\\n \\\"_fw_primal\\\",\\n \\\"_make_dual\\\",\\n}\\n\\n# Currently MANUAL_AUTOGRAD and MANUAL_TRACER share the same set of ops:\\n# union(MANUAL_BACKEND, MANUAL_AUTOGRAD_AND_TRACER)\\n# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp\\nMANUAL_AUTOGRAD = MANUAL_TRACER = MANUAL_BACKEND | MANUAL_AUTOGRAD_AND_TRACER\\n\\n# These functions we don't want to record for tracing, because we always want\\n# to trace their constituent parts. This is a temporary hack in lieue\\n# of proper scopes, where subsequent compilation passes can ask for the unfolding\\n# on demand. Only concrete ATen methods can be disabled this way; it will have\\n# NO EFFECT otherwise.\\nDONT_RECORD_TRACE = {\\n \\\"convolution\\\",\\n \\\"conv1d\\\",\\n \\\"conv2d\\\",\\n \\\"conv3d\\\",\\n \\\"conv_transpose1d\\\",\\n \\\"conv_transpose2d\\\",\\n \\\"conv_transpose3d\\\",\\n \\\"lstm_cell\\\",\\n \\\"gru_cell\\\",\\n \\\"rnn_tanh_cell\\\",\\n \\\"rnn_relu_cell\\\",\\n # FIXME: figure out a better way when we support sparse tensors in jit\\n \\\"_coalesced\\\",\\n}\\n\\n\\ndef should_trace(f: NativeFunction) -> bool:\\n # Operations involving Storage or Type are not traceable at the moment\\n if any(\\n str(arg.type) in {\\\"Storage\\\", \\\"Type\\\", \\\"ConstQuantizerPtr\\\"}\\n for arg in f.func.schema_order_arguments()\\n ):\\n return False\\n # We can't trace functions which don't have any Tensor or TensorList returns\\n if not any(r.type.is_tensor_like() for r in f.func.returns):\\n return False\\n return f.func.name.name.base not in DONT_RECORD_TRACE\\n\\n\\nSELECT = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n\\nif (${cond}) {\\n ${true}\\n} else {\\n ${false}\\n}\\n\\\"\\\"\\\"\\n)\\n\\nOP_NAME = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nop_name = c10::Symbol::fromQualString(\\\"aten::${trace_name}\\\");\\n\\\"\\\"\\\"\\n)\\n\\n# These functions have their names recorded under trace renamed,\\nRENAME_TRACE = {\\n \\\"zero\\\": \\\"zeros_like\\\", # replacing aten::zero_ with aten::zeros_like\\n \\\"fill\\\": \\\"full_like\\\", # replacing aten::fill_ with aten::full_like\\n}\\n\\n\\ndef format_trace_op_name(f: NativeFunction) -> str:\\n # TODO: byte-for-byte compatible with old codegen behavior - should clean up\\n if (\\n f.func.kind() in (SchemaKind.functional, SchemaKind.out)\\n or f.func.name.name.dunder_method\\n ):\\n # special case for *_out functions: the in-place and out-of-place ops\\n # are overloaded with the same name in the JIT\\n trace_name = str(f.func.name.name)\\n trace_name = RENAME_TRACE.get(trace_name, trace_name)\\n return OP_NAME.substitute(trace_name=trace_name)\\n\\n # otherwise, this is an in-place op and we need to emit both in- and\\n # out-of-place versions\\n outplace_trace_name = f.func.name.name.base\\n inplace_trace_name = cpp.name(f.func)\\n outplace_trace_name = RENAME_TRACE.get(outplace_trace_name, outplace_trace_name)\\n inplace_trace_name = RENAME_TRACE.get(inplace_trace_name, inplace_trace_name)\\n\\n return SELECT.substitute(\\n cond=\\\"tracer_state->force_outplace\\\",\\n true=OP_NAME.substitute(trace_name=outplace_trace_name),\\n false=OP_NAME.substitute(trace_name=inplace_trace_name),\\n )\\n\\n\\nADD_TRACE_INPUT = CodeTemplate(\\\"\\\"\\\"jit::tracer::addInputs(node, \\\"${name}\\\", ${input});\\\"\\\"\\\")\\n\\n\\ndef format_trace_inputs(f: NativeFunction) -> str:\\n def dispatch_trace_input(arg: Argument | TensorOptionsArguments) -> Sequence[str]:\\n if isinstance(arg, TensorOptionsArguments):\\n name = \\\"options\\\"\\n return [\\n ADD_TRACE_INPUT.substitute(\\n name=name, input=\\\"c10::optTypeMetaToScalarType(options.dtype_opt())\\\"\\n ),\\n ADD_TRACE_INPUT.substitute(name=name, input=\\\"options.layout()\\\"),\\n ADD_TRACE_INPUT.substitute(name=name, input=\\\"options.device()\\\"),\\n ADD_TRACE_INPUT.substitute(name=name, input=\\\"options.pinned_memory()\\\"),\\n ]\\n else:\\n name = arg.name\\n if str(arg.type) == \\\"Tensor?[]\\\":\\n return [f'jit::tracer::addInputs(node, \\\"{name}\\\", {name});']\\n else:\\n return [ADD_TRACE_INPUT.substitute(name=name, input=name)]\\n\\n args: list[Argument | TensorOptionsArguments] = list(\\n f.func.schema_order_arguments()\\n )\\n\\n if f.func.is_out_fn():\\n # *_out functions take the result as a separate argument, but we don't want to\\n # trace that argument directly. Instead, we trace its TensorOptions.\\n # So first, we need to remove the out argument from the list of arguments to trace.\\n num_out_args = len(f.func.arguments.out)\\n args = args[:-num_out_args]\\n\\n trace_inputs = itertools.chain.from_iterable(\\n dispatch_trace_input(arg) for arg in args\\n )\\n\\n if f.func.is_out_fn():\\n # for *_out functions, handle the result argument differently for inplace/outplace.\\n # For inplace: just add the input to the end to confirm with the JIT schema\\n inplace = [\\n ADD_TRACE_INPUT.substitute(\\n name=f.func.arguments.out[i].name, input=f.func.arguments.out[i].name\\n )\\n for i in range(num_out_args)\\n ]\\n\\n # for outplace: do nothing, except if the function is a factory.\\n # Factories are a bit special because their out-of-place overloads\\n # take an extra TensorOptions argument, which is missing in the _out function\\n has_tensor_return = any(r.type.is_tensor_like() for r in f.func.returns)\\n has_tensor_input_arg = any(\\n a.type.is_tensor_like() for a in f.func.arguments.flat_non_out\\n )\\n is_factory_method = f.category_override == \\\"factory\\\" or (\\n has_tensor_return and not has_tensor_input_arg\\n )\\n\\n # HACK: preserve old codegen behavior - the old codegen set the `is_factory_method`\\n # flag for the whole family of ops with the same basename if any of them is a\\n # factory method. For most cases the whole family of ops are indeed all factory\\n # method - 'normal' is the only exception. So we handle it specially here to avoid\\n # cloning the old logic.\\n if f.func.name.name.base == \\\"normal\\\":\\n is_factory_method = True\\n\\n if is_factory_method:\\n outplace = [\\n ADD_TRACE_INPUT.substitute(\\n name=\\\"out\\\",\\n input=\\\"c10::optTypeMetaToScalarType(out.options().dtype_opt())\\\",\\n ),\\n ADD_TRACE_INPUT.substitute(name=\\\"out\\\", input=\\\"out.options().layout()\\\"),\\n ADD_TRACE_INPUT.substitute(name=\\\"out\\\", input=\\\"out.options().device()\\\"),\\n ADD_TRACE_INPUT.substitute(\\n name=\\\"out\\\", input=\\\"out.options().pinned_memory()\\\"\\n ),\\n ]\\n else:\\n outplace = []\\n\\n trace_inputs = itertools.chain(\\n trace_inputs,\\n [\\n SELECT.substitute(\\n cond=\\\"tracer_state->force_outplace\\\",\\n true=\\\"\\\\n\\\".join(outplace),\\n false=\\\"\\\\n\\\".join(inplace),\\n )\\n ],\\n )\\n\\n return \\\"\\\\n\\\".join(trace_inputs)\\n\\n\\n# `torch.jit.trace` have undocumented keyword argument `_force_outplace`,\\n# which force jit to replace functions with outplace variants (for\\n# example `aten::add_` becomes `aten::add`).\\n#\\n# This replacement implemented in-place with minimum modifications of\\n# arguments stack (as it assumes that outplace call has the same arguments\\n# as inplace version).\\n#\\n# However there are no such substitutions available for `aten::fill_`\\n# and `aten::zero_` operators, as we never implemented `aten::fill`\\n# and `aten::zero`. So jit tracing hack replacing `aten::zero_` with\\n# `aten::zeros_like` and replacing `aten::fill_` with `aten::full_like`.\\n#\\n# But as they potentially can have different arguments, we also have\\n# to hack into the stack and add missing ones.\\n#\\n# A possible alternative would be:\\n#\\n# - Add `aten::fill` and `aten::zero`\\n#\\n# - Or keep `aten::zeros_like` arguments aligned with `aten::zero_`\\n# arguments (inside of the `native_functions.yaml`)\\nRENAME_TRACE_ADD_ARGS = {\\n \\\"fill\\\": \\\"\\\"\\\"\\\\\\n jit::tracer::addInputs(node, \\\"options\\\", ::std::optional());\\n jit::tracer::addInputs(node, \\\"options\\\", layout_or_default(::std::nullopt));\\n jit::tracer::addInputs(node, \\\"options\\\", device_or_default(::std::nullopt));\\n jit::tracer::addInputs(node, \\\"options\\\", pinned_memory_or_default(::std::nullopt));\\n ::std::optional memory_format = c10::MemoryFormat::Preserve;\\n jit::tracer::addInputs(node, \\\"memory_format\\\", memory_format);\\n\\\"\\\"\\\",\\n \\\"zero\\\": \\\"\\\"\\\"\\\\\\n jit::tracer::addInputs(node, \\\"options\\\", ::std::optional());\\n jit::tracer::addInputs(node, \\\"options\\\", layout_or_default(::std::nullopt));\\n jit::tracer::addInputs(node, \\\"options\\\", device_or_default(::std::nullopt));\\n jit::tracer::addInputs(node, \\\"options\\\", pinned_memory_or_default(::std::nullopt));\\n ::std::optional memory_format = c10::MemoryFormat::Preserve;\\n jit::tracer::addInputs(node, \\\"memory_format\\\", memory_format);\\n\\\"\\\"\\\",\\n}\\n\\nINPLACE_GUARD = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\njit::tracer::ensureUniqueIfOutOfPlaced(\\\"${name}\\\", ${mutable_input});\\n\\\"\\\"\\\"\\n)\\n\\nPRE_RECORD_TRACE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\ntorch::jit::Node* node = nullptr;\\nstd::shared_ptr tracer_state;\\nif (jit::tracer::isTracing()) {\\n tracer_state = jit::tracer::getTracingState();\\n at::Symbol op_name;\\n ${set_op_name}\\n node = tracer_state->createNode(op_name, /*num_outputs=*/0);\\n jit::tracer::recordSourceLocation(node);\\n ${add_trace_inputs}\\n tracer_state->insertNode(node);\\n ${inplace_guard}\\n jit::tracer::setTracingState(nullptr);\\n}\\n\\\"\\\"\\\"\\n)\\n\\n\\ndef format_prerecord_trace(f: NativeFunction) -> str:\\n if not should_trace(f):\\n return \\\"\\\"\\n\\n # TODO: clean up old codegen behavior\\n is_inplace = (\\n f.func.kind() in (SchemaKind.inplace, SchemaKind.out)\\n and not f.func.name.name.dunder_method\\n )\\n add_args = (\\n RENAME_TRACE_ADD_ARGS.get(f.func.name.name.base, \\\"\\\") if is_inplace else \\\"\\\"\\n )\\n additional_inputs = (\\n SELECT.substitute(\\n cond=\\\"tracer_state->force_outplace\\\",\\n true=add_args,\\n false=\\\"\\\",\\n )\\n if add_args\\n else \\\"\\\"\\n )\\n\\n return PRE_RECORD_TRACE.substitute(\\n set_op_name=format_trace_op_name(f),\\n add_trace_inputs=format_trace_inputs(f) + additional_inputs,\\n inplace_guard=INPLACE_GUARD.substitute(\\n name=cpp.name(f.func),\\n mutable_input=f.func.arguments.out[0].name\\n if f.func.arguments.out\\n else \\\"self\\\",\\n )\\n if is_inplace\\n else \\\"\\\",\\n )\\n\\n\\nPOST_RECORD_TRACE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (tracer_state) {\\n jit::tracer::setTracingState(std::move(tracer_state));\\n ${add_trace_outputs}\\n}\\n\\\"\\\"\\\"\\n)\\n\\n\\ndef format_postrecord_trace(f: NativeFunction) -> str:\\n if not should_trace(f):\\n return \\\"\\\"\\n\\n # For outplacing ops, *_out overloads require special handling to move the\\n # output *argument* to a return value\\n if f.func.is_out_fn():\\n output_names_outplace = [arg.name for arg in f.func.arguments.out]\\n output_names_inplace = cpp.return_names(f)\\n\\n # Code size optimization: the common case is that the return value is\\n # the same for both variants\\n if output_names_outplace == output_names_inplace:\\n outputs = [\\n f\\\"jit::tracer::addOutput(node, {n});\\\" for n in output_names_outplace\\n ]\\n return POST_RECORD_TRACE.substitute(add_trace_outputs=outputs)\\n\\n selection = SELECT.substitute(\\n cond=\\\"force_outplace\\\",\\n true=\\\"\\\\n\\\".join(\\n f\\\"jit::tracer::addOutput(node, {n});\\\" for n in output_names_outplace\\n ),\\n false=\\\"\\\\n\\\".join(\\n f\\\"jit::tracer::addOutput(node, {n});\\\" for n in output_names_inplace\\n ),\\n )\\n return POST_RECORD_TRACE.substitute(add_trace_outputs=selection)\\n else:\\n output_names = cpp.return_names(f)\\n outputs = [f\\\"jit::tracer::addOutput(node, {n});\\\" for n in output_names]\\n return POST_RECORD_TRACE.substitute(add_trace_outputs=outputs)\\n\\n\\ndef tie_return_values(f: NativeFunction) -> str:\\n if len(f.func.returns) == 1:\\n return f'auto {f.func.returns[0].name or \\\"result\\\"}'\\n names = cpp.return_names(f)\\n return f'auto [{\\\", \\\".join(names)}]'\\n\\n\\ndef get_return_value(f: NativeFunction) -> str:\\n names = cpp.return_names(f)\\n if len(f.func.returns) == 1:\\n return names[0]\\n if f.func.kind() == SchemaKind.out:\\n return f'std::forward_as_tuple({\\\", \\\".join(names)})'\\n else:\\n moved = \\\", \\\".join(f\\\"std::move({name})\\\" for name in names)\\n return f\\\"std::make_tuple({moved})\\\"\\n\\n\\nTRACE_DISPATCH = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${assign_return_values}at::_ops::${unambiguous_name}::redispatch(${unpacked_args});\\\"\\\"\\\"\\n)\\n\\n\\ndef emit_trace_body(f: NativeFunction) -> list[str]:\\n trace_body: list[str] = []\\n\\n trace_body.append(format_prerecord_trace(f))\\n\\n dispatcher_sig = DispatcherSignature.from_schema(f.func)\\n dispatcher_exprs = dispatcher_sig.exprs()\\n\\n # code-generated tracing kernels plumb and recompute dispatch keys directly through the kernel for performance.\\n # See Note [Plumbing Keys Through The Dispatcher] for details.\\n dispatch_key_set = \\\"ks & c10::DispatchKeySet(c10::DispatchKeySet::FULL_AFTER, c10::DispatchKey::Tracer)\\\"\\n redispatch_args = \\\", \\\".join([dispatch_key_set] + [a.expr for a in dispatcher_exprs])\\n\\n assign_return_values = (\\n f\\\"{tie_return_values(f)} = \\\"\\n if f.func.kind() in [SchemaKind.functional, SchemaKind.mutable]\\n and f.func.returns\\n else \\\"\\\"\\n )\\n\\n # Note that this calls the slow, dispatching variants of manual_cpp_binding ops.\\n # We could probably work harder to ensure that the fast variants are\\n # called instead, but the perf benefit would be minimal.\\n trace_body.append(\\n TRACE_DISPATCH.substitute(\\n assign_return_values=assign_return_values,\\n unambiguous_name=f.func.name.unambiguous_name(),\\n unpacked_args=redispatch_args,\\n )\\n )\\n\\n trace_body.append(format_postrecord_trace(f))\\n if f.func.returns:\\n trace_body.append(f\\\"return {get_return_value(f)};\\\")\\n return trace_body\\n\\n\\nMETHOD_DEFINITION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${return_type} ${type_wrapper_name}(${formals}) {\\n ${type_definition_body}\\n}\\n\\\"\\\"\\\"\\n)\\n\\n\\ndef type_wrapper_name(f: NativeFunction, key: str = \\\"Default\\\") -> str:\\n if f.func.name.overload_name:\\n name = f\\\"{cpp.name(f.func)}_{f.func.name.overload_name}\\\"\\n else:\\n name = cpp.name(f.func)\\n\\n # The key argument is only used in gen_variable_type where we need fns per autograd dispatch key.\\n # In gen_trace_type and gen_inplace_view_type where only one fn per native_fn must be generated,\\n # the key argument should not be passed.\\n # We do not append key if it is Default so that generated functions from\\n # before per-dispatch-key derivatives were added retain the same names.\\n if key != \\\"Default\\\":\\n name = name + f\\\"_{key}\\\"\\n return name\\n\\n\\n@with_native_function\\ndef method_definition(f: NativeFunction) -> str:\\n assert cpp.name(f.func) not in MANUAL_TRACER\\n\\n formals = \\\", \\\".join(\\n # code-generated tracing kernels plumb and recompute dispatch keys directly through the kernel for performance.\\n # See Note [Plumbing Keys Through The Dispatcher] for details.\\n [\\\"c10::DispatchKeySet ks\\\"]\\n + [\\n f'{cpp.argument_type(a, binds=\\\"__placeholder__\\\", symint=True).cpp_type()} {a.name}'\\n for a in f.func.schema_order_arguments()\\n ]\\n )\\n\\n return METHOD_DEFINITION.substitute(\\n return_type=cpp.returns_type(f.func.returns, symint=True).cpp_type(),\\n type_wrapper_name=type_wrapper_name(f),\\n formals=formals,\\n type_definition_body=emit_trace_body(f),\\n )\\n\\n\\nWRAPPER_REGISTRATION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nm.impl(\\\"${name}\\\",\\n TORCH_FN(${class_type}::${type_wrapper_name})\\n);\\n\\\"\\\"\\\"\\n)\\n\\n\\n@with_native_function\\ndef method_registration(f: NativeFunction) -> str:\\n assert cpp.name(f.func) not in MANUAL_TRACER\\n\\n return WRAPPER_REGISTRATION.substitute(\\n name=f.func.name,\\n type_wrapper_name=type_wrapper_name(f),\\n class_type=\\\"TraceType\\\",\\n )\\n\\n\\ndef gen_trace_type_func(fn: NativeFunction) -> dict[str, list[str]]:\\n return {\\n \\\"ops_headers\\\": [f\\\"#include \\\"],\\n \\\"trace_method_definitions\\\": [method_definition(fn)],\\n \\\"trace_wrapper_registrations\\\": [method_registration(fn)],\\n }\\n\\n\\ndef gen_trace_type(\\n out: str, native_functions: list[NativeFunction], template_path: str\\n) -> None:\\n # NOTE: see Note [Sharded File] at the top of the VariableType.cpp\\n # template regarding sharding of the generated files.\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n fm.write_sharded(\\n \\\"TraceType.cpp\\\",\\n [fn for fn in native_functions if cpp.name(fn.func) not in MANUAL_TRACER],\\n key_fn=lambda fn: fn.root_name,\\n base_env={\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/TraceType.cpp\\\",\\n },\\n env_callable=gen_trace_type_func,\\n num_shards=5,\\n sharded_keys={\\n \\\"ops_headers\\\",\\n \\\"trace_method_definitions\\\",\\n \\\"trace_wrapper_registrations\\\",\\n },\\n )\\n\\n\\n# Parses derivatives.yaml into autograd functions\\n#\\n# Each autograd function is represented by `DifferentiabilityInfo` containing\\n# a list of `Derivative`. See `torchgen.api.autograd` for the data models.\\n\\nfrom __future__ import annotations\\n\\nimport re\\nfrom collections import defaultdict\\nfrom typing import Any, Counter, Dict, Sequence, Set, Tuple\\n\\nimport yaml\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.autograd import (\\n Derivative,\\n DifferentiabilityInfo,\\n ForwardDerivative,\\n SavedAttribute,\\n)\\nfrom torchgen.api.types import (\\n BaseCType,\\n Binding,\\n boolT,\\n CppSignatureGroup,\\n layoutT,\\n longT,\\n NamedCType,\\n OptionalCType,\\n scalarTypeT,\\n SpecialArgName,\\n stringT,\\n symIntArrayRefT,\\n SymIntT,\\n tensorGeometryT,\\n tensorOptionsT,\\n typeAndSizeT,\\n VectorCType,\\n)\\nfrom torchgen.context import with_native_function\\nfrom torchgen.gen import get_grouped_by_view_native_functions, parse_native_yaml\\nfrom torchgen.model import (\\n AUTOGRAD_KEYS,\\n FunctionSchema,\\n NativeFunction,\\n NativeFunctionsViewGroup,\\n OperatorName,\\n SchemaKind,\\n Type,\\n Variant,\\n)\\nfrom torchgen.utils import concatMap, IDENT_REGEX, split_name_params\\nfrom torchgen.yaml_utils import YamlLoader\\n\\n\\nDerivativeRet = Tuple[Dict[FunctionSchema, Dict[str, DifferentiabilityInfo]], Set[str]]\\n\\n_GLOBAL_LOAD_DERIVATIVE_CACHE: dict[tuple[str, str], DerivativeRet] = {}\\n\\n_VALID_AUTOGRAD_KEYS = set(AUTOGRAD_KEYS)\\n\\n\\n# This function directly adds per-dispatchkey derivative entries for {view}_copy variants of each view op.\\n# Since every {view} and {view}_copy op shares the same derivative formula,\\n# we generate them here instead of duplicating them in the yaml.\\n# See Note [Codegen'd {view}_copy Operators]\\ndef add_view_copy_derivatives(\\n infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]],\\n view_groups: list[NativeFunctionsViewGroup],\\n) -> None:\\n # Get the map from each view op's name to its corresponding view group\\n view_name_to_group: dict[OperatorName, NativeFunctionsViewGroup] = {\\n g.view.func.name: g for g in view_groups\\n }\\n\\n view_infos = {}\\n\\n for info_dispatch_dict in infos.values():\\n # maybe_view_group only needs to be calculated once per info_dispatch_dict\\n maybe_view_group = None\\n view_copy_differentiability_infos = {}\\n for dispatch_key, info in info_dispatch_dict.items():\\n maybe_view_group = view_name_to_group.get(info.func.func.name, None)\\n if maybe_view_group is not None and maybe_view_group.view_copy is not None:\\n view_copy_info = info.create_view_copy_from_view_derivative(\\n maybe_view_group\\n )\\n if view_copy_info is not None:\\n fn_schema = view_copy_info.func.func\\n view_copy_differentiability_infos[dispatch_key] = view_copy_info\\n else:\\n break\\n # prefer manually-defined derivatives if any\\n if len(view_copy_differentiability_infos) > 0 and fn_schema not in infos:\\n assert fn_schema is not None\\n view_infos[fn_schema] = view_copy_differentiability_infos\\n\\n infos.update(view_infos)\\n\\n\\ndef load_derivatives(\\n derivatives_yaml_path: str, native_yaml_path: str, tags_yaml_path: str\\n) -> DerivativeRet:\\n # Do some caching as this is a deterministic function\\n global _GLOBAL_LOAD_DERIVATIVE_CACHE\\n key = (derivatives_yaml_path, native_yaml_path)\\n if key not in _GLOBAL_LOAD_DERIVATIVE_CACHE:\\n with open(derivatives_yaml_path) as f:\\n definitions = yaml.load(f, Loader=YamlLoader)\\n\\n funcs = parse_native_yaml(native_yaml_path, tags_yaml_path).native_functions\\n # From the parsed native functions, separate out the (generated) view_copy functions,\\n # so we can generate derivatives for them separately.\\n native_functions_with_view_groups = get_grouped_by_view_native_functions(funcs)\\n native_functions = concatMap(\\n lambda g: [g]\\n if isinstance(g, NativeFunction)\\n else list(g.functions(include_copy=True)),\\n native_functions_with_view_groups,\\n )\\n view_groups = [\\n g\\n for g in native_functions_with_view_groups\\n if isinstance(g, NativeFunctionsViewGroup)\\n ]\\n\\n # What's the difference between function schema v.s. signature?\\n # function schema is the complete declaration including mutability annotation / default value and etc.\\n # signature is the canonical schema for a group of functions (in-place/out/functional variants)\\n # that are semantically related.\\n functions_by_signature: dict[\\n FunctionSchema, list[NativeFunction]\\n ] = defaultdict(list)\\n functions_by_schema: dict[str, NativeFunction] = {}\\n for function in native_functions:\\n functions_by_signature[function.func.signature()].append(function)\\n assert str(function.func) not in functions_by_schema\\n functions_by_schema[str(function.func)] = function\\n\\n # Keep track of how many of which ops we've seen so we can\\n # disambiguate them with a numeric suffix.\\n op_counter = Counter[str]()\\n\\n # infos is a dict that maps FunctionSchema -> a dict of per dispatch key DifferentiabilityInfos\\n # this is useful because in tools/autograd/gen_autograd.py:match_differentiability_info\\n # we ultimately need to categorize the DifferentiabilityInfos by FunctionSchema\\n infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]] = {}\\n used_dispatch_keys: set[str] = set()\\n for defn_dict in definitions:\\n # Ensure that the old derivatives.yaml schema with no dispatch key can be loaded.\\n if \\\"dispatch\\\" not in defn_dict:\\n specification = defn_dict.pop(\\\"name\\\")\\n output_differentiability = defn_dict.pop(\\n \\\"output_differentiability\\\", None\\n )\\n defn_dict = {\\\"name\\\": specification, \\\"dispatch\\\": {\\\"Default\\\": defn_dict}}\\n if output_differentiability:\\n defn_dict[\\\"output_differentiability\\\"] = output_differentiability\\n name, per_dispatch_diffinfos = create_differentiability_info(\\n defn_dict,\\n functions_by_signature,\\n functions_by_schema,\\n op_counter,\\n used_dispatch_keys,\\n )\\n infos[name] = per_dispatch_diffinfos\\n\\n add_view_copy_derivatives(infos, view_groups)\\n\\n # cache both loaded infos as well a a set of all the dispatch_keys/aliases\\n # that appear in derivatives.yaml. used_dispatch_keys is useful for generating\\n # VariableType.cpp where we need a TORCH_LIBRARY_IMPL for every autograd dispatch key used\\n _GLOBAL_LOAD_DERIVATIVE_CACHE[key] = infos, used_dispatch_keys\\n\\n return _GLOBAL_LOAD_DERIVATIVE_CACHE[key]\\n\\n\\n# TODO: Why is this going through CppSignatureGroup, that doesn't make sense...\\n@with_native_function\\ndef cpp_arguments(f: NativeFunction) -> Sequence[Binding]:\\n sigs = CppSignatureGroup.from_native_function(f, method=False)\\n if sigs.symint_signature is not None:\\n return sigs.symint_signature.arguments()\\n else:\\n return sigs.signature.arguments()\\n\\n\\ndef create_derivative(\\n f: NativeFunction,\\n formula: str,\\n var_names: tuple[str, ...],\\n available_named_gradients: Sequence[str],\\n) -> Derivative:\\n original_formula = formula\\n arguments: list[NamedCType] = [\\n a.nctype.remove_const_ref() for a in cpp_arguments(f)\\n ]\\n\\n return_names = tuple(n if n != \\\"self\\\" else \\\"result\\\" for n in cpp.return_names(f))\\n return_types = tuple(\\n cpp.return_type(r, symint=True).remove_const_ref() for r in f.func.returns\\n )\\n\\n named_returns = [\\n NamedCType(name, type) for name, type in zip(return_names, return_types)\\n ]\\n\\n formula, saved_inputs = saved_variables(formula, arguments, var_names)\\n formula, saved_outputs = saved_variables(formula, named_returns, var_names)\\n\\n used_named_gradients = {\\n name\\n for name in available_named_gradients\\n if re.search(IDENT_REGEX.format(name), formula)\\n }\\n\\n # Check that the referenced derivatives in the formula are in bounds\\n for i in used_gradient_indices(formula):\\n if i >= len(f.func.returns):\\n raise RuntimeError(\\n f\\\"Out of bounds grads access: derivative formula for {cpp.name(f.func)} \\\"\\n f\\\"used grads[{i}], but the forward only returns {len(f.func.returns)} outputs.\\\"\\n )\\n\\n return Derivative(\\n formula=formula,\\n original_formula=original_formula,\\n var_names=var_names,\\n saved_inputs=saved_inputs,\\n saved_outputs=saved_outputs,\\n named_gradients=used_named_gradients,\\n )\\n\\n\\ndef create_forward_derivative(\\n f: NativeFunction, formula: str, names: tuple[str, ...]\\n) -> ForwardDerivative:\\n var_names = names\\n var_types: tuple[Type, ...] | None = None\\n for r in f.func.returns:\\n if r.name in var_names:\\n if var_types is None:\\n var_types = ()\\n var_types = var_types + (r.type,)\\n\\n # Handle default return names\\n if var_types is None:\\n if var_names == (\\\"result\\\",):\\n assert len(f.func.returns) == 1\\n var_types = (f.func.returns[0].type,)\\n else:\\n for var_name in var_names:\\n res = re.findall(r\\\"^result(\\\\d+)$\\\", var_name)\\n if len(res) == 1:\\n if var_types is None:\\n var_types = ()\\n arg_idx = int(res[0])\\n var_types = var_types + (f.func.returns[arg_idx].type,)\\n\\n assert var_types is not None, \\\"No matching output for forward derivative definition\\\"\\n return ForwardDerivative(\\n formula=formula,\\n var_names=var_names,\\n var_types=var_types,\\n required_inputs_fw_grad=None,\\n required_inputs_primal=None,\\n required_original_self_value=False,\\n is_reusing_outplace_formula=False,\\n )\\n\\n\\ndef postprocess_forward_derivatives(\\n f: NativeFunction,\\n defn_name: str,\\n all_arg_names: list[str],\\n derivatives: list[Derivative],\\n forward_derivatives: list[ForwardDerivative],\\n args_with_derivatives: Sequence[Binding],\\n) -> list[ForwardDerivative]:\\n def find_required_inputs(formula: str, postfix: str) -> tuple[str, ...]:\\n is_foreach = f.func.name.name.base.startswith(\\\"_foreach_\\\")\\n required_inputs = set()\\n for arg in args_with_derivatives:\\n if (\\n arg.type in (\\\"at::TensorList\\\", \\\"const at::ITensorListRef &\\\")\\n and not is_foreach\\n ):\\n # The functions taking TensorList handle everything internally\\n continue\\n arg_name = arg.name\\n\\n found = re.search(IDENT_REGEX.format(arg_name), formula)\\n if found:\\n raise RuntimeError(\\n f\\\"The forward formula for {defn_name} is using the base name of the {arg_name} \\\"\\n f\\\"argument which is ambiguous. You should use {arg_name}_p to access the primal \\\"\\n f\\\"value and {arg_name}_t to access the tangent.\\\"\\n )\\n\\n found = re.search(IDENT_REGEX.format(arg_name + postfix), formula)\\n if found:\\n required_inputs.add(arg_name)\\n\\n return tuple(required_inputs)\\n\\n updated_derivatives: list[ForwardDerivative] = []\\n\\n for defn in forward_derivatives:\\n formula = defn.formula\\n required_inputs_tangent = find_required_inputs(formula, \\\"_t\\\")\\n if formula == \\\"auto_element_wise\\\":\\n assert (\\n f.func.kind() != SchemaKind.inplace\\n ), f\\\"Cannot use auto_element_wise with {f.func.name} because it is an in-place variant\\\"\\n if (\\n (not len(args_with_derivatives) == 1)\\n or len(forward_derivatives) > 1\\n or len(forward_derivatives[0].var_names) > 1\\n ):\\n raise RuntimeError(\\n f\\\"Derivative definition of {defn_name} in derivatives.yaml defines the \\\"\\n \\\"forward definition of gradient as element_wise but this only \\\"\\n \\\"works for functions with a single differentiable input and a \\\"\\n \\\"single differentiable output.\\\"\\n )\\n if not len(derivatives) == 1:\\n raise RuntimeError(\\n f\\\"Derivative definition of {defn_name} in derivatives.yaml defines the \\\"\\n \\\"forward definition of gradient as element_wise but it does not \\\"\\n \\\"defines the gradient formula for its argument which is required.\\\"\\n )\\n # This transformation is based on the observation that for element-wise functions, the Jacobian\\n # matrix is diagonal and thus doing J * v is the same as (v^T J)^T (in practice, we ignore the transpositions)\\n # For the complex case, we use hermitian transpose and get (v.conj() J).conj()\\n # So here we are going to re-use the backward formula and replace two things:\\n # 1) all occurrences of \\\"grad\\\" with \\\"foo_t.conj()\\\", where foo is the name of the unique differentiable input.\\n # 2) all usage of an original input \\\"foo\\\" with its primal value \\\"foo_p\\\".\\n # 3) conjugate the final result\\n # For example, for abs, the backward formula is:\\n # grad * self.sgn()\\n # And this function generates a forward formula that is:\\n # (self_t.conj() * self_p.sgn()).conj()\\n\\n backward_formula = derivatives[0].original_formula\\n input_name = args_with_derivatives[0].name\\n\\n # Do replacement 1) of the grad\\n def repl(m: Any) -> str:\\n return f\\\"{m.group(1)}{input_name}_t.conj(){m.group(2)}\\\"\\n\\n fw_formula = re.sub(IDENT_REGEX.format(\\\"grad\\\"), repl, backward_formula)\\n\\n # Do replacement 2) of the input variables\\n for arg in args_with_derivatives:\\n arg_name = arg.name\\n\\n def repl(m: Any) -> str:\\n return f\\\"{m.group(1)}{arg_name}_p{m.group(2)}\\\"\\n\\n fw_formula = re.sub(IDENT_REGEX.format(arg_name), repl, fw_formula)\\n\\n # Do the final conjugate 3)\\n fw_formula = f\\\"({fw_formula}).conj()\\\"\\n\\n # Since there is a single differentiable inputs and we necessarily need its tangent we can\\n # simply require all differentiable input's tangent.\\n required_inputs_tangent = tuple(all_arg_names)\\n formula = fw_formula\\n elif formula == \\\"auto_linear\\\":\\n if (\\n len(forward_derivatives) > 1\\n or len(forward_derivatives[0].var_names) > 1\\n ):\\n raise RuntimeError(\\n f\\\"Derivative definition of {defn_name} in derivatives.yaml defines the \\\"\\n \\\"forward definition of gradient as linear but this only works \\\"\\n \\\"for functions with a single differentiable output.\\\"\\n )\\n # This transformation is based on the observation that linear functions can be written as:\\n # y = f(x) = A * x\\n # For some matrix A and the Jacobian of the function f is also A.\\n # So doing J * v = A * v = f(v).\\n # Hence to do the jvp, we simply need to evaluate the function at the point v instead of x.\\n # We do this by calling the forward again by replacing any occurrence of the differentiable\\n # input \\\"foo\\\" by it's tangent \\\"foo_t\\\".\\n # Note that multiple inputs are not a problem as long as the function is truly linear wrt to\\n # the vector where all the differentiable inputs are stacked.\\n\\n diff_arg_names = [arg.name for arg in args_with_derivatives]\\n assert len(diff_arg_names) > 0\\n\\n # Do replacement of input variables\\n new_args = []\\n for arg_name in all_arg_names:\\n if arg_name in diff_arg_names:\\n arg_name = arg_name + \\\"_t\\\"\\n new_args.append(arg_name)\\n\\n # TODO we are trolling\\n if f.func.has_symint():\\n defn_name += \\\"_symint\\\"\\n\\n # Call into the forward again. We need two cases here to handle both Tensor methods and at:: functions.\\n if Variant.function in f.variants:\\n fw_formula = f\\\"at::{defn_name}({', '.join(new_args)})\\\"\\n else:\\n assert Variant.method in f.variants\\n fw_formula = f\\\"{new_args[0]}.{defn_name}({', '.join(new_args[1:])})\\\"\\n\\n # All of the input tangents are always used so all of them are required here.\\n required_inputs_tangent = tuple(diff_arg_names)\\n formula = fw_formula\\n\\n # At this point, the formula is final and is not modified anymore.\\n\\n # During forward formula, we use the primal instead of the input Tensors.\\n # This call inspects the formula to find for which input's primal are used.\\n required_inputs_primal = find_required_inputs(formula, \\\"_p\\\")\\n\\n updated_derivatives.append(\\n ForwardDerivative(\\n formula=formula,\\n var_names=defn.var_names,\\n var_types=defn.var_types,\\n required_inputs_fw_grad=required_inputs_tangent,\\n required_inputs_primal=required_inputs_primal,\\n required_original_self_value=False,\\n is_reusing_outplace_formula=False,\\n )\\n )\\n\\n return updated_derivatives\\n\\n\\ndef is_forward_derivative_definition(\\n all_arg_names: list[str], names: tuple[str, ...]\\n) -> bool:\\n for name in names:\\n return name not in all_arg_names\\n raise RuntimeError(\\\"Expected `names` to be non-empty\\\")\\n\\n\\ndef create_differentiability_info(\\n defn_dict: dict[Any, Any],\\n functions_by_signature: dict[FunctionSchema, list[NativeFunction]],\\n functions_by_schema: dict[str, NativeFunction],\\n op_counter: Counter[str],\\n used_dispatch_keys: set[str],\\n) -> tuple[FunctionSchema, dict[str, DifferentiabilityInfo]]:\\n \\\"\\\"\\\"Processes a single entry `defn` in derivatives.yaml\\\"\\\"\\\"\\n\\n def canonical_function(\\n functions: Sequence[NativeFunction], name: str\\n ) -> NativeFunction:\\n for f in functions:\\n if (\\n not f.func.is_functional_fn()\\n and not f.func.is_out_fn()\\n and name == str(f.func.name.name)\\n ):\\n return f\\n # some functions only have in-place variants\\n assert name + \\\"_\\\" == cpp.name(functions[0].func)\\n return functions[0]\\n\\n def split_names(raw_names: str) -> tuple[str, ...]:\\n \\\"\\\"\\\"Given \\\"foo, bar\\\", return [\\\"foo\\\", \\\"bar\\\"].\\\"\\\"\\\"\\n return tuple(x.strip() for x in raw_names.split(\\\",\\\"))\\n\\n def check_grad_usage(defn_name: str, derivatives: Sequence[Derivative]) -> None:\\n \\\"\\\"\\\"\\n Check for some subtle mistakes one might make when writing derivatives.\\n These mistakes will compile, but will be latent until a function is\\n used with double backwards.\\n \\\"\\\"\\\"\\n\\n uses_grad = False # true if any derivative uses \\\"grad\\\"\\n num_grads_uses = 0 # count of uses of \\\"grads\\\" or \\\"grads[INDEX]\\\"\\n uses_named_grads = False # true if any derivative uses \\\"grad_{name}\\\"\\n used_grads_indices: list[int] = [] # which indices of grads are used\\n for d in derivatives:\\n formula = d.formula\\n uses_grad = uses_grad or bool(\\n re.findall(IDENT_REGEX.format(\\\"grad\\\"), formula)\\n )\\n num_grads_uses += len(re.findall(IDENT_REGEX.format(\\\"grads\\\"), formula))\\n uses_named_grads = uses_named_grads or bool(d.named_gradients)\\n used_grads_indices.extend(used_gradient_indices(formula))\\n # This is a basic sanity check: the number of places we see\\n # \\\"grads\\\" should be no fewer than the number of indices we see\\n # inside \\\"grads\\\". They may not be equal because we may use\\n # \\\"grads\\\" without an index.\\n assert num_grads_uses >= len(used_grads_indices)\\n # Thus if the number is equal, every use of grads is also\\n # indexed.\\n only_used_grads_indices = num_grads_uses == len(used_grads_indices)\\n\\n if uses_grad and num_grads_uses > 0:\\n raise RuntimeError(\\n f\\\"Derivative definition of {defn_name} in derivatives.yaml illegally \\\"\\n \\\"mixes use of 'grad' and 'grads'. Consider replacing \\\"\\n \\\"occurrences of 'grad' with 'grads[0]'\\\"\\n )\\n\\n if only_used_grads_indices and set(used_grads_indices) == {0}:\\n raise RuntimeError(\\n f\\\"Derivative definition of {defn_name} in derivatives.yaml solely \\\"\\n \\\"refers to 'grads[0]'. If the first output is indeed the \\\"\\n \\\"only differentiable output, replace 'grads[0]' with 'grad'; \\\"\\n \\\"otherwise, there is a likely error in your derivatives \\\"\\n \\\"declaration.\\\"\\n )\\n\\n if uses_named_grads and (uses_grad or num_grads_uses > 0):\\n raise RuntimeError(\\n f\\\"Derivative definition of {defn_name} in derivatives.yaml illegally \\\"\\n 'mixes use of \\\"grad_RETURN_NAME\\\" and \\\"grad\\\" or \\\"grads[x]\\\". Use '\\n \\\"only one method for identifying gradients.\\\"\\n )\\n\\n @with_native_function\\n def set_up_derivatives(\\n f: NativeFunction,\\n ) -> tuple[\\n Sequence[Derivative],\\n Sequence[ForwardDerivative],\\n Sequence[Binding],\\n Sequence[str],\\n Sequence[str],\\n ]:\\n # Set up the derivative information\\n derivatives: list[Derivative] = []\\n forward_derivatives: list[ForwardDerivative] = []\\n non_differentiable_arg_names: list[str] = []\\n args_with_derivatives_set: set[str] = set()\\n\\n all_arg_names = [a.name for a in cpp_arguments(f)]\\n all_ret_names = [\\n r.name for r in f.func.returns\\n ] # only used for the assert below\\n # output_differentiability is captured from the enclosed\\n # scope. Don't modify it.\\n #\\n # If it is not present, then no output is explicitly\\n # undifferentiable.\\n #\\n # It may be present and shorter than the length of return\\n # values. If that's the case, any return value that does not\\n # have a corresponding entry is considered not differentiable.\\n differentiability = output_differentiability or [True] * len(f.func.returns)\\n # A return is available as a named gradient ...\\n available_named_gradients = [\\n f\\\"grad_{ret.name}\\\"\\n for ret, differentiable in zip(f.func.returns, differentiability)\\n # if it has not been explicitly made undifferentiable\\n if differentiable\\n # and if it has a name\\n and ret.name is not None\\n # and if its type is differentiable\\n and ret.type.is_tensor_like()\\n ]\\n\\n for raw_names in sorted(defn.keys()):\\n formula = defn[raw_names]\\n names = split_names(raw_names)\\n\\n for name in names:\\n assert not (name in all_arg_names and name in all_ret_names), (\\n f\\\"While processing the derivative formula for '{f.func.name}' wrt '{name}', \\\"\\n f\\\"expected '{name}' to not be both an input arg and named return. \\\"\\n )\\n\\n if is_forward_derivative_definition(all_arg_names, names):\\n forward_derivatives.append(create_forward_derivative(f, formula, names))\\n else:\\n if formula.lower().strip() == \\\"non_differentiable\\\":\\n non_differentiable_arg_names += names\\n else:\\n derivative = create_derivative(\\n f, formula, names, available_named_gradients\\n )\\n derivatives.append(derivative)\\n args_with_derivatives_set |= set(names)\\n\\n overlap = args_with_derivatives_set.intersection(non_differentiable_arg_names)\\n if overlap:\\n raise RuntimeError(\\n f\\\"derivatives definition for {defn} have overlapped non_differentiable \\\"\\n f\\\"and differentiable variables: {overlap}\\\"\\n )\\n\\n # Next, let us determine the list of inputs in order.\\n # TODO: do we need eagerly calculate and save it here? Can it be derived\\n # from NativeFunction and `derivatives` on callsites instead?\\n args_with_derivatives = [\\n a for a in cpp_arguments(f) if a.name in args_with_derivatives_set\\n ]\\n\\n # Postprocess forward derivatives definitions now that we know the differentiable arguments\\n forward_derivatives = postprocess_forward_derivatives(\\n f,\\n defn_name,\\n all_arg_names,\\n derivatives,\\n forward_derivatives,\\n args_with_derivatives,\\n )\\n\\n # Test to see if the use of 'grads' makes sense.\\n check_grad_usage(defn_name, derivatives)\\n\\n return (\\n derivatives,\\n forward_derivatives,\\n args_with_derivatives,\\n non_differentiable_arg_names,\\n available_named_gradients,\\n )\\n\\n # NB: Removes 'name' from defn dictionary\\n specification = defn_dict.pop(\\\"name\\\")\\n defn_name, _ = split_name_params(specification)\\n # NB: Removes 'output_differentiability' from defn dictionary\\n # `None` means all differentiable.\\n output_differentiability = defn_dict.pop(\\\"output_differentiability\\\", None)\\n output_differentiability_conditions = None\\n if output_differentiability and any(\\n isinstance(diff, str) for diff in output_differentiability\\n ):\\n if len(output_differentiability) != 1:\\n raise RuntimeError(\\n f\\\"Not supported: for {specification},\\\"\\n f\\\"output_differentiability must either be \\\"\\n f\\\"List[bool] or a List[str] where each str is a \\\"\\n f\\\"condition. In the case where it is a condition, \\\"\\n f\\\"we only support single-output functions. \\\"\\n f\\\"Please file us an issue. \\\"\\n )\\n output_differentiability_conditions = output_differentiability\\n output_differentiability = [True]\\n\\n schema_function = functions_by_schema.get(specification)\\n if not schema_function:\\n avail = \\\"\\\\n\\\".join(\\n k for k, v in functions_by_schema.items() if cpp.name(v.func) == defn_name\\n )\\n raise RuntimeError(\\n f\\\"could not find ATen function for schema: {specification} \\\"\\n f\\\". Available signatures:\\\\n{avail}\\\"\\n )\\n\\n # now map this to the legacy schema; this isn't technically necessary, but we'd need some logic here\\n # to map in-place schemas to the out-of-place variants.\\n # TODO: maybe the logic to handle the legacy schema is no longer necessary?\\n signature = schema_function.func.signature()\\n functions = functions_by_signature[signature]\\n if len(functions) == 0:\\n avail = \\\"\\\\n\\\".join(\\n str(k)\\n for k, v in functions_by_signature.items()\\n if cpp.name(k) == defn_name\\n )\\n raise RuntimeError(\\n f\\\"could not find ATen function for legacy signature: {signature} \\\"\\n f\\\"corresponding to schema {specification}. Please report a bug to PyTorch. \\\"\\n f\\\"Available signatures:\\\\n{avail}\\\"\\n )\\n\\n canonical = canonical_function(functions, defn_name)\\n if \\\"grad_input_mask\\\" in (a.name for a in cpp_arguments(canonical)):\\n raise RuntimeError(\\n f\\\"Schema for {defn_name} has an argument named grad_input_mask, \\\"\\n \\\"but this name would be shadowed by our codegen. \\\"\\n \\\"Please use a different name in native_functions.yaml.\\\"\\n )\\n\\n if \\\"result\\\" in (a.name for a in cpp_arguments(canonical)):\\n raise RuntimeError(\\n f\\\"Schema for {defn_name} has an argument named result, \\\"\\n \\\"but this is only allowed for outputs.\\\"\\n \\\"Please use a different name in native_functions.yaml.\\\"\\n )\\n\\n diffinfo_dict = {}\\n for key, defn in defn_dict[\\\"dispatch\\\"].items():\\n if key != \\\"Default\\\" and key not in _VALID_AUTOGRAD_KEYS:\\n raise RuntimeError(\\n f\\\"Invalid dispatch key {key} in derivatives.yaml for {specification},\\\"\\n f\\\" expected key to be one of {_VALID_AUTOGRAD_KEYS}\\\"\\n )\\n if key not in used_dispatch_keys:\\n used_dispatch_keys.add(key)\\n\\n (\\n derivatives,\\n forward_derivatives,\\n args_with_derivatives,\\n non_differentiable_arg_names,\\n available_named_gradients,\\n ) = set_up_derivatives(canonical)\\n\\n used_named_gradients: set[str] = set()\\n for d in derivatives:\\n used_named_gradients |= d.named_gradients\\n\\n # only assign an op name if we are actually going to calculate a derivative\\n op = None\\n if args_with_derivatives:\\n op_prefix = _create_op_prefix(defn_name)\\n if key != \\\"Default\\\":\\n op_prefix = op_prefix + key\\n op = f\\\"{op_prefix}{op_counter[op_prefix]}\\\"\\n op_counter[op_prefix] += 1\\n\\n diffinfo_dict[key] = DifferentiabilityInfo(\\n name=defn_name,\\n func=canonical,\\n op=op,\\n derivatives=derivatives,\\n forward_derivatives=forward_derivatives,\\n all_saved_inputs=dedup_vars(\\n [v for d in derivatives for v in d.saved_inputs]\\n ),\\n all_saved_outputs=dedup_vars(\\n [v for d in derivatives for v in d.saved_outputs]\\n ),\\n available_named_gradients=available_named_gradients,\\n used_named_gradients=used_named_gradients,\\n args_with_derivatives=args_with_derivatives,\\n non_differentiable_arg_names=non_differentiable_arg_names,\\n output_differentiability=output_differentiability,\\n output_differentiability_conditions=output_differentiability_conditions,\\n )\\n\\n return canonical.func, diffinfo_dict\\n\\n\\nGRAD_INDEX_REGEX = r\\\"(?:^|\\\\W)grads\\\\[(\\\\d+)\\\\]\\\"\\n\\n\\ndef used_gradient_indices(formula: str) -> list[int]:\\n \\\"\\\"\\\"Determine a list of gradient indices (the i in grads[i]) that\\n are used by the formula.\\n\\n >>> used_gradient_indices(\\\"foo(grads[0], grads[1])\\\")\\n [0, 1]\\n \\\"\\\"\\\"\\n return [int(i) for i in re.findall(GRAD_INDEX_REGEX, formula)]\\n\\n\\ndef saved_variables(\\n formula: str,\\n nctypes: list[NamedCType],\\n var_names: tuple[str, ...],\\n) -> tuple[str, tuple[SavedAttribute, ...]]:\\n def stride_expr(name: str) -> str:\\n assert var_names == (name,), (\\n 'Replacement for \\\".strides()\\\" is currently only supported for single derivatives of the same tensor '\\n 'that \\\".strides()\\\" is being called on.'\\n )\\n return f'strides_or_error({name}, \\\"{name}\\\")'\\n\\n REPLACEMENTS: list[tuple[str, dict[str, Any]]] = [\\n # replace self.sym_sizes() with self_sym_sizes\\n (\\n r\\\"{}.sym_sizes\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_sym_sizes\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(symIntArrayRefT)),\\n },\\n ),\\n # replace self->sym_sizes() with self_sym_sizes_opt\\n (\\n r\\\"{}->sym_sizes\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_sym_sizes_opt\\\",\\n \\\"nctype\\\": lambda name: NamedCType(\\n name, OptionalCType(BaseCType(symIntArrayRefT))\\n ),\\n \\\"expr\\\": lambda name: f\\\"{name}.has_value() ? std::optional({name}->sym_sizes()) : std::nullopt\\\",\\n },\\n ),\\n # replace self.sym_blocksize() with self_sym_blocksize_opt\\n (\\n r\\\"{}.sym_blocksize\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_self_sym_blocksize_opt\\\",\\n \\\"nctype\\\": lambda name: NamedCType(\\n name, OptionalCType(BaseCType(symIntArrayRefT))\\n ),\\n \\\"expr\\\": lambda name: f\\\"at::sparse_csr::getSymIntBlockSize({name})\\\",\\n },\\n ),\\n # replace self.options() with self_options\\n (\\n r\\\"{}.options\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_options\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(tensorOptionsT)),\\n },\\n ),\\n # replace zeros_like(self) with self_info\\n (\\n r\\\"zeros_like\\\\({}\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_info\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(typeAndSizeT)),\\n \\\"expr\\\": lambda name: name, # at save-time\\n \\\"res\\\": lambda name: name + \\\"_info.zeros()\\\", # at eval-time\\n },\\n ),\\n # replace self.sym_size(2) with self_sym_size_2\\n (\\n r\\\"{}.sym_size\\\\((-?\\\\w+)\\\\)\\\",\\n {\\n \\\"suffix\\\": lambda m: f\\\"_sym_argsize_{m.groups()[0].replace('-', 'minus_')}\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(SymIntT)),\\n },\\n ),\\n # replace self.numel() with self_numel\\n (\\n r\\\"{}.numel\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_numel\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(longT)),\\n },\\n ),\\n # replace self.sym_numel() with self_sym_numel\\n (\\n r\\\"{}.sym_numel\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_sym_numel\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(SymIntT)),\\n },\\n ),\\n # replace to_args_sizes(self) with self_args_sizes\\n (\\n r\\\"to_args_sizes\\\\({}\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_args_sizes\\\",\\n \\\"nctype\\\": lambda name: NamedCType(\\n name, VectorCType(VectorCType(BaseCType(longT)))\\n ),\\n },\\n ),\\n # replace to_args_sizes_symint(self) with self_args_sizes\\n (\\n r\\\"to_args_sizes_symint\\\\({}\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_args_sizes_symint\\\",\\n \\\"nctype\\\": lambda name: NamedCType(\\n name, VectorCType(VectorCType(BaseCType(SymIntT)))\\n ),\\n },\\n ),\\n # replace to_args_scalartypes(self) with self_args_scalartypes\\n (\\n r\\\"to_args_scalartypes\\\\({}\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_args_scalartypes\\\",\\n \\\"nctype\\\": lambda name: NamedCType(\\n name, VectorCType(BaseCType(scalarTypeT))\\n ),\\n },\\n ),\\n # replace TensorGeometry(self) with self_geometry\\n (\\n r\\\"TensorGeometry\\\\({}\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_geometry\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(tensorGeometryT)),\\n },\\n ),\\n (\\n r\\\"{}.scalar_type\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_scalar_type\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(scalarTypeT)),\\n },\\n ),\\n # replace self.dim() with self_dim\\n (\\n r\\\"{}.dim\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_dim\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(longT)),\\n },\\n ),\\n # replace self.sym_strides() with self_sym_strides\\n (\\n r\\\"{}.sym_strides\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_sym_strides\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(symIntArrayRefT)),\\n \\\"expr\\\": stride_expr,\\n },\\n ),\\n # replace self.layout() with self_layout\\n (\\n r\\\"{}.layout\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_layout\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(layoutT)),\\n },\\n ),\\n # replace self.is_conj() with self_conjugate\\n (\\n r\\\"{}.is_conj\\\\(\\\\)\\\",\\n {\\n \\\"suffix\\\": \\\"_conjugate\\\",\\n \\\"nctype\\\": lambda name: NamedCType(name, BaseCType(boolT)),\\n },\\n ),\\n ]\\n\\n # find which arguments need to be saved\\n saved: list[SavedAttribute] = []\\n\\n if \\\".sizes()\\\" in formula or \\\"->sizes()\\\" in formula:\\n raise RuntimeError(\\n \\\".sizes() is not supported in derivative formulas. Instead, please use the SymInt version,\\\"\\n + f\\\".sym_sizes(), which returned a c10::SymIntArrayRef. formula={formula}\\\"\\n )\\n if re.search(r\\\"\\\\.size\\\\([-]?\\\\d+\\\\)\\\", formula) or re.search(\\n r\\\"->size\\\\([-]?\\\\d+\\\\)\\\", formula\\n ):\\n raise RuntimeError(\\n \\\".size(int) is not supported in derivative formulas. Instead, please use the SymInt version,\\\"\\n + f\\\".sym_size(int), which returned a c10::SymIntArrayRef. formula={formula}\\\"\\n )\\n if \\\".strides()\\\" in formula or \\\"->strides()\\\" in formula:\\n raise RuntimeError(\\n \\\".strides() is not supported in derivative formulas. Instead, please use the SymInt version,\\\"\\n + f\\\".sym_strides(), which returned a c10::SymIntArrayRef. formula={formula}\\\"\\n )\\n for nctype in nctypes:\\n name = (\\n nctype.name.name if isinstance(nctype.name, SpecialArgName) else nctype.name\\n )\\n # First search the formula for expressions which can be evaluated\\n # when the autograd Function is created to avoid saving variables\\n for regex, info in REPLACEMENTS:\\n\\n def repl(m: re.Match[str]) -> str:\\n suffix: str = (\\n info[\\\"suffix\\\"](m) if callable(info[\\\"suffix\\\"]) else info[\\\"suffix\\\"]\\n )\\n expr: str = info[\\\"expr\\\"](name) if \\\"expr\\\" in info else m.group(0)\\n saved.append(\\n SavedAttribute(\\n nctype=info[\\\"nctype\\\"](name + suffix),\\n expr=expr,\\n )\\n )\\n if \\\"res\\\" in info:\\n replacement: str = info[\\\"res\\\"](name)\\n return replacement\\n return name + suffix\\n\\n formula = re.sub(regex.format(name), repl, formula)\\n\\n # std::optional types stored in Backward nodes must be\\n # converted to std::optional before being passed into\\n # the backward function\\n if nctype.type == OptionalCType(BaseCType(stringT)):\\n formula = re.sub(\\n rf\\\"\\\\b{name}\\\\b\\\",\\n f\\\"{name}.has_value() ? std::optional({name}.value()) : std::nullopt\\\",\\n formula,\\n )\\n\\n # Find any variables which remain in the formula and save them\\n if re.search(IDENT_REGEX.format(name), formula):\\n saved.append(\\n SavedAttribute(\\n nctype=nctype,\\n expr=name,\\n )\\n )\\n\\n return formula, tuple(saved)\\n\\n\\ndef _create_op_prefix(name: str) -> str:\\n \\\"\\\"\\\"Takes a native function name converts to a op prefix name.\\n\\n Note that the \\\"name\\\" parameter must be the native function name\\n without the optional variant suffix, so \\\"add\\\" instead of\\n \\\"add.out\\\".\\n\\n OP names correspond to classes, hence the change to title case.\\n\\n Example::\\n >>> _create_op_prefix('add')\\n 'AddBackward'\\n \\\"\\\"\\\"\\n camel_case = \\\"\\\".join([p.title() for p in name.split(\\\"_\\\")])\\n return (camel_case + \\\"Backward\\\").replace(\\\"ForwardBackward\\\", \\\"Backward\\\")\\n\\n\\ndef dedup_vars(vars: Sequence[SavedAttribute]) -> Sequence[SavedAttribute]:\\n seen: set[str] = set()\\n saved: list[SavedAttribute] = []\\n for var in vars:\\n name = (\\n var.nctype.name.name\\n if isinstance(var.nctype.name, SpecialArgName)\\n else var.nctype.name\\n )\\n if name in seen:\\n continue\\n seen.add(name)\\n saved.append(var)\\n return saved\\n\\n\\n# Generates ViewFuncs.h/cpp\\n#\\n# NOTE: If any changes are being made to the ViewFunc codegen please also check\\n# if updates are needed in torch/csrc/autograd/autograd_not_implemented_fallback.cpp\\n# The fallback is expected to mimic this codegen, so we should keep the two in sync.\\n\\nfrom __future__ import annotations\\n\\nfrom typing import TYPE_CHECKING\\n\\nimport torchgen.api.dispatcher as dispatcher\\nfrom torchgen.api.translate import translate\\nfrom torchgen.api.types import (\\n BaseCType,\\n Binding,\\n NamedCType,\\n SymIntT,\\n tensorT,\\n VectorCType,\\n)\\nfrom torchgen.code_template import CodeTemplate\\nfrom torchgen.model import Argument, NativeFunction, OptionalType\\nfrom torchgen.utils import FileManager\\n\\nfrom .gen_inplace_or_view_type import (\\n CALL_DISPATCH,\\n extract_bindings,\\n get_view_info,\\n modifies_arguments,\\n use_derived,\\n)\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.api.autograd import NativeFunctionWithDifferentiabilityInfo\\n\\n\\nFUNCTION_DECLARATION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n#define ${uppercase_op}_AVAILABLE\\nstruct ${op} : public ${superclass} {\\n ${op}(${constructor_args}) ${initializer_list}\\n {};\\n virtual ~${op}() override {};\\n virtual std::vector get_symints() const override;\\n virtual size_t num_symints() const override;\\n virtual std::vector get_tensors() const override;\\n virtual size_t num_tensors() const override;\\n virtual at::Tensor operator()(const at::Tensor&) const override;\\n virtual std::unique_ptr clone_and_set(\\n std::optional> = ::std::nullopt,\\n std::optional> = ::std::nullopt) const override;\\n\\nprotected:\\n virtual void set_symints(std::vector) override;\\n virtual void set_tensors(std::vector) override;\\n\\nprivate:\\n ${state}\\n};\\n\\n\\\"\\\"\\\"\\n)\\n\\nFUNCTION_DEFINITION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstd::vector ${op}::get_symints() const {\\n ${get_symints}\\n}\\n\\nsize_t ${op}::num_symints() const {\\n return static_cast(${num_symints});\\n}\\n\\nvoid ${op}::set_symints(std::vector ${symints_vec}) {\\n TORCH_INTERNAL_ASSERT(${symints_vec}.size() == num_symints());\\n ${set_symints}\\n}\\n\\nstd::vector ${op}::get_tensors() const {\\n ${get_tensors}\\n}\\n\\nsize_t ${op}::num_tensors() const {\\n return static_cast(${num_tensors});\\n}\\n\\nvoid ${op}::set_tensors(std::vector ${tensors_vec}) {\\n TORCH_INTERNAL_ASSERT(${tensors_vec}.size() == num_tensors());\\n ${set_tensors}\\n}\\n\\nat::Tensor ${op}::operator()(const at::Tensor& ${call_input_name}) const {\\n return ${op_call};\\n}\\n\\nstd::unique_ptr ${op}::clone_and_set(\\n std::optional> ${symints_vec},\\n std::optional> ${tensors_vec}) const {\\n auto output = std::make_unique<${op}>(${clone_args});\\n if (${symints_vec}.has_value()) {\\n output->set_symints(std::move(*(${symints_vec})));\\n }\\n if (${tensors_vec}.has_value()) {\\n output->set_tensors(std::move(*(${tensors_vec})));\\n }\\n return output;\\n}\\n\\n\\\"\\\"\\\"\\n)\\n\\n\\n# e.g. as_strided -> AsStridedViewFunc for camel case or\\n# as_strided_view_func otherwise\\ndef view_func_name(\\n f: NativeFunction, include_namespace: bool = False, camel_case: bool = True\\n) -> str:\\n name = f.func.name.unambiguous_name()\\n view_func_name = f\\\"{name.replace('.', '_')}_view_func\\\"\\n if camel_case:\\n is_private = view_func_name.startswith(\\\"_\\\")\\n view_func_name = \\\"\\\".join(\\n [p.title() for p in view_func_name.replace(\\\".\\\", \\\"_\\\").split(\\\"_\\\")]\\n )\\n if is_private:\\n # put the leading underscore back in\\n view_func_name = f\\\"_{view_func_name}\\\"\\n namespace = \\\"torch::autograd::generated::\\\" if include_namespace else \\\"\\\"\\n return f\\\"{namespace}{view_func_name}\\\"\\n\\n\\ndef is_symint_or_tensor(arg: Argument) -> bool:\\n return arg.type.is_tensor_like() or arg.type.is_symint_like()\\n\\n\\ndef remove_const_ref(binding: Binding) -> Binding:\\n return Binding(\\n name=binding.name,\\n nctype=binding.nctype.remove_const_ref(),\\n argument=binding.argument,\\n default=binding.default,\\n )\\n\\n\\ndef returns_multi_tensor(fn: NativeFunction) -> bool:\\n returns = fn.func.returns\\n assert len(returns) == 1\\n returns_list_like = returns[0].type.is_list_like() is not None\\n returns_tensor_like = returns[0].type.is_tensor_like()\\n return returns_list_like and returns_tensor_like\\n\\n\\n# Generates strings with logic for getting / setting state of a particular type.\\n#\\n# Args:\\n# bindings (list): List of state bindings of interest (may be empty)\\n# state_vec_type (NamedCType): Type of vector to either return or copy from\\n#\\n# Returns:\\n# tuple: (list of getter logic strings, list of setter logic strings, string\\n# with num items expression)\\ndef generate_state_getter_setter(\\n bindings: list[Binding],\\n state_vec_type: NamedCType,\\n) -> tuple[list[str], list[str], str]:\\n getter_logic = []\\n setter_logic = []\\n\\n state_vec = state_vec_type.name\\n getter_logic.append(f\\\"{state_vec_type.cpp_type()} {state_vec};\\\")\\n if len(bindings) > 0:\\n setter_logic.append(\\\"auto i = 0;\\\")\\n\\n num_exprs = []\\n for i, b in enumerate(bindings):\\n assert isinstance(b.argument, Argument)\\n if b.argument.type.is_list_like():\\n # Handle list-likes.\\n num_expr = f\\\"{b.name}.size()\\\"\\n num_exprs.append(num_expr)\\n getter = f\\\"{state_vec}.insert({state_vec}.end(), {b.name}.begin(), {b.name}.end());\\\"\\n setter = f\\\"std::copy({state_vec}.begin() + i, {state_vec}.begin() + i + {b.name}.size(), {b.name}.begin());\\\"\\n elif isinstance(b.argument.type, OptionalType):\\n # Handle optionals.\\n num_expr = f\\\"({b.name}.has_value() ? 1 : 0)\\\"\\n num_exprs.append(num_expr)\\n conditional = f\\\"if({b.name}.has_value())\\\"\\n getter = (\\n f\\\"{conditional} {state_vec}.insert({state_vec}.end(), *({b.name}));\\\"\\n )\\n setter = f\\\"{conditional} {b.name} = {state_vec}[i];\\\"\\n else:\\n num_expr = \\\"1\\\"\\n num_exprs.append(num_expr)\\n getter = f\\\"{state_vec}.push_back({b.name});\\\"\\n setter = f\\\"{b.name} = {state_vec}[i];\\\"\\n\\n getter_logic.append(getter)\\n setter_logic.append(setter)\\n if i < len(bindings) - 1:\\n setter_logic.append(f\\\"i += {num_expr};\\\")\\n\\n # Reserve / assert based on the total number of items expression.\\n num_items = \\\"0\\\" if len(num_exprs) == 0 else \\\" + \\\".join(num_exprs)\\n if len(bindings) > 0:\\n getter_logic.insert(1, f\\\"{state_vec}.reserve({num_items});\\\")\\n\\n getter_logic.append(f\\\"return {state_vec};\\\")\\n\\n return getter_logic, setter_logic, num_items\\n\\n\\ndef process_function(fn: NativeFunction, template: CodeTemplate) -> str:\\n bindings = extract_bindings(fn)\\n non_self_bindings = [b for b in bindings if b.name != \\\"self\\\"]\\n\\n non_self_args = fn.func.arguments.flat_all[1:]\\n non_self_value_bindings = [\\n dispatcher.argument(a, remove_non_owning_ref_types=True) for a in non_self_args\\n ]\\n\\n # Generate constructor / clone args for the generated struct.\\n constructor_args = [b.defn() for b in non_self_bindings]\\n clone_args = [b.name for b in non_self_bindings]\\n\\n # Generate state variable declarations for the generated struct.\\n state_variables = [\\n f\\\"{remove_const_ref(b).defn()};\\\" for b in non_self_value_bindings\\n ]\\n\\n # Generate initializer list expressions for the generated struct.\\n # allow_expensive_conversions=True because we need to store e.g. SymIntArrayRefs as\\n # vectors.\\n init_exprs = translate(\\n non_self_bindings, non_self_value_bindings, allow_expensive_conversions=True\\n )\\n initializers = []\\n for b, init_expr in zip(non_self_bindings, init_exprs):\\n name = b.nctype.name\\n assert isinstance(name, str)\\n initializers.append(f\\\"{name}({init_expr.expr})\\\")\\n\\n # Generate call to underlying view op\\n call_input_name = \\\"input_base\\\"\\n op_call_args = [call_input_name, *(b.name for b in non_self_bindings)]\\n op_call = CALL_DISPATCH.substitute(\\n unambiguous_name=fn.func.name.unambiguous_name(),\\n unpacked_args=op_call_args,\\n )\\n\\n # Multi-output views additionally require a view_idx for disambiguation.\\n if returns_multi_tensor(fn):\\n view_idx_name = \\\"view_idx\\\"\\n view_idx_typename = \\\"int64_t\\\"\\n view_idx_decl = f\\\"{view_idx_typename} {view_idx_name}\\\"\\n constructor_args.append(view_idx_decl)\\n clone_args.append(view_idx_name)\\n state_variables.append(f\\\"{view_idx_decl};\\\")\\n initializers.append(f\\\"{view_idx_name}({view_idx_name})\\\")\\n op_call += f\\\"[{view_idx_name}]\\\"\\n\\n # Generate initializer list for the generated struct.\\n initializer_list = f\\\": {', '.join(initializers)}\\\" if len(initializers) > 0 else \\\"\\\"\\n\\n # Generate getter / setter logic for any symints.\\n symint_bindings = [\\n b\\n for b in non_self_bindings\\n if isinstance(b.argument, Argument) and b.argument.type.is_symint_like()\\n ]\\n symints_vec_type = NamedCType(\\\"symints\\\", VectorCType(BaseCType(SymIntT)))\\n get_symints, set_symints, num_symints = generate_state_getter_setter(\\n symint_bindings, symints_vec_type\\n )\\n\\n # Generate getter / setter logic for any tensors.\\n tensor_bindings = [\\n b\\n for b in non_self_bindings\\n if isinstance(b.argument, Argument) and b.argument.type.is_tensor_like()\\n ]\\n tensors_vec_type = NamedCType(\\\"tensors\\\", VectorCType(BaseCType(tensorT)))\\n get_tensors, set_tensors, num_tensors = generate_state_getter_setter(\\n tensor_bindings, tensors_vec_type\\n )\\n\\n return template.substitute(\\n op=view_func_name(fn),\\n uppercase_op=view_func_name(fn, camel_case=False).upper(),\\n superclass=\\\"torch::autograd::ViewFunc\\\",\\n initializer_list=initializer_list,\\n state=state_variables,\\n constructor_args=constructor_args,\\n clone_args=clone_args,\\n symints_vec=symints_vec_type.name,\\n get_symints=get_symints,\\n set_symints=set_symints,\\n num_symints=num_symints,\\n tensors_vec=tensors_vec_type.name,\\n get_tensors=get_tensors,\\n set_tensors=set_tensors,\\n num_tensors=num_tensors,\\n call_input_name=call_input_name,\\n op_call=op_call,\\n )\\n\\n\\ndef gen_view_funcs(\\n out: str,\\n fns_with_infos: list[NativeFunctionWithDifferentiabilityInfo],\\n template_path: str,\\n) -> None:\\n # don't need the info parts, just the function\\n fns = [fn.func for fn in fns_with_infos if use_derived(fn)]\\n # only want out-of-place views\\n view_fns = [\\n fn for fn in fns if get_view_info(fn) is not None and not modifies_arguments(fn)\\n ]\\n\\n declarations = [process_function(fn, FUNCTION_DECLARATION) for fn in view_fns]\\n definitions = [process_function(fn, FUNCTION_DEFINITION) for fn in view_fns]\\n ops_headers = [f\\\"#include \\\" for fn in view_fns]\\n\\n file_basename = \\\"ViewFuncs\\\"\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n for suffix in [\\\".h\\\", \\\".cpp\\\"]:\\n fname = file_basename + suffix\\n fm.write_with_template(\\n fname,\\n fname,\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/\\\"\\n + fname,\\n \\\"view_func_declarations\\\": declarations,\\n \\\"view_func_definitions\\\": definitions,\\n \\\"ops_headers\\\": ops_headers,\\n },\\n )\\n\\n\\nIf you add a file to this directory, you **MUST** update\\n`torch/CMakeLists.txt` and add the file as a dependency to\\nthe `add_custom_command` call.\\n\\n\\n# Generates C++ functions that wrap ATen tensor factory methods to turn them into Variables.\\n#\\n# This writes one file: variable_factories.h\\n\\nfrom __future__ import annotations\\n\\nimport re\\n\\nimport torchgen.api.python as python\\nfrom torchgen.api import cpp\\nfrom torchgen.api.types import CppSignatureGroup\\nfrom torchgen.context import with_native_function\\nfrom torchgen.gen import parse_native_yaml\\nfrom torchgen.model import NativeFunction, TensorOptionsArguments, Variant\\nfrom torchgen.utils import FileManager, mapMaybe\\n\\n\\nOPTIONAL_TYPE_PATTERN = re.compile(r\\\"std::optional<(.+)>\\\")\\nTYPE_PATTERN = re.compile(r\\\"(?:const\\\\s+)?([A-Z]\\\\w+)\\\")\\n\\n\\n# Add 'at::' to types defined in ATen namespace, e.g. Tensor, TensorList, IntArrayRef and etc.\\n# TODO: maybe update the cpp argument API to take optional namespace argument?\\ndef fully_qualified_type(argument_type: str) -> str:\\n def maybe_optional_type(type: str, is_opt: bool) -> str:\\n return f\\\"std::optional<{type}>\\\" if is_opt else type\\n\\n opt_match = OPTIONAL_TYPE_PATTERN.match(argument_type)\\n is_opt = opt_match is not None\\n if opt_match:\\n argument_type = argument_type[opt_match.start(1) : opt_match.end(1)]\\n match = TYPE_PATTERN.match(argument_type)\\n if match is None:\\n return maybe_optional_type(argument_type, is_opt)\\n index = match.start(1)\\n qualified_type = f\\\"{argument_type[:index]}at::{argument_type[index:]}\\\"\\n return maybe_optional_type(qualified_type, is_opt)\\n\\n\\ndef gen_variable_factories(\\n out: str, native_yaml_path: str, tags_yaml_path: str, template_path: str\\n) -> None:\\n native_functions = parse_native_yaml(\\n native_yaml_path, tags_yaml_path\\n ).native_functions\\n factory_functions = [fn for fn in native_functions if is_factory_function(fn)]\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n fm.write_with_template(\\n \\\"variable_factories.h\\\",\\n \\\"variable_factories.h\\\",\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/variable_factories.h\\\",\\n \\\"ops_headers\\\": [\\n f\\\"#include \\\" for fn in factory_functions\\n ],\\n \\\"function_definitions\\\": list(mapMaybe(process_function, factory_functions)),\\n },\\n )\\n\\n\\n@with_native_function\\ndef is_factory_function(f: NativeFunction) -> bool:\\n if Variant.function not in f.variants:\\n return False\\n\\n name = cpp.name(f.func)\\n has_tensor_options = python.has_tensor_options(f)\\n return has_tensor_options or name.endswith(\\\"_like\\\")\\n\\n\\n@with_native_function\\ndef process_function(f: NativeFunction) -> str | None:\\n name = cpp.name(f.func)\\n has_tensor_options = python.has_tensor_options(f)\\n is_factory = has_tensor_options or name.endswith(\\\"_like\\\")\\n\\n if Variant.function not in f.variants or not is_factory:\\n return None\\n\\n cpp_sigs = CppSignatureGroup.from_native_function(f, method=False)\\n sigs = [cpp_sigs.signature]\\n if cpp_sigs.symint_signature is not None:\\n sigs.append(cpp_sigs.symint_signature)\\n r = \\\"\\\"\\n for sig in sigs:\\n formals: list[str] = []\\n exprs: list[str] = []\\n requires_grad = \\\"false\\\"\\n for arg in sig.arguments():\\n qualified_type = fully_qualified_type(arg.type)\\n if arg.default:\\n formals.append(f\\\"{qualified_type} {arg.name} = {arg.default}\\\")\\n else:\\n formals.append(f\\\"{qualified_type} {arg.name}\\\")\\n\\n if isinstance(arg.argument, TensorOptionsArguments):\\n # note: we remove the requires_grad setting from the TensorOptions because\\n # it is ignored anyways (and we actually have an assertion that it isn't set\\n # which would fail otherwise). We handle requires_grad explicitly here\\n # instead of passing it through to the kernel.\\n exprs.append(\\n f\\\"at::TensorOptions({arg.name}).requires_grad(::std::nullopt)\\\"\\n )\\n # Manually set the requires_grad bit on the result tensor.\\n requires_grad = f\\\"{arg.name}.requires_grad()\\\"\\n else:\\n exprs.append(arg.name)\\n\\n r += f\\\"\\\"\\\"\\\\\\ninline at::Tensor {sig.name()}({', '.join(formals)}) {{\\n at::AutoDispatchBelowADInplaceOrView guard;\\n return autograd::make_variable(at::{sig.name()}({', '.join(exprs)}), /*requires_grad=*/{requires_grad});\\n}}\\n\\\"\\\"\\\"\\n return r\\n\\n\\n\\\"\\\"\\\"\\nFor procedural tests needed for __torch_function__, we use this function\\nto export method names and signatures as needed by the tests in\\ntest/test_overrides.py.\\n\\npython -m tools.autograd.gen_annotated_fn_args \\\\\\n aten/src/ATen/native/native_functions.yaml \\\\\\n aten/src/ATen/native/tags.yaml \\\\\\n $OUTPUT_DIR \\\\\\n tools/autograd\\n\\nWhere $OUTPUT_DIR is where you would like the files to be\\ngenerated. In the full build system, OUTPUT_DIR is\\ntorch/testing/_internal/generated\\n\\\"\\\"\\\"\\n\\nfrom __future__ import annotations\\n\\nimport argparse\\nimport os\\nimport textwrap\\nfrom collections import defaultdict\\nfrom typing import Any, Sequence, TYPE_CHECKING\\n\\nimport torchgen.api.python as python\\nfrom torchgen.context import with_native_function\\nfrom torchgen.gen import parse_native_yaml\\nfrom torchgen.utils import FileManager\\n\\nfrom .gen_python_functions import (\\n is_py_fft_function,\\n is_py_linalg_function,\\n is_py_nn_function,\\n is_py_special_function,\\n is_py_torch_function,\\n is_py_variable_method,\\n should_generate_py_binding,\\n)\\n\\n\\nif TYPE_CHECKING:\\n from torchgen.model import Argument, BaseOperatorName, NativeFunction\\n\\n\\ndef gen_annotated(\\n native_yaml_path: str, tags_yaml_path: str, out: str, autograd_dir: str\\n) -> None:\\n native_functions = parse_native_yaml(\\n native_yaml_path, tags_yaml_path\\n ).native_functions\\n mappings = (\\n (is_py_torch_function, \\\"torch._C._VariableFunctions\\\"),\\n (is_py_nn_function, \\\"torch._C._nn\\\"),\\n (is_py_linalg_function, \\\"torch._C._linalg\\\"),\\n (is_py_special_function, \\\"torch._C._special\\\"),\\n (is_py_fft_function, \\\"torch._C._fft\\\"),\\n (is_py_variable_method, \\\"torch.Tensor\\\"),\\n )\\n annotated_args: list[str] = []\\n for pred, namespace in mappings:\\n groups: dict[BaseOperatorName, list[NativeFunction]] = defaultdict(list)\\n for f in native_functions:\\n if not should_generate_py_binding(f) or not pred(f):\\n continue\\n groups[f.func.name.name].append(f)\\n for group in groups.values():\\n for f in group:\\n annotated_args.append(f\\\"{namespace}.{gen_annotated_args(f)}\\\")\\n\\n template_path = os.path.join(autograd_dir, \\\"templates\\\")\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n fm.write_with_template(\\n \\\"annotated_fn_args.py\\\",\\n \\\"annotated_fn_args.py.in\\\",\\n lambda: {\\n \\\"annotated_args\\\": textwrap.indent(\\\"\\\\n\\\".join(annotated_args), \\\" \\\"),\\n },\\n )\\n\\n\\n@with_native_function\\ndef gen_annotated_args(f: NativeFunction) -> str:\\n def _get_kwargs_func_exclusion_list() -> list[str]:\\n # functions that currently don't work with kwargs in test_overrides.py\\n return [\\n \\\"diagonal\\\",\\n \\\"round_\\\",\\n \\\"round\\\",\\n \\\"scatter_\\\",\\n ]\\n\\n def _add_out_arg(\\n out_args: list[dict[str, Any]], args: Sequence[Argument], *, is_kwarg_only: bool\\n ) -> None:\\n for arg in args:\\n if arg.default is not None:\\n continue\\n out_arg: dict[str, Any] = {}\\n out_arg[\\\"is_kwarg_only\\\"] = str(is_kwarg_only)\\n out_arg[\\\"name\\\"] = arg.name\\n out_arg[\\\"simple_type\\\"] = python.argument_type_str(\\n arg.type, simple_type=True\\n )\\n size_t = python.argument_type_size(arg.type)\\n if size_t:\\n out_arg[\\\"size\\\"] = size_t\\n out_args.append(out_arg)\\n\\n out_args: list[dict[str, Any]] = []\\n _add_out_arg(out_args, f.func.arguments.flat_positional, is_kwarg_only=False)\\n if f\\\"{f.func.name.name}\\\" not in _get_kwargs_func_exclusion_list():\\n _add_out_arg(out_args, f.func.arguments.flat_kwarg_only, is_kwarg_only=True)\\n\\n return f\\\"{f.func.name.name}: {repr(out_args)},\\\"\\n\\n\\ndef main() -> None:\\n parser = argparse.ArgumentParser(description=\\\"Generate annotated_fn_args script\\\")\\n parser.add_argument(\\n \\\"native_functions\\\", metavar=\\\"NATIVE\\\", help=\\\"path to native_functions.yaml\\\"\\n )\\n parser.add_argument(\\\"tags\\\", metavar=\\\"TAGS\\\", help=\\\"path to tags.yaml\\\")\\n parser.add_argument(\\\"out\\\", metavar=\\\"OUT\\\", help=\\\"path to output directory\\\")\\n parser.add_argument(\\n \\\"autograd\\\", metavar=\\\"AUTOGRAD\\\", help=\\\"path to template directory\\\"\\n )\\n args = parser.parse_args()\\n gen_annotated(args.native_functions, args.tags, args.out, args.autograd)\\n\\n\\nif __name__ == \\\"__main__\\\":\\n main()\\n\\n\\n# Generates C++ autograd functions for the derivatives of ATen operations\\n#\\n# This writes two files:\\n# Functions.h/cpp: subclasses of autograd::Node\\n# python_functions.h/cpp: Python bindings for the above classes\\n#\\n\\nfrom __future__ import annotations\\n\\nfrom typing import Sequence\\n\\nfrom torchgen.api.autograd import (\\n Derivative,\\n DifferentiabilityInfo,\\n SavedAttribute,\\n uses_retain_variables,\\n uses_single_grad,\\n)\\nfrom torchgen.api.types import (\\n ArrayRefCType,\\n BaseCppType,\\n BaseCType,\\n Binding,\\n boolT,\\n doubleT,\\n intArrayRefT,\\n iTensorListRefT,\\n ListCType,\\n longT,\\n MutRefCType,\\n OptionalCType,\\n optionalIntArrayRefT,\\n optionalSymIntArrayRefT,\\n scalarT,\\n stringT,\\n symIntArrayRefT,\\n SymIntT,\\n TENSOR_LIST_LIKE_CTYPES,\\n tensorListT,\\n tensorT,\\n VectorCType,\\n)\\nfrom torchgen.code_template import CodeTemplate\\nfrom torchgen.model import Argument, FunctionSchema\\nfrom torchgen.utils import FileManager\\n\\nfrom .gen_inplace_or_view_type import VIEW_FUNCTIONS\\n\\n\\nFUNCTION_DECLARATION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n#ifdef _WIN32\\nstruct ${op} : public ${superclass} {\\n TORCH_API ${op}() = default;\\n#else\\nstruct TORCH_API ${op} : public ${superclass} {\\n#endif\\n using ${superclass}::${superclass};\\n variable_list apply(variable_list&& grads) override;\\n std::string name() const override { return \\\"${op}\\\"; }\\n void release_variables() override {\\n ${thread_lock}\\n ${release_variables}\\n }\\n ${will_release_variables}\\n void compiled_args(CompiledNodeArgs& args) override;\\n variable_list apply_with_saved(const variable_list& inputs, SwapSavedVariables& saved) override;\\n ${saved_variables}\\n ${saved_list_sizes}\\n};\\n\\\"\\\"\\\"\\n)\\n\\nWILL_RELEASE_VARIABLES = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nbool retain_variables = true;\\nvoid will_release_variables() override {\\n retain_variables = false;\\n}\\n\\\"\\\"\\\"\\n)\\n\\nFUNCTION_DEFINITION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nvariable_list ${op}::apply(variable_list&& grads) {\\n ${thread_lock}\\n ${asserts}\\n IndexRangeGenerator gen;\\n ${compute_index_ranges}\\n variable_list grad_inputs(gen.size());\\n ${body}\\n return grad_inputs;\\n}\\nvoid ${op}::compiled_args(CompiledNodeArgs& args) {\\n ${compiled_args}\\n}\\nvariable_list ${op}::apply_with_saved(const variable_list& grads, SwapSavedVariables& saved) {\\n ${apply_with_saved_before}\\n variable_list result = apply(variable_list(grads));\\n ${apply_with_saved_after}\\n return result;\\n}\\n\\\"\\\"\\\"\\n)\\n\\nGRAD_INPUT_MASK = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n auto grad_input_mask = std::array{\\n ${masks}\\n };\\\\\\n\\\"\\\"\\\"\\n)\\n\\nDERIVATIVE_SINGLE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (task_should_compute_output({ ${name}_ix })) {\\n auto grad_result = ${derivative};\\n copy_range(grad_inputs, ${name}_ix, grad_result);\\n}\\n\\\"\\\"\\\"\\n)\\n\\n# note(crcrpar): `self` argument and other optional positional argument\\n# of foreach functions are basically a list of n `Tensor`s thus iterating over\\n# `grads` in order to utilize and apply the existing derivative definitions\\n# to each `Tensor`(s) of `self`, and the others.\\nDERIVATIVE_SINGLE_FOREACH = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (task_should_compute_output({ ${name}_ix })) {\\n std::vector grad_result;\\n grad_result.reserve(grads.size());\\n for (const auto & i : c10::irange(grads.size())) {\\n if (grads[i].defined()) {\\n grad_result.emplace_back(${derivative});\\n } else {\\n grad_result.emplace_back(Tensor());\\n }\\n }\\n copy_range(grad_inputs, ${name}_ix, grad_result);\\n}\\n\\\"\\\"\\\"\\n)\\n\\nDERIVATIVE_MULTI_COPY_RANGE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n if (task_should_compute_output({ ${name}_ix })) {\\n copy_range(grad_inputs, ${name}_ix, std::get<${i}>(grad_result));\\n }\\n\\\"\\\"\\\"\\n)\\n\\nDERIVATIVE_MULTI = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (task_should_compute_output({ ${idx_ranges} })) {\\n ${grad_input_mask}\\n auto grad_result = ${derivative};\\n ${copy_ranges}\\n}\\n\\\"\\\"\\\"\\n)\\n\\n# Generates python bindings\\n#\\n# This generates the definitions for:\\n# (1) The PyTypeObject for each backward grad_fn subclassing Node\\n# (2) The entry for PyTypeObject's tp_getset slot (an array of PyGetSetDef structs)\\n# We generate one PyGetSetDef struct for each of grad_fn's saved inputs and outputs\\n# Each PyGetSetDef has a function ptr to a getter, also defined here (3).\\n# (3) Getters for each of grad_fn's saved inputs and outputs.\\n#\\nPY_FUNCTION_DEFINITION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstatic PyTypeObject ${op}Class;\\naddClass<${op}>(module, ${op}Class, \\\"${op}\\\", ${op}_properties);\\n\\\"\\\"\\\"\\n)\\n\\nPY_FUNCTION_PROPS_AND_GETTERS = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${all_getter_definitions}\\n\\nstatic struct PyGetSetDef ${op}_properties[] = {\\n THP_FUNCTION_DEFAULT_PROPERTIES,\\n ${all_getsetdef_structs}\\n {nullptr} /* sentinel */\\n};\\n\\n\\\"\\\"\\\"\\n)\\n\\nPY_GETSETDEF_STRUCT = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n{(char*)\\\"_saved_${name}\\\", (getter)THP${op}_${name}_getter, nullptr, nullptr, nullptr}\\\"\\\"\\\"\\n)\\n\\nPY_RAW_GETSETDEF_STRUCT = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n{(char*)\\\"_raw_saved_${name}\\\", (getter)THP${op}_${name}_raw_getter, nullptr, nullptr, nullptr}\\\"\\\"\\\"\\n)\\n\\n# Getter templates\\nGETTER_DEFINITION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nPyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) {\\n HANDLE_TH_ERRORS\\n auto prop = static_cast<${op}*>(self->cdata.get())->${name};\\n ${body}\\n END_HANDLE_TH_ERRORS\\n}\\n\\\"\\\"\\\"\\n)\\n\\nGETTER_DEFINITION_SAVEDVAR = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nPyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) {\\n HANDLE_TH_ERRORS\\n const auto& prop = static_cast<${op}*>(self->cdata.get())->${name}_;\\n ${body}\\n END_HANDLE_TH_ERRORS\\n}\\n\\\"\\\"\\\"\\n)\\n\\nGETTER_DEFINITION_RAW_SAVEDVAR = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nPyObject* THP${op}_${name}_raw_getter(THPCppFunction *self, void *_unused) {\\n HANDLE_TH_ERRORS\\n const auto& prop = static_cast<${op}*>(self->cdata.get())->${name}_;\\n ${body}\\n END_HANDLE_TH_ERRORS\\n}\\n\\\"\\\"\\\"\\n)\\n\\nGETTER_DEFINITION_VEC_SAVEDVAR = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nPyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) {\\n HANDLE_TH_ERRORS\\n const auto *node = static_cast<${op}*>(self->cdata.get());\\n const auto& prop = node->${name}_;\\n if (node->${name}_released_) {\\n PyErr_SetString(PyExc_RuntimeError, ERR_BACKWARD_TWICE);\\n return nullptr;\\n }\\n ${body}\\n END_HANDLE_TH_ERRORS\\n}\\n\\\"\\\"\\\"\\n)\\n\\nGETTER_DEFINITION_RAW_VEC_SAVEDVAR = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nPyObject* THP${op}_${name}_raw_getter(THPCppFunction *self, void *_unused) {\\n HANDLE_TH_ERRORS\\n const auto *node = static_cast<${op}*>(self->cdata.get());\\n const auto& prop = node->${name}_;\\n if (node->${name}_released_) {\\n PyErr_SetString(PyExc_RuntimeError, ERR_BACKWARD_TWICE);\\n return nullptr;\\n }\\n ${body}\\n END_HANDLE_TH_ERRORS\\n}\\n\\\"\\\"\\\"\\n)\\n\\nGETTER_DEFINITION_OPT = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nPyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) {\\n HANDLE_TH_ERRORS\\n auto opt_prop = static_cast<${op}*>(self->cdata.get())->${name};\\n if (!opt_prop.has_value()) {\\n Py_RETURN_NONE;\\n }\\n auto prop = opt_prop.value();\\n ${body}\\n END_HANDLE_TH_ERRORS\\n}\\n\\\"\\\"\\\"\\n)\\n\\nGETTER_DEFINITION_OPT_ARRAYREF = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nPyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) {\\n HANDLE_TH_ERRORS\\n auto opt_prop = static_cast<${op}*>(self->cdata.get())->${name};\\n if (!opt_prop.list.has_value()) {\\n Py_RETURN_NONE;\\n }\\n auto prop = opt_prop.list.value();\\n ${body}\\n END_HANDLE_TH_ERRORS\\n}\\n\\\"\\\"\\\"\\n)\\n\\n# Getter body\\nGETTER_BODY_SAVEDVAR = \\\"\\\"\\\"\\\\\\nreturn THPVariable_Wrap(prop.unpack(self->cdata));\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_RAW_SAVEDVAR = \\\"\\\"\\\"\\\\\\npybind11::object obj = pybind11::cast(prop, pybind11::return_value_policy::reference);\\nreturn obj.release().ptr();\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_VEC_SAVEDVAR = \\\"\\\"\\\"\\\\\\nPyObject* tup = PyTuple_New((Py_ssize_t) prop.size());\\nfor (auto i: c10::irange(prop.size())) {\\n PyTuple_SetItem(tup, (Py_ssize_t) i, THPVariable_Wrap(prop[i].unpack(self->cdata)));\\n}\\nreturn tup;\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_RAW_VEC_SAVEDVAR = \\\"\\\"\\\"\\\\\\nPyObject* tup = PyTuple_New((Py_ssize_t) prop.size());\\nfor (auto i : c10::irange(prop.size())) {\\n pybind11::object obj = pybind11::cast(prop[i], pybind11::return_value_policy::reference);\\n PyTuple_SetItem(tup, (Py_ssize_t) i, obj.release().ptr());\\n}\\nreturn tup;\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_ARRAYREF_LONG = \\\"\\\"\\\"\\\\\\nPyObject* tup = PyTuple_New((Py_ssize_t) prop.size());\\nfor (auto i : c10::irange(prop.size())) {\\n PyTuple_SetItem(tup, (Py_ssize_t) i, PyLong_FromUnsignedLong((uint64_t) prop[i]));\\n}\\nreturn tup;\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_ARRAYREF_SYMINT = \\\"\\\"\\\"\\\\\\nPyObject* tup = PyTuple_New((Py_ssize_t) prop.size());\\nfor (auto i : c10::irange(prop.size())) {\\n auto si = prop[i];\\n if (auto m = si.maybe_as_int()) {\\n PyTuple_SetItem(tup, (Py_ssize_t) i, PyLong_FromUnsignedLong(*m));\\n } else {\\n auto py_symint = py::cast(si).release().ptr();\\n PyTuple_SetItem(tup, (Py_ssize_t) i, py_symint);\\n }\\n}\\nreturn tup;\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_ARRAYREF_DOUBLE = \\\"\\\"\\\"\\\\\\nPyObject* tup = PyTuple_New((Py_ssize_t) prop.size());\\nfor (auto i : c10::irange(prop.size())) {\\n PyTuple_SetItem(tup, (Py_ssize_t) i, PyFloat_FromDouble((double) prop[i]));\\n}\\nreturn tup;\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_INT64_T = \\\"\\\"\\\"\\\\\\nreturn PyLong_FromUnsignedLong((int64_t) prop);\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_SYMINT = \\\"\\\"\\\"\\\\\\nif (auto m = prop.maybe_as_int()) {\\n return PyLong_FromUnsignedLong(*m);\\n} else {\\n return py::cast(prop).release().ptr();\\n}\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_DOUBLE = \\\"\\\"\\\"\\\\\\nreturn PyFloat_FromDouble((double) prop);\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_BOOL = \\\"\\\"\\\"\\\\\\nif (prop) {\\n Py_RETURN_TRUE;\\n} else {\\n Py_RETURN_FALSE;\\n}\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_STRING = \\\"\\\"\\\"\\\\\\nreturn PyUnicode_FromStringAndSize(prop.data(), prop.size());\\n\\\"\\\"\\\"\\n\\nGETTER_BODY_SCALAR = \\\"\\\"\\\"\\\\\\nif (prop.isComplex()) {\\n auto cprop = prop.to>();\\n return PyComplex_FromDoubles(cprop.real(), cprop.imag());\\n} else if (prop.isFloatingPoint()) {\\n return PyFloat_FromDouble(prop.to());\\n} else if (prop.isIntegral(/*includeBool=*/false)) {\\n return PyLong_FromLong(prop.to());\\n} else if (prop.isBoolean()) {\\n if (prop.to()) {\\n Py_RETURN_TRUE;\\n } else {\\n Py_RETURN_FALSE;\\n }\\n} else {\\n PyErr_SetString(PyExc_RuntimeError, \\\"Unknown scalar type\\\");\\n return nullptr;\\n}\\n\\\"\\\"\\\"\\n\\n\\nGETTER_BODY_VEC_SCALAR = \\\"\\\"\\\"\\\\\\nPyObject* tup = PyTuple_New((Py_ssize_t) prop.size());\\nfor (auto i: c10::irange(prop.size())) {\\n if (prop[i].isComplex()) {\\n auto cprop = prop[i].to>();\\n PyTuple_SetItem(tup, (Py_ssize_t) i, PyComplex_FromDoubles(cprop.real(), cprop.imag()));\\n } else if (prop[i].isFloatingPoint()) {\\n auto double_prop = prop[i].to();\\n PyTuple_SetItem(tup, (Py_ssize_t) i, PyFloat_FromDouble(double_prop));\\n } else if (prop[i].isIntegral(/*includeBool=*/false)) {\\n auto long_prop = prop[i].to();\\n PyTuple_SetItem(tup, (Py_ssize_t) i, PyLong_FromLong(long_prop));\\n } else if (prop[i].isBoolean()) {\\n if (prop[i].to()) {\\n PyTuple_SetItem(tup, (Py_ssize_t) i, Py_True);\\n } else {\\n PyTuple_SetItem(tup, (Py_ssize_t) i, Py_False);\\n }\\n } else {\\n PyErr_SetString(PyExc_RuntimeError, \\\"Unknown scalar type\\\");\\n return nullptr;\\n }\\n}\\nreturn tup;\\n\\\"\\\"\\\"\\n\\n\\nMISC_GETTER_DEFS = {\\n OptionalCType(BaseCType(longT)): (GETTER_DEFINITION_OPT, GETTER_BODY_INT64_T),\\n OptionalCType(BaseCType(SymIntT)): (GETTER_DEFINITION_OPT, GETTER_BODY_SYMINT),\\n BaseCType(doubleT): (GETTER_DEFINITION, GETTER_BODY_DOUBLE),\\n OptionalCType(BaseCType(doubleT)): (GETTER_DEFINITION_OPT, GETTER_BODY_DOUBLE),\\n BaseCType(boolT): (GETTER_DEFINITION, GETTER_BODY_BOOL),\\n BaseCType(scalarT): (GETTER_DEFINITION, GETTER_BODY_SCALAR),\\n OptionalCType(BaseCType(scalarT)): (GETTER_DEFINITION_OPT, GETTER_BODY_SCALAR),\\n}\\n\\n# These functions have backwards which cannot be traced, and so must have\\n# their backward functions traced opaquely.\\n# VIEW_FUNCTIONS are not traceable because they use as_strided, which\\n# has an untraceable backwards, see\\n# https://github.com/pytorch/pytorch/issues/4250\\n# TODO: This is probably not exhaustive, but it's a start\\nUNTRACEABLE_FUNCTIONS = VIEW_FUNCTIONS\\n\\n\\ndef get_infos_with_derivatives_list(\\n differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]]\\n) -> list[DifferentiabilityInfo]:\\n diff_info_list = [\\n info\\n for diffinfo_dict in differentiability_infos.values()\\n for info in diffinfo_dict.values()\\n ]\\n\\n return list(filter(lambda info: info.args_with_derivatives, diff_info_list))\\n\\n\\ndef gen_autograd_functions_lib(\\n out: str,\\n differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]],\\n template_path: str,\\n) -> None:\\n \\\"\\\"\\\"Functions.h and Functions.cpp body\\n\\n These contain the auto-generated subclasses of torch::autograd::Node\\n for each every differentiable torch function.\\n \\\"\\\"\\\"\\n\\n # get a 1D list of diffinfos, we do not need them to be per FunctionSchema/DispatchKey here\\n # infos with the diff dispatchkeys but the same name will still be in the same shard.\\n infos = get_infos_with_derivatives_list(differentiability_infos)\\n declarations = [process_function(f, FUNCTION_DECLARATION) for f in infos]\\n definitions = [process_function(f, FUNCTION_DEFINITION) for f in infos]\\n\\n file_basename = \\\"Functions\\\"\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n for suffix in [\\\".h\\\", \\\".cpp\\\"]:\\n fname = file_basename + suffix\\n fm.write_with_template(\\n fname,\\n fname,\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/\\\"\\n + fname,\\n \\\"autograd_function_declarations\\\": declarations,\\n \\\"autograd_function_definitions\\\": definitions,\\n },\\n )\\n\\n\\ndef gen_autograd_functions_python(\\n out: str,\\n differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]],\\n template_path: str,\\n) -> None:\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n num_shards = 5\\n fm.write(\\n \\\"python_functions.h\\\",\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/python_functions.h\\\",\\n \\\"shard_forward_declare\\\": [\\n f\\\"void initialize_autogenerated_functions_{i}(PyObject* module);\\\"\\n for i in range(num_shards)\\n ],\\n \\\"shard_call\\\": [\\n f\\\"initialize_autogenerated_functions_{i}(module);\\\"\\n for i in range(num_shards)\\n ],\\n },\\n )\\n\\n # get a 1D list of diffinfos, we do not need them to be per FunctionSchema/DispatchKey here\\n # infos with the diff dispatchkeys but the same name will still be in the same shard.\\n infos = get_infos_with_derivatives_list(differentiability_infos)\\n fm.write_sharded(\\n \\\"python_functions.cpp\\\",\\n infos,\\n key_fn=lambda info: info.name,\\n base_env={\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/python_functions.cpp\\\",\\n },\\n env_callable=lambda info: {\\n \\\"py_function_initializers\\\": [\\n process_function(info, PY_FUNCTION_DEFINITION)\\n ],\\n \\\"py_function_props_and_getters\\\": [\\n process_function(info, PY_FUNCTION_PROPS_AND_GETTERS)\\n ],\\n },\\n num_shards=num_shards,\\n sharded_keys={\\\"py_function_initializers\\\", \\\"py_function_props_and_getters\\\"},\\n )\\n\\n\\ndef process_function(info: DifferentiabilityInfo, template: CodeTemplate) -> str:\\n saved_variables: list[str] = []\\n release_variables: list[str] = []\\n saved_list_sizes: list[str] = []\\n unpack: list[str] = []\\n asserts: list[str] = []\\n compute_index_ranges: list[str] = []\\n getter_definitions: list[str] = []\\n py_getsetdef_structs: list[str] = []\\n compiled_args: list[str] = []\\n apply_with_saved_before: list[str] = []\\n apply_with_saved_after: list[str] = []\\n\\n for arg in info.args_with_derivatives:\\n if arg.type in TENSOR_LIST_LIKE_CTYPES:\\n size = f\\\"{arg.name}_size_\\\"\\n saved_list_sizes.append(f\\\"size_t {arg.name}_size_;\\\")\\n else:\\n size = \\\"1\\\"\\n compute_index_ranges.append(f\\\"auto {arg.name}_ix = gen.range({size});\\\")\\n\\n def save_var(var: SavedAttribute, is_output: bool) -> None:\\n name = var.nctype.name\\n type = var.nctype.type\\n should_append_getsetdef = True\\n should_append_raw_getsetdef = False\\n visit_name = name\\n uses_cpp_saved_variable_cls = False\\n\\n if (\\n type == BaseCType(tensorT)\\n or type == OptionalCType(BaseCType(tensorT))\\n or type == MutRefCType(OptionalCType(BaseCType(tensorT)))\\n or (type == BaseCType(scalarT) and is_output)\\n ):\\n uses_cpp_saved_variable_cls = True\\n saved_variables.append(f\\\"SavedVariable {name}_;\\\")\\n release_variables.append(f\\\"{name}_.reset_data();\\\")\\n ptr = \\\"shared_from_this()\\\" if is_output else \\\"\\\"\\n unpack.append(f\\\"auto {name} = {name}_.unpack({ptr});\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_SAVEDVAR.substitute(\\n op=info.op, name=name, body=GETTER_BODY_SAVEDVAR\\n )\\n )\\n getter_definitions.append(\\n GETTER_DEFINITION_RAW_SAVEDVAR.substitute(\\n op=info.op, name=name, body=GETTER_BODY_RAW_SAVEDVAR\\n )\\n )\\n should_append_raw_getsetdef = True\\n visit_name = f\\\"{name}_\\\"\\n elif (\\n type == BaseCType(tensorListT)\\n or type == BaseCType(iTensorListRefT)\\n or type == VectorCType(BaseCType(tensorT))\\n ):\\n # note(crcrpar): [nuanced return type of out-of-place foreach functions]\\n # When an out-of-place foreach function whose return signature is `Tensor[]`\\n # spells out its backward definitions in `derivatives.yaml`, and some of them depend on\\n # `result`, `result`'s type is interpreted and treated as `std::vector`.\\n # An out-of-place foreach whose backwards rely on their output doesn't suffer from this\\n # difference if the definitions are codegen'ed.\\n # This special case is needed for `_foreach_pow.List` and `_foreach_pow.ScalarAndTensor`\\n # as of https://github.com/pytorch/pytorch/pull/105504.\\n if type == VectorCType(BaseCType(tensorT)):\\n assert (\\n info.func.func.name.name.base.startswith(\\\"_foreach\\\") and is_output\\n )\\n uses_cpp_saved_variable_cls = True\\n saved_variables.append(f\\\"std::vector {name}_;\\\")\\n saved_variables.append(f\\\"bool {name}_released_ = false;\\\")\\n # Just clear() is sufficient, we don't need to loop and clear each variable.\\n # Because the SavedVariable owns a tensor and a grad_fn, removing the SavedVariable makes them go away as well.\\n release_variables.append(f\\\"{name}_.clear();\\\")\\n release_variables.append(f\\\"{name}_released_ = true;\\\")\\n ptr = \\\"shared_from_this()\\\" if is_output else \\\"nullptr\\\"\\n unpack.append(f\\\"auto {name} = unpack_list({name}_, {ptr});\\\")\\n asserts.append(f\\\"TORCH_CHECK(!{name}_released_, ERR_BACKWARD_TWICE);\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_VEC_SAVEDVAR.substitute(\\n op=info.op, name=name, body=GETTER_BODY_VEC_SAVEDVAR\\n )\\n )\\n getter_definitions.append(\\n GETTER_DEFINITION_RAW_VEC_SAVEDVAR.substitute(\\n op=info.op, name=name, body=GETTER_BODY_RAW_VEC_SAVEDVAR\\n )\\n )\\n should_append_raw_getsetdef = True\\n visit_name = f\\\"{name}_\\\"\\n elif type == ListCType(OptionalCType(BaseCType(tensorT))):\\n uses_cpp_saved_variable_cls = True\\n saved_variables.append(f\\\"std::vector {name}_;\\\")\\n saved_variables.append(f\\\"bool {name}_released_ = false;\\\")\\n # Just clear() is sufficient, we don't need to loop and clear each variable.\\n # Because the SavedVariable owns a tensor and a grad_fn, removing the SavedVariable makes them go away as well.\\n release_variables.append(f\\\"{name}_.clear();\\\")\\n release_variables.append(f\\\"{name}_released_ = true;\\\")\\n unpack.append(f\\\"auto {name} = unpack_opt_list({name}_);\\\")\\n asserts.append(f\\\"TORCH_CHECK(!{name}_released_, ERR_BACKWARD_TWICE);\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_VEC_SAVEDVAR.substitute(\\n op=info.op, name=name, body=GETTER_BODY_VEC_SAVEDVAR\\n )\\n )\\n getter_definitions.append(\\n GETTER_DEFINITION_RAW_VEC_SAVEDVAR.substitute(\\n op=info.op, name=name, body=GETTER_BODY_RAW_VEC_SAVEDVAR\\n )\\n )\\n should_append_raw_getsetdef = True\\n visit_name = f\\\"{name}_\\\"\\n elif type == BaseCType(intArrayRefT):\\n saved_variables.append(f\\\"std::vector {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION.substitute(\\n op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG\\n )\\n )\\n elif type == BaseCType(symIntArrayRefT):\\n saved_variables.append(f\\\"std::vector {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION.substitute(\\n op=info.op, name=name, body=GETTER_BODY_ARRAYREF_SYMINT\\n )\\n )\\n elif type == BaseCType(optionalIntArrayRefT):\\n saved_variables.append(f\\\"c10::OptionalArray {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_OPT_ARRAYREF.substitute(\\n op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG\\n )\\n )\\n elif type == BaseCType(optionalSymIntArrayRefT):\\n saved_variables.append(f\\\"c10::OptionalArray {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_OPT_ARRAYREF.substitute(\\n op=info.op, name=name, body=GETTER_BODY_ARRAYREF_SYMINT\\n )\\n )\\n elif type == OptionalCType(BaseCType(intArrayRefT)):\\n saved_variables.append(f\\\"c10::OptionalArray {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_OPT_ARRAYREF.substitute(\\n op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG\\n )\\n )\\n elif type == OptionalCType(BaseCType(symIntArrayRefT)):\\n saved_variables.append(f\\\"c10::OptionalArray {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_OPT_ARRAYREF.substitute(\\n op=info.op, name=name, body=GETTER_BODY_ARRAYREF_SYMINT\\n )\\n )\\n elif type == OptionalCType(ArrayRefCType(BaseCType(doubleT))):\\n saved_variables.append(f\\\"c10::OptionalArray {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_OPT_ARRAYREF.substitute(\\n op=info.op, name=name, body=GETTER_BODY_ARRAYREF_DOUBLE\\n )\\n )\\n elif type == BaseCType(longT):\\n saved_variables.append(f\\\"{type.cpp_type()} {name} = 0;\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION.substitute(\\n op=info.op, name=name, body=GETTER_BODY_INT64_T\\n )\\n )\\n elif type == BaseCType(SymIntT):\\n saved_variables.append(f\\\"c10::SymInt {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION.substitute(\\n op=info.op, name=name, body=GETTER_BODY_SYMINT\\n )\\n )\\n elif type == BaseCType(stringT):\\n saved_variables.append(f\\\"std::string {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION.substitute(\\n op=info.op, name=name, body=GETTER_BODY_STRING\\n )\\n )\\n elif type == OptionalCType(BaseCType(stringT)):\\n saved_variables.append(f\\\"std::optional {name};\\\")\\n getter_definitions.append(\\n GETTER_DEFINITION_OPT.substitute(\\n op=info.op, name=name, body=GETTER_BODY_STRING\\n )\\n )\\n elif type == ArrayRefCType(\\n elem=BaseCType(type=BaseCppType(ns=\\\"at\\\", name=\\\"Scalar\\\"))\\n ):\\n saved_variables.append(f\\\"std::vector {name};\\\")\\n saved_variables.append(f\\\"bool {name}_released_ = false;\\\")\\n # Just clear() is sufficient, we don't need to loop and clear each variable.\\n # Because the SavedVariable owns a tensor and a grad_fn, removing the SavedVariable makes them go away as well.\\n release_variables.append(f\\\"{name}.clear();\\\")\\n # release_variables.append(f\\\"{name}_released_ = true;\\\")\\n # unpack.append(f\\\"auto {name} = unpack_list({name}_);\\\")\\n # asserts.append(f\\\"TORCH_CHECK(!{name}_released_, ERR_BACKWARD_TWICE);\\\")\\n getter_definitions.append(\\n CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nPyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) {\\n HANDLE_TH_ERRORS\\n const auto *node = static_cast<${op}*>(self->cdata.get());\\n const auto& prop = node->${name};\\n if (node->${name}_released_) {\\n PyErr_SetString(PyExc_RuntimeError, ERR_BACKWARD_TWICE);\\n return nullptr;\\n }\\n ${body}\\n END_HANDLE_TH_ERRORS\\n}\\n \\\"\\\"\\\"\\n ).substitute(\\n op=info.op,\\n name=name,\\n body=GETTER_BODY_VEC_SCALAR,\\n )\\n )\\n else:\\n # Check for indicators that you're putting a non-owning reference\\n # into the saved variable field. If this is spuriously firing,\\n # edit this field. Otherwise, you probably need to add a case\\n # above.\\n assert (\\n \\\"ref\\\" not in type.cpp_type().lower()\\n and \\\"view\\\" not in type.cpp_type().lower()\\n and \\\"*\\\" not in type.cpp_type()\\n and \\\"&\\\" not in type.cpp_type()\\n ), f\\\"{type.cpp_type()} looks like it contains a non-owning reference\\\"\\n saved_variables.append(f\\\"{type.cpp_type()} {name};\\\")\\n\\n if type in MISC_GETTER_DEFS:\\n getter_def, body = MISC_GETTER_DEFS[type]\\n getter_definitions.append(\\n getter_def.substitute(op=info.op, name=name, body=body)\\n )\\n else:\\n # Types we don't expose python bindings to yet:\\n # TypeAndSize, at::ScalarType, TensorOptions, TensorGeometry,\\n # std::vector>, std::vector\\n should_append_getsetdef = False\\n\\n if should_append_getsetdef:\\n py_getsetdef_structs.append(\\n PY_GETSETDEF_STRUCT.substitute(op=info.op, name=name)\\n )\\n if should_append_raw_getsetdef:\\n py_getsetdef_structs.append(\\n PY_RAW_GETSETDEF_STRUCT.substitute(op=info.op, name=name)\\n )\\n\\n if uses_cpp_saved_variable_cls:\\n compiled_args.append(\\n f\\\"args.collect({visit_name}, {'true' if is_output else 'false'});\\\"\\n )\\n else:\\n compiled_args.append(f\\\"args.collect({visit_name});\\\")\\n apply_with_saved_before.append(f\\\"saved.before({visit_name});\\\")\\n apply_with_saved_after.append(f\\\"saved.after({visit_name});\\\")\\n\\n for var in sorted(info.all_saved_inputs, key=lambda sa: str(sa.nctype.name)):\\n save_var(var, is_output=False)\\n for var in sorted(info.all_saved_outputs, key=lambda sa: str(sa.nctype.name)):\\n save_var(var, is_output=True)\\n\\n # lock the mutex when we release variables and in Node::apply to protect thread safety\\n # see Note [Thread Safety on Autograd Node]\\n if len(release_variables) > 0:\\n thread_lock = \\\"std::lock_guard lock(mutex_);\\\"\\n else:\\n thread_lock = \\\"\\\"\\n\\n if uses_retain_variables(info):\\n will_release_variables = WILL_RELEASE_VARIABLES.substitute()\\n else:\\n will_release_variables = \\\"\\\"\\n\\n body: list[str] = []\\n\\n if uses_single_grad(info):\\n body.append(\\\"const auto& grad = grads[0];\\\")\\n else:\\n # Generate aliases for gradients named for returned values.\\n body.extend(\\n f\\\"const auto& {name} = grads[{info.available_named_gradients.index(name)}];\\\"\\n for name in sorted(info.used_named_gradients)\\n )\\n\\n def emit_derivative(\\n derivative: Derivative,\\n args_with_derivatives: Sequence[Binding],\\n ) -> tuple[bool, str]:\\n formula = derivative.formula\\n var_names = derivative.var_names\\n if len(var_names) == 1:\\n checks_any_grad_defined = False\\n if \\\"not_implemented\\\" not in formula:\\n matching_args = [\\n arg for arg in args_with_derivatives if arg.name == var_names[0]\\n ]\\n if len(matching_args) == 1:\\n # We can add undefined grad support if the input variable is a Tensor\\n arg = matching_args[0]\\n if isinstance(arg.argument, Argument) and str(\\n arg.argument.type\\n ) in (\\\"Tensor\\\", \\\"Tensor?\\\"):\\n formula = \\\"any_grad_defined ? (\\\" + formula + \\\") : Tensor()\\\"\\n checks_any_grad_defined = True\\n if info.name.startswith(\\\"_foreach_\\\"):\\n derivative_template = DERIVATIVE_SINGLE_FOREACH\\n else:\\n derivative_template = DERIVATIVE_SINGLE\\n return (\\n checks_any_grad_defined,\\n derivative_template.substitute(name=var_names[0], derivative=formula),\\n )\\n else:\\n if \\\"grad_input_mask\\\" in formula:\\n masks = [\\n f\\\"task_should_compute_output({{ {n}_ix }}),\\\" for n in var_names\\n ]\\n grad_input_mask = GRAD_INPUT_MASK.substitute(\\n masks=masks, n=len(var_names)\\n )\\n else:\\n grad_input_mask = \\\"\\\"\\n idx_ranges = \\\", \\\".join(f\\\"{n}_ix\\\" for n in var_names)\\n copy_ranges: list[str] = []\\n for i, n in enumerate(var_names):\\n copy_ranges.append(DERIVATIVE_MULTI_COPY_RANGE.substitute(name=n, i=i))\\n return False, DERIVATIVE_MULTI.substitute(\\n idx_ranges=idx_ranges,\\n copy_ranges=copy_ranges,\\n derivative=formula,\\n grad_input_mask=grad_input_mask,\\n )\\n\\n body.extend(unpack)\\n need_any_grad_defined_var = False\\n for derivative in info.derivatives:\\n checks_any_grad_defined, derivative_text = emit_derivative(\\n derivative, info.args_with_derivatives\\n )\\n body.append(derivative_text)\\n need_any_grad_defined_var |= checks_any_grad_defined\\n # Since single-output derivative formulas need to check if grads are\\n # defined, only perform the check once, before all the formulas\\n if need_any_grad_defined_var:\\n body.insert(\\n -len(info.derivatives),\\n \\\"bool any_grad_defined = any_variable_defined(grads);\\\",\\n )\\n\\n if info.name in UNTRACEABLE_FUNCTIONS:\\n superclass = \\\"Node\\\"\\n else:\\n superclass = \\\"TraceableFunction\\\"\\n\\n all_getsetdef_structs = (\\n \\\",\\\\n\\\".join(py_getsetdef_structs) + \\\",\\\" if len(py_getsetdef_structs) != 0 else \\\"\\\"\\n )\\n all_getter_definitions = \\\"\\\\n\\\".join(getter_definitions)\\n\\n return template.substitute(\\n op=info.op,\\n compute_index_ranges=compute_index_ranges,\\n saved_variables=saved_variables,\\n release_variables=release_variables,\\n saved_list_sizes=saved_list_sizes,\\n asserts=asserts,\\n thread_lock=thread_lock,\\n will_release_variables=will_release_variables,\\n body=body,\\n superclass=superclass,\\n all_getter_definitions=all_getter_definitions,\\n all_getsetdef_structs=all_getsetdef_structs,\\n compiled_args=compiled_args,\\n apply_with_saved_before=apply_with_saved_before,\\n apply_with_saved_after=apply_with_saved_after,\\n )\\n\\n\\n# Generates Python bindings for ATen functions\\n#\\n# The bindings are generated as methods on python_variable or functions on the\\n# torch._C._nn. torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._sparse\\n# or torch._C._special objects.\\n#\\n\\n# Code tries to stick to the following rules:\\n#\\n# - templates should be colocated with the functions that use them.\\n# no templates are currently shared between functions, but if that\\n# happens, maybe put the template with the first one\\n#\\n# - don't use environment dictionaries when calling template.substitute().\\n# pass named arguments directly for everything, otherwise it's much too\\n# hard to track what's actually being used and by who\\n#\\n# - colocate any new hacks/adjustments with existing ones of the same kind.\\n# ideally in a data structure rather than code if possible. See e.g.\\n# SCHEMA_DEFAULT_CONVERSION_HACKS, etc.\\n#\\n# - similarly, conversions from one format to another should ideally happen\\n# all at once in a single place.\\n#\\n# - no nontrivial nested functions. couple-liners are ok but please no more.\\n# especially avoid functions that read/write outer variables defined far away.\\n#\\n# - raise RuntimeError instead of asserting, and put as much\\n# information as is available into the message. I.e. no need to\\n# plumb in new params whose only purpose is to fill out an error\\n# message, but use what's there\\n#\\n\\nfrom __future__ import annotations\\n\\nimport itertools\\nimport re\\nfrom collections import defaultdict\\nfrom typing import Callable, Iterable, Sequence\\n\\nimport yaml\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.python import (\\n arg_parser_output_exprs,\\n cpp_dispatch_exprs,\\n cpp_dispatch_target,\\n dispatch_lambda_args,\\n dispatch_lambda_exprs,\\n dispatch_lambda_return_str,\\n has_tensor_options,\\n PythonSignature,\\n PythonSignatureDeprecated,\\n PythonSignatureGroup,\\n PythonSignatureNativeFunctionPair,\\n signature,\\n signature_from_schema,\\n structseq_fieldnames,\\n)\\nfrom torchgen.code_template import CodeTemplate\\nfrom torchgen.context import with_native_function\\nfrom torchgen.gen import cpp_string, parse_native_yaml, parse_tags_yaml\\nfrom torchgen.model import (\\n Argument,\\n BaseOperatorName,\\n FunctionSchema,\\n NativeFunction,\\n SchemaKind,\\n Type,\\n Variant,\\n)\\nfrom torchgen.utils import FileManager, split_name_params\\nfrom torchgen.yaml_utils import YamlLoader\\n\\nfrom .gen_inplace_or_view_type import is_tensor_list_type\\nfrom .gen_trace_type import should_trace\\n\\n\\n#\\n# declarations blocklist\\n# We skip codegen for these functions, for various reasons.\\n# Future PRs will categorize this list and eliminate or hoist\\n# them out of eager-only codegen.\\n# See https://github.com/pytorch/pytorch/issues/30788\\n#\\n\\n# These functions require manual Python bindings or are not exposed to Python\\n_SKIP_PYTHON_BINDINGS = [\\n \\\"alias\\\",\\n \\\"contiguous\\\",\\n \\\"is_cuda\\\",\\n \\\"is_sparse\\\",\\n \\\"is_sparse_csr\\\",\\n \\\"size\\\",\\n \\\"stride\\\",\\n \\\"sym_size\\\",\\n \\\"sym_stride\\\",\\n \\\"sym_storage_offset\\\",\\n \\\"sym_numel\\\",\\n \\\".*_backward\\\",\\n \\\".*_backward_(out|input|weight|bias)\\\",\\n \\\".*_forward\\\",\\n \\\".*_forward_out\\\",\\n \\\".*_jvp\\\",\\n \\\"_unsafe_view\\\",\\n \\\"tensor\\\",\\n \\\"_?sparse_(coo|compressed|csr|csc|bsr|bsc)_tensor.*\\\",\\n \\\"_range.*\\\",\\n \\\"_sparse_add_out\\\",\\n \\\"_sparse_div.*\\\",\\n \\\"_sparse_mul.*\\\",\\n \\\"_sparse_sub.*\\\",\\n \\\"_sparse_dense_add_out\\\",\\n \\\"index\\\",\\n \\\"index_out\\\",\\n \\\"unique_dim_consecutive\\\",\\n \\\"_cumsum.*\\\",\\n \\\"_cumprod.*\\\",\\n \\\"_sum.*\\\",\\n \\\"_prod.*\\\",\\n \\\"_th_.*\\\",\\n \\\"_thnn_.*\\\",\\n \\\"range.*\\\",\\n \\\"_solve.*\\\",\\n \\\"_inverse.*\\\",\\n \\\"_cholesky.*\\\",\\n \\\"_triangular_solve.*\\\",\\n \\\"_qr.*\\\",\\n \\\"_svd.*\\\",\\n \\\"slice\\\",\\n \\\"item\\\",\\n \\\"_local_scalar_dense\\\",\\n \\\"to\\\",\\n \\\"_to_copy\\\",\\n \\\"_to_copy_out\\\",\\n \\\"_reshape_copy\\\",\\n \\\"_reshape_copy_out\\\",\\n \\\"copy_sparse_to_sparse_\\\",\\n \\\"copy_\\\",\\n \\\"_foreach_copy\\\",\\n \\\"numpy_T\\\",\\n \\\"matrix_H\\\",\\n \\\"mT\\\",\\n \\\"mH\\\", # these need to be an attributes in Python, not functions\\n \\\"nonzero(_(out|numpy))?\\\",\\n \\\"set_data\\\",\\n \\\".*_overrideable\\\", # overrideable functions for backend extension\\n \\\"data\\\",\\n \\\"is_leaf\\\",\\n \\\"output_nr\\\",\\n \\\"_version\\\",\\n \\\"requires_grad_\\\",\\n \\\"retains_grad\\\",\\n \\\"set_\\\",\\n \\\"_fw_primal\\\",\\n \\\"fake_quantize_per_tensor_affine_cachemask\\\",\\n \\\"fake_quantize_per_channel_affine_cachemask\\\",\\n \\\"_new_zeros_with_same_feature_meta\\\",\\n \\\"_has_same_storage_numel\\\", # used for forward AD internals\\n \\\"_reshape_alias\\\",\\n \\\"replace_\\\", # only used by the functionalization pass, doesn't need to be exposed to python\\n \\\"copy\\\", # only used by the functionalization pass\\n \\\"fill.Tensor\\\", # only used by the functionalization pass\\n \\\"fill.Scalar\\\", # only used by the functionalization pass\\n \\\"lift.*\\\",\\n \\\"normal_functional\\\", # only used by the functionalization pass\\n \\\"nbytes\\\",\\n \\\"itemsize\\\",\\n \\\"_batch_norm_with_update\\\",\\n \\\"_batch_norm_with_update_out\\\",\\n \\\"_batch_norm_no_update\\\",\\n]\\n\\nSKIP_PYTHON_BINDINGS = [\\n re.compile(rf\\\"^{pattern}$\\\") for pattern in _SKIP_PYTHON_BINDINGS\\n]\\n\\n# These function signatures are not exposed to Python. Note that this signature\\n# list does not support regex.\\nSKIP_PYTHON_BINDINGS_SIGNATURES = [\\n \\\"add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor\\\",\\n \\\"add_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)\\\",\\n \\\"sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor\\\",\\n \\\"sub_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)\\\",\\n \\\"mul.Scalar(Tensor self, Scalar other) -> Tensor\\\",\\n \\\"mul_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)\\\",\\n \\\"div.Scalar(Tensor self, Scalar other) -> Tensor\\\",\\n \\\"div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)\\\",\\n]\\n\\n\\n@with_native_function\\ndef should_generate_py_binding(f: NativeFunction) -> bool:\\n # NativeFunctions that are entirely code-generated should not get python bindings\\n # because these codegen implementations are often inefficient. A handful of\\n # view_copy style ops were exposed accidentally when they were handwritten and now\\n # that we are moving them to codegen for bc reasons we need to keep them exposed in\\n # python.\\n if \\\"generated\\\" in f.tags and \\\"view_copy\\\" not in f.tags:\\n return False\\n\\n name = cpp.name(f.func)\\n for skip_regex in SKIP_PYTHON_BINDINGS:\\n if skip_regex.match(name):\\n return False\\n\\n signature = str(f.func)\\n for pattern in SKIP_PYTHON_BINDINGS_SIGNATURES:\\n if pattern == signature:\\n return False\\n return True\\n\\n\\ndef get_pycname(name: BaseOperatorName) -> str:\\n return f\\\"THPVariable_{name}\\\"\\n\\n\\ndef is_noarg(overloads: Sequence[PythonSignatureNativeFunctionPair]) -> bool:\\n return len(overloads) == 1 and overloads[0].signature.arguments_count() == 0\\n\\n\\ndef is_py_variable_method(f: NativeFunction) -> bool:\\n return f.python_module is None and Variant.method in f.variants\\n\\n\\ndef is_py_torch_function(f: NativeFunction) -> bool:\\n return f.python_module is None and Variant.function in f.variants\\n\\n\\ndef is_py_nn_function(f: NativeFunction) -> bool:\\n return f.python_module == \\\"nn\\\"\\n\\n\\ndef is_py_fft_function(f: NativeFunction) -> bool:\\n return f.python_module == \\\"fft\\\"\\n\\n\\ndef is_py_linalg_function(f: NativeFunction) -> bool:\\n return f.python_module == \\\"linalg\\\"\\n\\n\\ndef is_py_nested_function(f: NativeFunction) -> bool:\\n return f.python_module == \\\"nested\\\"\\n\\n\\ndef is_py_sparse_function(f: NativeFunction) -> bool:\\n return f.python_module == \\\"sparse\\\"\\n\\n\\ndef is_py_special_function(f: NativeFunction) -> bool:\\n return f.python_module == \\\"special\\\"\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Main Function\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef gen(\\n out: str,\\n native_yaml_path: str,\\n tags_yaml_path: str,\\n deprecated_yaml_path: str,\\n template_path: str,\\n *,\\n symint: bool = True,\\n) -> None:\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n native_functions = parse_native_yaml(\\n native_yaml_path, tags_yaml_path\\n ).native_functions\\n native_functions = list(filter(should_generate_py_binding, native_functions))\\n\\n methods = load_signatures(native_functions, deprecated_yaml_path, method=True)\\n create_python_bindings(\\n fm,\\n methods,\\n is_py_variable_method,\\n None,\\n \\\"python_variable_methods.cpp\\\",\\n method=True,\\n symint=symint,\\n )\\n\\n # NOTE: num_shards here must be synced with gatherTorchFunctions in\\n # torch/csrc/autograd/python_torch_functions_manual.cpp\\n functions = load_signatures(native_functions, deprecated_yaml_path, method=False)\\n create_python_bindings_sharded(\\n fm,\\n functions,\\n is_py_torch_function,\\n \\\"torch\\\",\\n \\\"python_torch_functions.cpp\\\",\\n method=False,\\n num_shards=3,\\n symint=symint,\\n )\\n\\n create_python_bindings(\\n fm,\\n functions,\\n is_py_nn_function,\\n \\\"torch.nn\\\",\\n \\\"python_nn_functions.cpp\\\",\\n method=False,\\n symint=symint,\\n )\\n\\n create_python_bindings(\\n fm,\\n functions,\\n is_py_fft_function,\\n \\\"torch.fft\\\",\\n \\\"python_fft_functions.cpp\\\",\\n method=False,\\n symint=symint,\\n )\\n\\n create_python_bindings(\\n fm,\\n functions,\\n is_py_linalg_function,\\n \\\"torch.linalg\\\",\\n \\\"python_linalg_functions.cpp\\\",\\n method=False,\\n symint=symint,\\n )\\n\\n create_python_bindings(\\n fm,\\n functions,\\n is_py_nested_function,\\n \\\"torch.nested\\\",\\n \\\"python_nested_functions.cpp\\\",\\n method=False,\\n )\\n\\n create_python_bindings(\\n fm,\\n functions,\\n is_py_sparse_function,\\n \\\"torch.sparse\\\",\\n \\\"python_sparse_functions.cpp\\\",\\n method=False,\\n symint=symint,\\n )\\n\\n create_python_bindings(\\n fm,\\n functions,\\n is_py_special_function,\\n \\\"torch.special\\\",\\n \\\"python_special_functions.cpp\\\",\\n method=False,\\n symint=symint,\\n )\\n\\n # Currently, we only use `functions` to generate `return_types` bindings.\\n # All methods which return structseq have function variant at this point.\\n # If any method only operator with structseq is added in the future,\\n # we will have to address that.\\n create_python_return_type_bindings(\\n fm, functions, lambda fn: True, \\\"python_return_types.cpp\\\"\\n )\\n create_python_return_type_bindings_header(\\n fm, functions, lambda fn: True, \\\"python_return_types.h\\\"\\n )\\n\\n valid_tags = parse_tags_yaml(tags_yaml_path)\\n\\n def gen_tags_enum() -> dict[str, str]:\\n return {\\n \\\"enum_of_valid_tags\\\": (\\n \\\"\\\".join(\\n [f'\\\\n.value(\\\"{tag}\\\", at::Tag::{tag})' for tag in sorted(valid_tags)]\\n )\\n )\\n }\\n\\n fm.write(\\\"python_enum_tag.cpp\\\", gen_tags_enum)\\n\\n\\ndef group_filter_overloads(\\n pairs: Sequence[PythonSignatureNativeFunctionPair],\\n pred: Callable[[NativeFunction], bool],\\n) -> dict[BaseOperatorName, list[PythonSignatureNativeFunctionPair]]:\\n grouped: dict[\\n BaseOperatorName, list[PythonSignatureNativeFunctionPair]\\n ] = defaultdict(list)\\n for pair in pairs:\\n if pred(pair.function):\\n grouped[pair.function.func.name.name].append(pair)\\n return grouped\\n\\n\\ndef create_python_bindings(\\n fm: FileManager,\\n pairs: Sequence[PythonSignatureNativeFunctionPair],\\n pred: Callable[[NativeFunction], bool],\\n module: str | None,\\n filename: str,\\n *,\\n method: bool,\\n symint: bool = True,\\n) -> None:\\n \\\"\\\"\\\"Generates Python bindings to ATen functions\\\"\\\"\\\"\\n py_methods: list[str] = []\\n ops_headers: list[str] = []\\n py_method_defs: list[str] = []\\n py_forwards: list[str] = []\\n\\n grouped = group_filter_overloads(pairs, pred)\\n\\n for name in sorted(grouped.keys(), key=str):\\n overloads = grouped[name]\\n py_methods.append(\\n method_impl(name, module, overloads, method=method, symint=symint)\\n )\\n py_method_defs.append(method_def(name, module, overloads, method=method))\\n py_forwards.extend(forward_decls(name, overloads, method=method))\\n ops_headers.append(f\\\"#include \\\")\\n\\n fm.write_with_template(\\n filename,\\n filename,\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/{filename}\\\",\\n \\\"ops_headers\\\": ops_headers,\\n \\\"py_forwards\\\": py_forwards,\\n \\\"py_methods\\\": py_methods,\\n \\\"py_method_defs\\\": py_method_defs,\\n },\\n )\\n\\n\\ndef create_python_return_type_bindings(\\n fm: FileManager,\\n pairs: Sequence[PythonSignatureNativeFunctionPair],\\n pred: Callable[[NativeFunction], bool],\\n filename: str,\\n) -> None:\\n \\\"\\\"\\\"\\n Generate function to initialize and return named tuple for native functions\\n which returns named tuple and registration invocations in `python_return_types.cpp`.\\n \\\"\\\"\\\"\\n py_return_types_definition: list[str] = []\\n py_return_types_registrations: list[str] = []\\n\\n grouped = group_filter_overloads(pairs, pred)\\n\\n for name in sorted(grouped.keys(), key=str):\\n overloads = grouped[name]\\n definitions, registrations = generate_return_type_definition_and_registrations(\\n overloads\\n )\\n py_return_types_definition.append(\\n \\\"\\\" if not definitions else \\\"\\\\n\\\".join(definitions)\\n )\\n py_return_types_registrations.append(\\n \\\"\\\" if not registrations else \\\"\\\\n\\\".join(registrations)\\n )\\n\\n fm.write_with_template(\\n filename,\\n filename,\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/{filename}\\\",\\n \\\"py_return_types\\\": py_return_types_definition,\\n \\\"py_return_types_registrations\\\": py_return_types_registrations,\\n },\\n )\\n\\n\\ndef create_python_return_type_bindings_header(\\n fm: FileManager,\\n pairs: Sequence[PythonSignatureNativeFunctionPair],\\n pred: Callable[[NativeFunction], bool],\\n filename: str,\\n) -> None:\\n \\\"\\\"\\\"\\n Generate function to initialize and return named tuple for native functions\\n which returns named tuple and relevant entry for the map in `python_return_types.cpp`.\\n \\\"\\\"\\\"\\n py_return_types_declarations: list[str] = []\\n\\n grouped = group_filter_overloads(pairs, pred)\\n\\n for name in sorted(grouped.keys(), key=str):\\n overloads = grouped[name]\\n declarations = generate_return_type_declarations(overloads)\\n py_return_types_declarations.append(\\n \\\"\\\" if not declarations else \\\"\\\\n\\\".join(declarations)\\n )\\n\\n fm.write_with_template(\\n filename,\\n filename,\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/{filename}\\\",\\n \\\"py_return_types_declarations\\\": py_return_types_declarations,\\n },\\n )\\n\\n\\ndef create_python_bindings_sharded(\\n fm: FileManager,\\n pairs: Sequence[PythonSignatureNativeFunctionPair],\\n pred: Callable[[NativeFunction], bool],\\n module: str | None,\\n filename: str,\\n *,\\n method: bool,\\n num_shards: int,\\n symint: bool = True,\\n) -> None:\\n \\\"\\\"\\\"Generates Python bindings to ATen functions\\\"\\\"\\\"\\n grouped = group_filter_overloads(pairs, pred)\\n\\n def key_func(\\n kv: tuple[BaseOperatorName, list[PythonSignatureNativeFunctionPair]]\\n ) -> str:\\n return kv[0].base\\n\\n def env_func(\\n kv: tuple[BaseOperatorName, list[PythonSignatureNativeFunctionPair]]\\n ) -> dict[str, list[str]]:\\n name, fn_pairs = kv\\n return {\\n \\\"ops_headers\\\": [f\\\"#include \\\"],\\n \\\"py_forwards\\\": list(forward_decls(name, fn_pairs, method=method)),\\n \\\"py_methods\\\": [\\n method_impl(name, module, fn_pairs, method=method, symint=symint)\\n ],\\n \\\"py_method_defs\\\": [method_def(name, module, fn_pairs, method=method)],\\n }\\n\\n fm.write_sharded(\\n filename,\\n grouped.items(),\\n base_env={\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/{filename}\\\",\\n },\\n key_fn=key_func,\\n env_callable=env_func,\\n num_shards=num_shards,\\n sharded_keys={\\\"ops_headers\\\", \\\"py_forwards\\\", \\\"py_methods\\\", \\\"py_method_defs\\\"},\\n )\\n\\n\\ndef load_signatures(\\n native_functions: list[NativeFunction],\\n deprecated_yaml_path: str,\\n *,\\n method: bool,\\n skip_deprecated: bool = False,\\n pyi: bool = False,\\n) -> Sequence[PythonSignatureNativeFunctionPair]:\\n @with_native_function\\n def gen_signature_pairs(f: NativeFunction) -> PythonSignatureNativeFunctionPair:\\n return PythonSignatureNativeFunctionPair(\\n signature=signature(f, method=method, pyi=pyi),\\n function=f,\\n )\\n\\n pairs = list(map(gen_signature_pairs, native_functions))\\n deprecated = load_deprecated_signatures(\\n pairs, deprecated_yaml_path, method=method, pyi=pyi\\n )\\n return pairs if skip_deprecated else pairs + deprecated\\n\\n\\ndef load_deprecated_signatures(\\n pairs: Sequence[PythonSignatureNativeFunctionPair],\\n deprecated_yaml_path: str,\\n *,\\n method: bool,\\n pyi: bool,\\n) -> list[PythonSignatureNativeFunctionPair]:\\n # The deprecated.yaml doesn't have complete type information, we need\\n # find and leverage the original ATen signature (to which it delegates\\n # the call) to generate the full python signature.\\n # We join the deprecated and the original signatures using type-only form.\\n\\n # group the original ATen signatures by name\\n grouped: dict[str, list[PythonSignatureNativeFunctionPair]] = defaultdict(list)\\n for pair in pairs:\\n grouped[pair.signature.name].append(pair)\\n\\n # find matching original signatures for each deprecated signature\\n results: list[PythonSignatureNativeFunctionPair] = []\\n\\n with open(deprecated_yaml_path) as f:\\n deprecated_defs = yaml.load(f, Loader=YamlLoader)\\n\\n for deprecated in deprecated_defs:\\n schema = FunctionSchema.parse(deprecated[\\\"name\\\"])\\n aten_name, call_args = split_name_params(deprecated[\\\"aten\\\"])\\n is_out = aten_name.endswith(\\\"_out\\\")\\n if is_out:\\n aten_name = aten_name.replace(\\\"_out\\\", \\\"\\\")\\n\\n # HACK: these are fixed constants used to pass the aten function.\\n # The type must be known ahead of time\\n known_constants = {\\n \\\"1\\\": Type.parse(\\\"Scalar\\\"),\\n }\\n schema_args_by_name = {a.name: a for a in schema.arguments.flat_all}\\n for name in call_args:\\n assert (\\n name in schema_args_by_name or name in known_constants\\n ), f\\\"deprecation definiton: Unrecognized value {name}\\\"\\n\\n # Map deprecated signature arguments to their aten signature and test\\n # if the types and alias annotation match.\\n def is_schema_compatible(\\n aten_schema: FunctionSchema,\\n ) -> bool:\\n arguments: Iterable[Argument]\\n if is_out:\\n arguments = itertools.chain(\\n aten_schema.arguments.out, aten_schema.arguments.flat_non_out\\n )\\n else:\\n arguments = aten_schema.arguments.flat_all\\n\\n for i, arg in enumerate(arguments):\\n if i < len(call_args):\\n arg_name = call_args[i]\\n if arg_name in known_constants:\\n schema_type = known_constants[arg_name]\\n schema_annotation = None\\n else:\\n schema_arg = schema_args_by_name[arg_name]\\n schema_type = schema_arg.type\\n schema_annotation = schema_arg.annotation\\n\\n if schema_type != arg.type or schema_annotation != arg.annotation:\\n return False\\n else:\\n if arg.default is None:\\n return False\\n\\n return len(schema.returns) == len(aten_schema.returns) and all(\\n a == b for a, b in zip(schema.returns, aten_schema.returns)\\n )\\n\\n any_schema_found = False\\n for pair in grouped[aten_name]:\\n if not is_schema_compatible(pair.function.func):\\n continue\\n any_schema_found = True\\n\\n python_sig = signature_from_schema(\\n schema,\\n category_override=pair.function.category_override,\\n method=method,\\n pyi=pyi,\\n )\\n\\n results.append(\\n PythonSignatureNativeFunctionPair(\\n signature=PythonSignatureDeprecated(\\n name=python_sig.name,\\n input_args=python_sig.input_args,\\n input_kwargs=python_sig.input_kwargs,\\n output_args=python_sig.output_args,\\n tensor_options_args=python_sig.tensor_options_args,\\n method=python_sig.method,\\n deprecated_schema=schema,\\n deprecated_args_exprs=tuple(call_args),\\n returns=python_sig.returns,\\n ),\\n function=pair.function,\\n )\\n )\\n assert (\\n any_schema_found\\n ), f\\\"No native function with name {aten_name} matched signature:\\\\n {str(schema)}\\\"\\n\\n return results\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Named Tuple Codegen\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\n@with_native_function\\ndef gen_structseq_typename_key(f: NativeFunction) -> str:\\n name = cpp.name(f.func)\\n fieldnames = structseq_fieldnames(f.func.returns)\\n return \\\"_\\\".join([name] + fieldnames)\\n\\n\\ndef emit_structseq_call(\\n overloads: Sequence[PythonSignatureNativeFunctionPair],\\n) -> tuple[list[str], dict[str, str]]:\\n \\\"\\\"\\\"\\n Generate block of named tuple type def inits, and add typeref snippets\\n to declarations that use them\\n \\\"\\\"\\\"\\n typenames: dict[\\n str, str\\n ] = {} # map from unique name + field name lists to typedef name\\n typedefs: list[str] = [] # typedef declarations and init code\\n\\n for overload in overloads:\\n fieldnames = structseq_fieldnames(overload.function.func.returns)\\n if not fieldnames:\\n continue\\n\\n name = cpp.name(overload.function.func) # use @with_native_function?\\n tn_key = gen_structseq_typename_key(overload.function)\\n typename = typenames.get(tn_key)\\n if typename is None:\\n typename = f'NamedTuple{\\\"\\\" if not typedefs else len(typedefs)}'\\n typenames[tn_key] = typename\\n typedefs.append(\\n f\\\"\\\"\\\"\\\\\\nstatic PyTypeObject* {typename} = generated::get_{name}_structseq();\\\"\\\"\\\"\\n )\\n\\n return typedefs, typenames\\n\\n\\ndef generate_return_type_definition_and_registrations(\\n overloads: Sequence[PythonSignatureNativeFunctionPair],\\n) -> tuple[list[str], list[str]]:\\n \\\"\\\"\\\"\\n Generate block of function in `python_return_types.cpp` to initialize\\n and return named tuple for a native function which returns named tuple\\n and registration invocations in same file.\\n \\\"\\\"\\\"\\n typenames: dict[\\n str, str\\n ] = {} # map from unique name + field name lists to typedef name\\n definitions: list[str] = [] # function definition to register the typedef\\n registrations: list[str] = [] # register call for the typedef\\n\\n for overload in overloads:\\n fieldnames = structseq_fieldnames(overload.function.func.returns)\\n if not fieldnames:\\n continue\\n\\n fields = \\\", \\\".join(f'{{\\\"{fn}\\\", \\\"\\\"}}' for fn in fieldnames)\\n\\n name = cpp.name(overload.function.func) # use @with_native_function?\\n tn_key = gen_structseq_typename_key(overload.function)\\n typename = typenames.get(tn_key)\\n\\n if typename is None:\\n typename = f'{name}NamedTuple{\\\"\\\" if not definitions else len(definitions)}'\\n typenames[tn_key] = typename\\n definitions.append(\\n f\\\"\\\"\\\"\\\\\\nPyTypeObject* get_{name}_structseq() {{\\n static PyStructSequence_Field NamedTuple_fields[] = {{ {fields}, {{nullptr}} }};\\n static PyTypeObject {typename};\\n static bool is_initialized = false;\\n static PyStructSequence_Desc desc = {{ \\\"torch.return_types.{name}\\\", nullptr, NamedTuple_fields, {len(fieldnames)} }};\\n if (!is_initialized) {{\\n PyStructSequence_InitType(&{typename}, &desc);\\n {typename}.tp_repr = (reprfunc)torch::utils::returned_structseq_repr;\\n is_initialized = true;\\n }}\\n return &{typename};\\n}}\\n\\\"\\\"\\\"\\n )\\n registrations.append(\\n f'addReturnType(return_types_module, \\\"{name}\\\", generated::get_{name}_structseq());'\\n )\\n\\n return definitions, registrations\\n\\n\\ndef generate_return_type_declarations(\\n overloads: Sequence[PythonSignatureNativeFunctionPair],\\n) -> list[str]:\\n \\\"\\\"\\\"\\n Generate block of function declarations in `python_return_types.h` to initialize\\n and return named tuple for a native function.\\n \\\"\\\"\\\"\\n typenames: dict[\\n str, str\\n ] = {} # map from unique name + field name lists to typedef name\\n declarations: list[str] = [] # function declaration to register the typedef\\n\\n for overload in overloads:\\n fieldnames = structseq_fieldnames(overload.function.func.returns)\\n if not fieldnames:\\n continue\\n\\n name = cpp.name(overload.function.func) # use @with_native_function?\\n tn_key = gen_structseq_typename_key(overload.function)\\n typename = typenames.get(tn_key)\\n\\n if typename is None:\\n typename = (\\n f'{name}NamedTuple{\\\"\\\" if not declarations else len(declarations)}'\\n )\\n typenames[tn_key] = typename\\n declarations.append(f\\\"PyTypeObject* get_{name}_structseq();\\\")\\n\\n return declarations\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Method Impl Codegen\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n# python binding for all overloads of a particular function/method\\nPY_VARIABLE_METHOD_VARARGS = CodeTemplate(\\n r\\\"\\\"\\\"\\\\\\n// ${name}\\nstatic PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs)\\n{\\n ${method_header}\\n static PythonArgParser parser({\\n ${signatures}\\n }, /*traceable=*/${traceable});\\n\\n ParsedArgs<${max_args}> parsed_args;\\n auto _r = parser.parse(${self_}, args, kwargs, parsed_args);\\n ${check_has_torch_function}\\n switch (_r.idx) {\\n ${dispatch}\\n }\\n ${method_footer}\\n}\\n\\n\\\"\\\"\\\"\\n)\\n\\n# handler for a single parsed signature - may be a single overload or\\n# a pair of overloads that whose signatures only differ in output params\\n# (plugged into PY_VARIABLE_METHOD_VARARGS as an item in ${dispatch})\\nPY_VARIABLE_CASE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\ncase ${overload_index}: {\\n ${body}\\n}\\n\\\"\\\"\\\"\\n)\\n\\n# python binding for single-overload function/method\\nPY_VARIABLE_METHOD_VARARGS_SINGLETON = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n// ${name}\\nstatic PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs)\\n{\\n ${method_header}\\n static PythonArgParser parser({\\n ${signatures}\\n }, /*traceable=*/${traceable});\\n\\n ParsedArgs<${max_args}> parsed_args;\\n auto _r = parser.parse(${self_}, args, kwargs, parsed_args);\\n ${check_has_torch_function}\\n ${dispatch}\\n ${method_footer}\\n}\\n\\n\\\"\\\"\\\"\\n)\\n\\n# python binding for a method with no args, shortcuts parsing\\nPY_VARIABLE_METHOD_NOARGS = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n// ${name}\\nstatic PyObject * ${pycname}(PyObject* self_, PyObject* args)\\n{\\n ${method_header}\\n ${check_has_torch_function}\\n ${dispatch}\\n ${method_footer}\\n}\\n\\n\\\"\\\"\\\"\\n)\\n\\n\\ndef method_impl(\\n name: BaseOperatorName,\\n module: str | None,\\n overloads: Sequence[PythonSignatureNativeFunctionPair],\\n *,\\n method: bool,\\n symint: bool = True,\\n) -> str:\\n \\\"\\\"\\\"\\n Generate a python binding for all overloads of an op.\\n \\\"\\\"\\\"\\n pycname = get_pycname(name)\\n noarg = is_noarg(overloads)\\n structseq_inits, structseq_typenames = emit_structseq_call(overloads)\\n\\n method_header = [\\\"HANDLE_TH_ERRORS\\\"]\\n method_header += structseq_inits\\n method_header += (\\n [\\\"const Tensor& self = THPVariable_Unpack(self_);\\\"] if method else []\\n )\\n\\n method_footer = ([] if noarg else [\\\"Py_RETURN_NONE;\\\"]) + [\\\"END_HANDLE_TH_ERRORS\\\"]\\n\\n traceable = \\\"true\\\" if all(should_trace(o.function) for o in overloads) else \\\"false\\\"\\n\\n grouped_overloads: Sequence[PythonSignatureGroup] = group_overloads(\\n overloads, symint=symint\\n )\\n is_singleton = len(grouped_overloads) == 1\\n signatures: list[str] = []\\n dispatch: list[str] = []\\n for overload_index, overload in enumerate(grouped_overloads):\\n signature = overload.signature.signature_str(symint=symint)\\n signatures.append(f\\\"{cpp_string(str(signature))},\\\")\\n dispatch_body = emit_dispatch_case(overload, structseq_typenames, symint=symint)\\n dispatch.append(\\n PY_VARIABLE_CASE.substitute(\\n overload_index=overload_index, body=dispatch_body\\n )\\n if not is_singleton\\n else dispatch_body\\n )\\n\\n if noarg:\\n template = PY_VARIABLE_METHOD_NOARGS\\n elif is_singleton:\\n template = PY_VARIABLE_METHOD_VARARGS_SINGLETON\\n else:\\n template = PY_VARIABLE_METHOD_VARARGS\\n\\n return template.substitute(\\n name=name,\\n pycname=pycname,\\n method_header=method_header,\\n max_args=max(o.signature.arguments_count() for o in overloads),\\n signatures=signatures,\\n traceable=traceable,\\n check_has_torch_function=gen_has_torch_function_check(\\n name=name,\\n module=module,\\n noarg=noarg,\\n method=method,\\n ),\\n dispatch=dispatch,\\n method_footer=method_footer,\\n self_=\\\"self_\\\" if method else \\\"nullptr\\\",\\n )\\n\\n\\ndef gen_has_torch_function_check(\\n name: BaseOperatorName, module: str | None, *, noarg: bool, method: bool\\n) -> str:\\n if noarg:\\n if method:\\n return f\\\"\\\"\\\"\\\\\\nif(check_has_torch_function(self_)) {{\\n return handle_torch_function(self_, \\\"{name}\\\");\\n}}\\n\\\"\\\"\\\"\\n else:\\n return \\\"\\\"\\n\\n self_ = \\\"self_\\\" if method else \\\"nullptr\\\"\\n namespace = (\\n {\\n \\\"torch\\\": \\\"THPVariableFunctionsModule\\\",\\n \\\"torch.nn\\\": \\\"THPNNVariableFunctionsModule\\\",\\n \\\"torch.fft\\\": \\\"THPFFTVariableFunctionsModule\\\",\\n \\\"torch.linalg\\\": \\\"THPLinalgVariableFunctionsModule\\\",\\n \\\"torch.nested\\\": \\\"THPNestedVariableFunctionsModule\\\",\\n \\\"torch.sparse\\\": \\\"THPSparseVariableFunctionsModule\\\",\\n \\\"torch.special\\\": \\\"THPSpecialVariableFunctionsModule\\\",\\n }[module]\\n if module\\n else \\\"THPVariableClass\\\"\\n )\\n\\n return f\\\"\\\"\\\"\\\\\\nif(_r.has_torch_function()) {{\\n return handle_torch_function(_r, {self_}, args, kwargs, {namespace}, \\\"{module or \\\"torch.Tensor\\\"}\\\");\\n}}\\n\\\"\\\"\\\"\\n\\n\\n# handler for output/no-output overload pair\\nPY_VARIABLE_OUT = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (_r.isNone(${out_idx})) {\\n ${call_dispatch}\\n} else {\\n ${call_dispatch_out}\\n}\\n\\\"\\\"\\\"\\n)\\n\\n\\ndef emit_dispatch_case(\\n overload: PythonSignatureGroup,\\n structseq_typenames: dict[str, str],\\n *,\\n symint: bool = True,\\n) -> str:\\n \\\"\\\"\\\"\\n Emit dispatch code for a single parsed signature. This corresponds to either\\n a single native function, or a pair that differ only in output params. In the\\n latter case, a single python signature is used for both and dispatching\\n switches on the presence/absence of passed output args.\\n \\\"\\\"\\\"\\n if overload.outplace is not None:\\n # dispatch output and no-output variants, branch on _r.isNone()\\n return PY_VARIABLE_OUT.substitute(\\n out_idx=overload.signature.output_idx(),\\n call_dispatch=emit_single_dispatch(\\n overload.signature, overload.base, structseq_typenames, symint=symint\\n ),\\n call_dispatch_out=emit_single_dispatch(\\n overload.signature,\\n overload.outplace,\\n structseq_typenames,\\n symint=symint,\\n ),\\n )\\n else:\\n # no-output version only\\n return emit_single_dispatch(\\n overload.signature, overload.base, structseq_typenames, symint=symint\\n )\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Forward Declarations Codegen\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef forward_decls(\\n name: BaseOperatorName,\\n overloads: Sequence[PythonSignatureNativeFunctionPair],\\n *,\\n method: bool,\\n) -> tuple[str, ...]:\\n if method:\\n return ()\\n\\n pycname = get_pycname(name)\\n if is_noarg(overloads):\\n return (\\n f\\\"\\\"\\\"\\\\\\nstatic PyObject * {pycname}(PyObject* self_, PyObject* args);\\n\\\"\\\"\\\",\\n )\\n else:\\n return (\\n f\\\"\\\"\\\"\\\\\\nstatic PyObject * {pycname}(PyObject* self_, PyObject* args, PyObject* kwargs);\\n\\\"\\\"\\\",\\n )\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Method Def (Binding Table Entry) Codegen\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef method_def(\\n name: BaseOperatorName,\\n module: str | None,\\n overloads: Sequence[PythonSignatureNativeFunctionPair],\\n *,\\n method: bool,\\n) -> str:\\n \\\"\\\"\\\"\\n Generate method def entry.\\n \\\"\\\"\\\"\\n pycname = get_pycname(name)\\n\\n if name.dunder_method:\\n # PyMethodDef entry for binary op, throws not implemented error\\n pycname = f\\\"TypeError_to_NotImplemented_<{pycname}>\\\"\\n\\n if is_noarg(overloads):\\n flags = \\\"METH_NOARGS\\\" if method else \\\"METH_VARARGS | METH_KEYWORDS\\\"\\n else:\\n pycname = f\\\"castPyCFunctionWithKeywords({pycname})\\\"\\n flags = \\\"METH_VARARGS | METH_KEYWORDS\\\"\\n\\n if module == \\\"torch\\\":\\n flags += \\\" | METH_STATIC\\\"\\n\\n return f'{{\\\"{name}\\\", {pycname}, {flags}, NULL}},'\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Overload Sorting and Grouping\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef group_overloads(\\n overloads: Sequence[PythonSignatureNativeFunctionPair], *, symint: bool = True\\n) -> Sequence[PythonSignatureGroup]:\\n bases: dict[str, PythonSignatureNativeFunctionPair] = {}\\n outplaces: dict[str, PythonSignatureNativeFunctionPair] = {}\\n\\n # first group by signature ignoring out arguments\\n for overload in overloads:\\n sig = overload.signature.signature_str(skip_outputs=True, symint=symint)\\n if overload.function.func.is_out_fn():\\n if sig in outplaces:\\n raise RuntimeError(\\n f\\\"Found duplicated function definition:\\\\n- {overload.function.func}.\\\\n\\\"\\n f\\\"Existing definition:\\\\n- {outplaces[sig].function.func}.\\\"\\n )\\n outplaces[sig] = overload\\n else:\\n if sig in bases:\\n raise RuntimeError(\\n f\\\"Found duplicated function definition:\\\\n- {overload.function.func}.\\\\n\\\"\\n f\\\"Existing definition:\\\\n- {bases[sig].function.func}.\\\"\\n )\\n bases[sig] = overload\\n\\n for sig, out in outplaces.items():\\n if sig not in bases:\\n candidates: list[str] = []\\n for overload in overloads:\\n if (\\n str(overload.function.func.name.name)\\n == str(out.function.func.name.name)\\n and not overload.function.func.is_out_fn()\\n and not overload.signature.deprecated\\n ):\\n candidates.append(\\n overload.signature.signature_str(\\n skip_outputs=True, symint=symint\\n )\\n )\\n out_sig = out.signature.signature_str(symint=symint)\\n raise RuntimeError(\\n f\\\"While identifying overloads, we found an out schema {out_sig} without a corresponding non-out variant. \\\"\\n f\\\"We expected the non-out variant to have schema: \\\\n- {sig}\\\\nPlease check that you spelled the schema \\\"\\n \\\"correctly in native_functions.yaml. We discovered the following candidate(s): \\\\n\\\"\\n + \\\"\\\\n\\\".join(f\\\"- {candidate}\\\" for candidate in candidates)\\n )\\n\\n grouped = [\\n PythonSignatureGroup.from_pairs(\\n functional=base,\\n out=outplaces.get(sig),\\n )\\n for sig, base in bases.items()\\n ]\\n return sort_overloads(grouped, symint=symint)\\n\\n\\n# This function declares a partial order on declarations, and sorts them according\\n# to its linear extension. This is necessary, because there's some ambiguity in the\\n# choice of overload, and we want a different order.\\n#\\n# See Note[Order of overloads matters]\\n#\\n# A few examples of ambiguous python signature pairs.\\n#\\n# All parameters have the same type, except one taking Tensor the other taking\\n# Scalar. A numeric PyObject can be casted into Tensor, and a zero-dim Tensor\\n# object can be accepted as Scalar type parameter (see python_arg_parser.cpp).\\n# Therefore, same input arguments might be accepted by either python signature.\\n# We want to always parse the one taking Tensor first.\\n#\\n# bitwise_and(Tensor input, Tensor other, *, Tensor out=None)\\n# bitwise_and(Tensor input, Scalar other, *, Tensor out=None)\\n#\\n# If they have different number of parameters then they are not ambiguous - but\\n# the difference on output param can be ignored as it's optional.\\n#\\n# multiply(Tensor input, Tensor other, *, Tensor out=None)\\n# multiply(Tensor input, Scalar other)\\n#\\n# Both positional args and keyword-only args are considered together.\\n#\\n# subtract(Tensor other, *, Scalar alpha=1)\\n# subtract(Scalar other, Scalar alpha=1)\\n#\\n# A few ambiguous cases which it does NOT handle yet.\\n#\\n# If there is any difference in other parameters besides the Tensor/Scalar\\n# difference, then they are not considered ambiguous by this method anymore.\\n# However, the difference could be too trivial to disambiguate.\\n#\\n# foo(Tensor input, Scalar other, Scalar bar)\\n# foo(Tensor input, Tensor other, double bar)\\n#\\n# If they are taking different number of parameters then they are not considered\\n# ambiguous anymore, even if the difference is only on optional kwargs.\\n#\\n# foo(Scalar other, Scalar alpha=1)\\n# foo(Tensor other, *, Scalar alpha=1, Scalar beta=1)\\n#\\n\\n\\ndef sort_overloads(\\n grouped_overloads: Sequence[PythonSignatureGroup], *, symint: bool = True\\n) -> Sequence[PythonSignatureGroup]:\\n # NB: Smaller here means lower priority\\n\\n def is_arg_smaller(t1: Type, t2: Type) -> bool:\\n return (\\n str(t1) == \\\"Scalar\\\"\\n and str(t2) == \\\"Tensor\\\"\\n or str(t1) == \\\"Scalar?\\\"\\n and str(t2) == \\\"Tensor?\\\"\\n or \\\"Dimname\\\" in str(t1)\\n and \\\"Dimname\\\" not in str(t2)\\n or\\n # In the discussion https://github.com/pytorch/pytorch/issues/54555 it has been\\n # discussed why it is important to prioritize int/int? over int[]\\n str(t1) == \\\"int[]\\\"\\n and (str(t2) == \\\"int\\\" or str(t2) == \\\"int?\\\")\\n or\\n # TensorList currently throws an error during argument parsing, that's why it needs to be\\n # last in signature ordering. See discussion: https://github.com/pytorch/pytorch/issues/58087\\n str(t1) == \\\"Tensor[]\\\"\\n and str(t2).find(\\\"[]\\\") != -1\\n or\\n # Prioritize IntArrayRef overload over SymIntArrayRef\\n str(t1) == \\\"SymInt[]\\\"\\n and str(t2) == \\\"int[]\\\"\\n or\\n # Make sure both in, SymInt are sorted consistently w.r.t. Tensor since Tensor can be implicitly\\n # converted to either int or SymInt. Prioritize the Tensor overload since it otherwise gets shadowed.\\n (str(t1) == \\\"SymInt\\\" or str(t1) == \\\"int\\\")\\n and str(t2) == \\\"Tensor\\\"\\n )\\n\\n def is_smaller(s1: PythonSignature, s2: PythonSignature) -> bool:\\n \\\"\\\"\\\"Returns True if s1 < s2 in the partial order.\\\"\\\"\\\"\\n args1, args2 = s1.arguments(skip_outputs=True), s2.arguments(skip_outputs=True)\\n if len(args1) != len(args2):\\n return False\\n # TODO: should use some canonical form instead of 'str(arg.type)' - see comments\\n # above. The old codegen used the deprecated 'dynamic_type(arg.type)', which\\n # ignores the optional annotation, i.e. 'Scalar' and 'Scalar?'.\\n equal = all(arg1.type == arg2.type for arg1, arg2 in zip(args1, args2))\\n smaller_or_equal = all(\\n str(arg1.type) == str(arg2.type) or is_arg_smaller(arg1.type, arg2.type)\\n for arg1, arg2 in zip(args1, args2)\\n )\\n return smaller_or_equal and not equal\\n\\n # First sort by signature\\n grouped_overloads = sorted(\\n grouped_overloads, key=lambda x: x.signature.signature_str(symint=symint)\\n )\\n\\n # Construct the relation graph\\n larger_than: dict[int, set[int]] = defaultdict(set)\\n for i1, overload1 in enumerate(grouped_overloads):\\n for i2, overload2 in enumerate(grouped_overloads):\\n if is_smaller(overload1.signature, overload2.signature):\\n larger_than[i1].add(i2)\\n\\n if not larger_than:\\n return list(grouped_overloads)\\n\\n # Use a topological sort to sort overloads according to the partial order.\\n N = len(grouped_overloads)\\n sorted_ids: list[int] = list(filter(lambda x: x not in larger_than, range(N)))\\n\\n for idx in range(N):\\n # The size of sorted_ids will grow to N eventually.\\n i = sorted_ids[idx]\\n for j in sorted(larger_than.keys()):\\n larger = larger_than[j]\\n larger.discard(i)\\n if not larger:\\n del larger_than[j]\\n sorted_ids.append(j)\\n\\n return [grouped_overloads[x] for x in sorted_ids]\\n\\n\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n#\\n# Codegen API Integration\\n#\\n# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #\\n\\n\\ndef emit_single_dispatch(\\n ps: PythonSignature,\\n f: NativeFunction,\\n structseq_typenames: dict[str, str],\\n *,\\n symint: bool = True,\\n) -> str:\\n \\\"\\\"\\\"\\n Emit dispatch code for a single native function.\\n \\\"\\\"\\\"\\n\\n @with_native_function\\n def go(f: NativeFunction) -> str:\\n # header comments\\n if isinstance(ps, PythonSignatureDeprecated):\\n schema_comment = f\\\"// [deprecated] aten::{ps.deprecated_schema}\\\"\\n else:\\n schema_comment = f\\\"// aten::{f.func}\\\"\\n\\n deprecated = \\\"[deprecated] \\\" if ps.deprecated else \\\"\\\"\\n\\n # dispatch lambda signature\\n name = cpp.name(f.func)\\n lambda_formals = \\\", \\\".join(\\n f\\\"{a.type_str} {a.name}\\\" for a in dispatch_lambda_args(ps, f, symint=symint)\\n )\\n lambda_return = dispatch_lambda_return_str(f)\\n\\n # dispatch lambda body\\n dispatch_callee = cpp_dispatch_target(f)\\n dispatch_args = \\\", \\\".join(cpp_dispatch_exprs(f, python_signature=ps))\\n\\n # from arg parser outputs to dispatch lambda arguments\\n parser_outputs = arg_parser_output_exprs(ps, f, symint=symint)\\n lambda_arg_exprs = dispatch_lambda_exprs(ps, f, symint=symint)\\n inits = \\\"\\\\n\\\".join(lambda_arg_exprs.inits)\\n lambda_args = \\\", \\\".join(lambda_arg_exprs.exprs)\\n\\n # scatter fields\\n # TODO: Checking `ps.method and ('requires_grad' in parser_outputs)` is a hacky\\n # solution for enabling the 'requires_grad' argument for tensor methods\\n # new_full, new_empty, and new_zeros. A much better but more difficult to\\n # implement solution involves refactoring according to Ed's description here:\\n # https://github.com/pytorch/pytorch/issues/36455#issuecomment-614767589\\n need_set_requires_grad = ps.tensor_options_args and (\\n not has_tensor_options(f)\\n or (ps.method and (\\\"requires_grad\\\" in parser_outputs))\\n )\\n set_requires_grad = (\\n f'.set_requires_grad({parser_outputs[\\\"requires_grad\\\"].expr})'\\n if need_set_requires_grad\\n else \\\"\\\"\\n )\\n\\n if lambda_return == \\\"void\\\":\\n # Make in-place foreach return `self` at python-binding level.\\n # ref: https://github.com/pytorch/pytorch/pull/118622#pullrequestreview-1904804954\\n self_arg = f.func.arguments.self_arg\\n return_stmt: str\\n if (\\n str(f.func.name).startswith(\\\"_foreach_\\\")\\n and f.func.kind() == SchemaKind.inplace\\n ):\\n # note(crcrpar): `_foreach_pow.ScalarAndTensor` does NOT have its in-place\\n # variant and it unlikely to have it in the future. Thus it's safe to have the following assert.\\n assert self_arg is not None and is_tensor_list_type(\\n self_arg.argument.type\\n )\\n return_stmt = \\\"\\\"\\\"PyObject* self_tensorlist = _r.args[0];\\nPy_INCREF(self_tensorlist);\\nreturn self_tensorlist;\\n\\\"\\\"\\\"\\n else:\\n return_stmt = \\\"Py_RETURN_NONE;\\\"\\n return f\\\"\\\"\\\"\\\\\\n{schema_comment}\\n{inits}\\nauto dispatch_{name} = []({lambda_formals}) -> {lambda_return} {{\\n pybind11::gil_scoped_release no_gil;\\n {dispatch_callee}({dispatch_args});\\n}};\\ndispatch_{name}({lambda_args}){set_requires_grad};\\n{return_stmt}\\n\\\"\\\"\\\"\\n else:\\n typename = structseq_typenames.get(gen_structseq_typename_key(f))\\n structseq_typeref = f\\\"{typename}, \\\" if typename is not None else \\\"\\\"\\n return f\\\"\\\"\\\"\\\\\\n{schema_comment}\\n{inits}\\nauto dispatch_{name} = []({lambda_formals}) -> {lambda_return} {{\\n pybind11::gil_scoped_release no_gil;\\n return {dispatch_callee}({dispatch_args});\\n}};\\nreturn wrap({structseq_typeref}dispatch_{name}({lambda_args}){set_requires_grad});\\n\\\"\\\"\\\"\\n\\n return go(f)\\n\\n\\nimport functools\\nfrom typing import Callable\\n\\nfrom torchgen.api.autograd import NativeFunctionWithDifferentiabilityInfo as NFWDI\\nfrom torchgen.context import native_function_manager\\nfrom torchgen.utils import T\\n\\n\\n# Like tools.api.context.with_native_function, but for\\n# NativeFunctionWithDifferentiabilityInfo.\\ndef with_native_function_with_differentiability_info(\\n func: Callable[[NFWDI], T]\\n) -> Callable[[NFWDI], T]:\\n @functools.wraps(func)\\n def wrapper(f: NFWDI) -> T:\\n with native_function_manager(f.func):\\n return func(f)\\n\\n return wrapper\\n\\n\\n# Like the above but with an additional dispatch key string argument\\ndef with_native_function_with_differentiability_info_and_key(\\n func: Callable[[NFWDI, str], T]\\n) -> Callable[[NFWDI, str], T]:\\n @functools.wraps(func)\\n def wrapper(f: NFWDI, key: str) -> T:\\n with native_function_manager(f.func):\\n return func(f, key)\\n\\n return wrapper\\n\\n\\n\\\"\\\"\\\"\\nTo run this file by hand from the root of the PyTorch\\nrepository, run:\\n\\npython -m tools.autograd.gen_autograd \\\\\\n aten/src/ATen/native/native_functions.yaml \\\\\\n aten/src/ATen/native/tags.yaml \\\\\\n $OUTPUT_DIR \\\\\\n tools/autograd\\n\\nWhere $OUTPUT_DIR is where you would like the files to be\\ngenerated. In the full build system, OUTPUT_DIR is\\ntorch/csrc/autograd/generated/\\n\\\"\\\"\\\"\\n\\n# gen_autograd.py generates C++ autograd functions and Python bindings.\\n#\\n# It delegates to the following scripts:\\n#\\n# gen_autograd_functions.py: generates subclasses of torch::autograd::Node\\n# gen_variable_type.py: generates VariableType.h which contains all tensor methods\\n# gen_python_functions.py: generates Python bindings to THPVariable\\n#\\n\\nfrom __future__ import annotations\\n\\nimport argparse\\nimport os\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.autograd import (\\n match_differentiability_info,\\n NativeFunctionWithDifferentiabilityInfo,\\n)\\nfrom torchgen.gen import parse_native_yaml\\nfrom torchgen.selective_build.selector import SelectiveBuilder\\n\\nfrom . import gen_python_functions\\nfrom .gen_autograd_functions import (\\n gen_autograd_functions_lib,\\n gen_autograd_functions_python,\\n)\\nfrom .gen_inplace_or_view_type import gen_inplace_or_view_type\\nfrom .gen_trace_type import gen_trace_type\\nfrom .gen_variable_factories import gen_variable_factories\\nfrom .gen_variable_type import gen_variable_type\\nfrom .gen_view_funcs import gen_view_funcs\\nfrom .load_derivatives import load_derivatives\\n\\n\\ndef gen_autograd(\\n native_functions_path: str,\\n tags_path: str,\\n out: str,\\n autograd_dir: str,\\n operator_selector: SelectiveBuilder,\\n disable_autograd: bool = False,\\n) -> None:\\n # Parse and load derivatives.yaml\\n differentiability_infos, used_dispatch_keys = load_derivatives(\\n os.path.join(autograd_dir, \\\"derivatives.yaml\\\"), native_functions_path, tags_path\\n )\\n\\n template_path = os.path.join(autograd_dir, \\\"templates\\\")\\n\\n native_funcs = parse_native_yaml(native_functions_path, tags_path).native_functions\\n fns = sorted(\\n filter(\\n operator_selector.is_native_function_selected_for_training, native_funcs\\n ),\\n key=lambda f: cpp.name(f.func),\\n )\\n fns_with_diff_infos: list[\\n NativeFunctionWithDifferentiabilityInfo\\n ] = match_differentiability_info(fns, differentiability_infos)\\n\\n # Generate VariableType.h/cpp\\n if not disable_autograd:\\n gen_variable_type(\\n out,\\n native_functions_path,\\n tags_path,\\n fns_with_diff_infos,\\n template_path,\\n used_dispatch_keys,\\n )\\n\\n gen_inplace_or_view_type(\\n out, native_functions_path, tags_path, fns_with_diff_infos, template_path\\n )\\n\\n # operator filter not applied as tracing sources are excluded in selective build\\n gen_trace_type(out, native_funcs, template_path)\\n # Generate Functions.h/cpp\\n gen_autograd_functions_lib(out, differentiability_infos, template_path)\\n\\n # Generate variable_factories.h\\n gen_variable_factories(out, native_functions_path, tags_path, template_path)\\n\\n # Generate ViewFuncs.h/cpp\\n gen_view_funcs(out, fns_with_diff_infos, template_path)\\n\\n\\ndef gen_autograd_python(\\n native_functions_path: str,\\n tags_path: str,\\n out: str,\\n autograd_dir: str,\\n) -> None:\\n differentiability_infos, _ = load_derivatives(\\n os.path.join(autograd_dir, \\\"derivatives.yaml\\\"), native_functions_path, tags_path\\n )\\n\\n template_path = os.path.join(autograd_dir, \\\"templates\\\")\\n\\n # Generate Functions.h/cpp\\n gen_autograd_functions_python(out, differentiability_infos, template_path)\\n\\n # Generate Python bindings\\n deprecated_path = os.path.join(autograd_dir, \\\"deprecated.yaml\\\")\\n gen_python_functions.gen(\\n out, native_functions_path, tags_path, deprecated_path, template_path\\n )\\n\\n\\ndef main() -> None:\\n parser = argparse.ArgumentParser(description=\\\"Generate autograd C++ files script\\\")\\n parser.add_argument(\\n \\\"native_functions\\\", metavar=\\\"NATIVE\\\", help=\\\"path to native_functions.yaml\\\"\\n )\\n parser.add_argument(\\\"tags\\\", metavar=\\\"NATIVE\\\", help=\\\"path to tags.yaml\\\")\\n parser.add_argument(\\\"out\\\", metavar=\\\"OUT\\\", help=\\\"path to output directory\\\")\\n parser.add_argument(\\n \\\"autograd\\\", metavar=\\\"AUTOGRAD\\\", help=\\\"path to autograd directory\\\"\\n )\\n args = parser.parse_args()\\n gen_autograd(\\n args.native_functions,\\n args.tags,\\n args.out,\\n args.autograd,\\n SelectiveBuilder.get_nop_selector(),\\n )\\n\\n\\nif __name__ == \\\"__main__\\\":\\n main()\\n\\n\\n\\n\\n# Generates ADInplaceOrViewType.h/cpp\\n#\\n# NOTE: If any changes are being made to the ADInplaceOrView codegen please also check\\n# if updates are needed in torch/csrc/autograd/autograd_not_implemented_fallback.cpp\\n# The fallback is expected to mimick this codegen, so we should keep the two in sync.\\n\\nfrom __future__ import annotations\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.autograd import (\\n dispatch_strategy,\\n gen_differentiable_outputs,\\n NativeFunctionWithDifferentiabilityInfo,\\n)\\nfrom torchgen.api.types import (\\n BaseCType,\\n Binding,\\n boolT,\\n ConstRefCType,\\n CType,\\n DispatcherSignature,\\n intArrayRefT,\\n longT,\\n OptionalCType,\\n symIntArrayRefT,\\n SymIntT,\\n tensorT,\\n)\\nfrom torchgen.code_template import CodeTemplate\\nfrom torchgen.context import with_native_function\\nfrom torchgen.model import (\\n NativeFunction,\\n SchemaKind,\\n SelfArgument,\\n TensorOptionsArguments,\\n Type,\\n)\\nfrom torchgen.utils import FileManager\\n\\nfrom .context import with_native_function_with_differentiability_info\\nfrom .gen_trace_type import (\\n get_return_value,\\n MANUAL_AUTOGRAD,\\n tie_return_values,\\n type_wrapper_name,\\n)\\n\\n\\n# See NOTE [ Autograd View Variables ] in variable.h for details.\\n# If you update list VIEW_FUNCTIONS or RETURNS_VIEWS_OF_INPUT,\\n# you **MUST** also update the public list of view ops accordingly in\\n# docs/source/tensor_view.rst. Note not all ATen functions are exposed to public,\\n# e.g alias & sparse_coo_tensor_with_dims_and_tensors.\\n#\\n# A map: function name => name of the argument that all outputs are view of\\n\\nVIEW_FUNCTIONS_WITH_METADATA_CHANGE = [\\n \\\"view_as_complex\\\",\\n \\\"view_as_real\\\",\\n \\\"_conj\\\",\\n \\\"_neg_view\\\",\\n \\\"_nested_get_values\\\",\\n \\\"_nested_view_from_buffer\\\",\\n \\\"_nested_view_from_jagged\\\",\\n]\\n\\nVIEW_FUNCTIONS = {\\n \\\"numpy_T\\\": \\\"self\\\",\\n \\\"alias\\\": \\\"self\\\",\\n \\\"as_strided\\\": \\\"self\\\",\\n \\\"diagonal\\\": \\\"self\\\",\\n \\\"expand\\\": \\\"self\\\",\\n \\\"permute\\\": \\\"self\\\",\\n \\\"select\\\": \\\"self\\\",\\n \\\"slice\\\": \\\"self\\\",\\n \\\"slice_inverse\\\": \\\"self\\\",\\n \\\"split\\\": \\\"self\\\",\\n \\\"split_with_sizes\\\": \\\"self\\\",\\n \\\"squeeze\\\": \\\"self\\\",\\n \\\"t\\\": \\\"self\\\",\\n \\\"transpose\\\": \\\"self\\\",\\n \\\"unfold\\\": \\\"self\\\",\\n \\\"unsqueeze\\\": \\\"self\\\",\\n \\\"flatten\\\": \\\"self\\\",\\n \\\"view\\\": \\\"self\\\",\\n \\\"unbind\\\": \\\"self\\\",\\n \\\"_indices\\\": \\\"self\\\",\\n \\\"_values\\\": \\\"self\\\",\\n \\\"indices\\\": \\\"self\\\",\\n \\\"values\\\": \\\"self\\\",\\n \\\"crow_indices\\\": \\\"self\\\",\\n \\\"col_indices\\\": \\\"self\\\",\\n \\\"ccol_indices\\\": \\\"self\\\",\\n \\\"row_indices\\\": \\\"self\\\",\\n # sparse_coo ctor output should really be views of both indices and values,\\n # but we only supports making as view of a single variable, and indices is\\n # discrete anyways.\\n # FIXME: clone indices on construction.\\n \\\"sparse_coo_tensor_with_dims_and_tensors\\\": \\\"values\\\",\\n \\\"_reshape_alias\\\": \\\"self\\\",\\n \\\"_test_autograd_multiple_dispatch_view\\\": \\\"self\\\",\\n}\\n\\nfor key in VIEW_FUNCTIONS_WITH_METADATA_CHANGE:\\n VIEW_FUNCTIONS[key] = \\\"self\\\"\\n\\n# note: some VIEW_FUNCTIONS are just compositions of the view functions above\\n# this list contains both the root view functions and any that are purely composed\\n# of viewing functions, and is used by the JIT to determine when an operator\\n# may return a view of its inputs; however they may sometimes return a copy.\\n# (e.g. `contiguous`)\\nRETURNS_VIEWS_OF_INPUT = set(VIEW_FUNCTIONS.keys()).union(\\n {\\n \\\"chunk\\\",\\n \\\"detach\\\",\\n \\\"contiguous\\\",\\n \\\"reshape\\\",\\n \\\"reshape_as\\\",\\n \\\"expand_as\\\",\\n \\\"view_as\\\",\\n \\\"real\\\",\\n \\\"imag\\\",\\n \\\"narrow\\\",\\n \\\"movedim\\\",\\n \\\"tensor_split\\\",\\n \\\"swapdims\\\",\\n \\\"swapaxes\\\",\\n \\\"mT\\\",\\n \\\"mH\\\",\\n \\\"adjoint\\\",\\n \\\"matrix_H\\\",\\n }\\n)\\n\\n# These are the functions we consider views for the purposes of validating\\n# StorageImpl and TensorImpl in gen_variable_type.\\n# `_unsafe_view` is not included in VIEW_FUNCTIONS above because it is not a\\n# view for the purposes of ADInplaceOrView kernel, we do not want to call as_view\\n# See NOTE [Unsafe View] for more info.\\nALL_VIEW_FUNCTIONS = {\\n **VIEW_FUNCTIONS,\\n \\\"_unsafe_view\\\": \\\"self\\\",\\n}\\n\\nARRAYREF_TO_VEC = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nauto ${vec} = ${arg}.vec();\\n\\\"\\\"\\\"\\n)\\n\\nOPTIONAL_TO_VAL = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nauto ${val} = ${arg}.value_or(${default});\\n\\\"\\\"\\\"\\n)\\n\\nCALL_DISPATCH = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nat::_ops::${unambiguous_name}::call(${unpacked_args})\\\"\\\"\\\"\\n)\\n\\nREVERSE_VIEW_DISPATCH = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${reverse_name}(${unpacked_args})\\\"\\\"\\\"\\n)\\n\\nMULTI_OUTPUT_VIEW_ITERATION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfor (auto ${view_idx} : c10::irange(${var}.size())) {\\n ${body}\\n}\\n\\\"\\\"\\\"\\n)\\n\\nSETUP_REPLAY_VIEW_IF_NOT_SUPPORT_AS_STRIDED_OR_VIEW_WITH_METADATA_CHANGE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstd::unique_ptr func(nullptr);\\nstd::function rev_func=nullptr;\\nif (${is_view_with_metadata_change} ||\\n !self.unsafeGetTensorImpl()->support_as_strided() ||\\n self.unsafeGetTensorImpl()->is_python_dispatch() ||\\n c10::AutogradState::get_tls_state().get_view_replay_enabled()) {\\n ${replay_view_func}\\n ${reverse_replay_view_func}\\n}\\n\\\"\\\"\\\"\\n)\\n\\nREPLAY_VIEW_FUNC = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfunc = std::make_unique<${view_func_name}>(${view_func_args});\\n\\\"\\\"\\\"\\n)\\n\\nREVERSE_REPLAY_VIEW_LAMBDA_FUNC = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nrev_func = [=](const at::Tensor& ${input_view}) {\\n return ${reverse_replay_view_call};\\n};\\n\\\"\\\"\\\"\\n)\\n\\nMETHOD_DEFINITION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${return_type} ${type_wrapper_name}(${formals}) {\\n ${type_definition_body}\\n}\\n\\\"\\\"\\\"\\n)\\n\\nWRAPPER_REGISTRATION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nm.impl(\\\"${unqual_operator_name_with_overload}\\\",\\n TORCH_FN(${class_type}::${type_wrapper_name})\\n);\\n\\\"\\\"\\\"\\n)\\n\\nAUTOGRAD_NOT_IMPLEMENTED_REGISTRATION = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nm.impl(\\\"${unqual_operator_name_with_overload}\\\", torch::autograd::autogradNotImplementedFallback());\\n\\\"\\\"\\\"\\n)\\n\\nINPLACE_REDISPATCH = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n{\\n at::AutoDispatchBelowADInplaceOrView guard;\\n at::_ops::${unambiguous_name}::redispatch(${unpacked_args});\\n}\\n\\\"\\\"\\\"\\n)\\n\\nASSIGN_RETURN_VALUE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${return_values} = ${rhs_value};\\n\\\"\\\"\\\"\\n)\\n\\nVIEW_REDISPATCH = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${assign_return_values} ([&]() {\\n at::AutoDispatchBelowADInplaceOrView guard;\\n return at::_ops::${unambiguous_name}::redispatch(${unpacked_args});\\n})();\\n\\\"\\\"\\\"\\n)\\n\\nTMP_VAR = \\\"_tmp\\\"\\n\\n\\n# FIXME: Ideally these functions should be methods on Type class, but we have a\\n# comment in codegen/model.py there saying these concepts are not well defined.\\n# Thus we put a version that commonly used by autograd codegen here.\\ndef is_tensor_type(t: Type) -> bool:\\n # TODO: Should handle optional here?\\n return t.is_tensor_like() and t.is_list_like() is None\\n\\n\\ndef is_tensor_list_type(t: Type) -> bool:\\n # TODO: Should handle optional here?\\n return t.is_tensor_like() and t.is_list_like() is not None\\n\\n\\nUNPACK_TENSOR = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nauto${ref} ${arg_name}_ = unpack${suffix}(${arg_name}, \\\"${arg_name}\\\", ${arg_pos});\\\"\\\"\\\"\\n)\\n\\n\\ndef unpacked_name(arg_name: str) -> str:\\n return arg_name + \\\"_\\\"\\n\\n\\n# e.g. select.int -> select_copy_int_inverse()\\ndef inverse_view_name(f: NativeFunction) -> str:\\n copy_variant = f\\\"{f.root_name}_copy\\\"\\n overload = f\\\"{f.func.name.overload_name}\\\"\\n if overload != \\\"\\\":\\n overload = \\\"_\\\" + overload\\n return f\\\"{copy_variant}{overload}_inverse\\\"\\n\\n\\ndef extract_bindings(f: NativeFunction) -> list[Binding]:\\n return [\\n r\\n for a in f.func.schema_order_arguments()\\n for r in cpp.argument(\\n a,\\n method=False,\\n symint=True,\\n cpp_no_default_args=set(),\\n faithful=False,\\n has_tensor_options=False,\\n )\\n ]\\n\\n\\n@with_native_function\\ndef unpack_args(f: NativeFunction) -> tuple[list[str], list[Binding]]:\\n body: list[str] = []\\n unpacked_bindings: list[Binding] = []\\n\\n for i, binding in enumerate(extract_bindings(f)):\\n assert not isinstance(binding.argument, SelfArgument)\\n if isinstance(binding.argument, TensorOptionsArguments):\\n raise RuntimeError(\\\"VariableKernel shouldn't take TensorOptions\\\")\\n\\n is_nullable = binding.argument.type.is_nullable()\\n if not binding.argument.type.is_tensor_like() or is_nullable:\\n unpacked_bindings.append(binding)\\n continue\\n\\n is_tensor_list = is_tensor_list_type(binding.argument.type)\\n ref = (not is_nullable) and not is_tensor_list\\n suffix = \\\"_opt\\\" if is_nullable and not is_tensor_list else \\\"\\\"\\n body.append(\\n UNPACK_TENSOR.substitute(\\n arg_name=binding.name,\\n arg_pos=i,\\n suffix=suffix,\\n ref=\\\"&\\\" if ref else \\\"\\\",\\n )\\n )\\n unpacked_bindings.append(\\n Binding(\\n name=unpacked_name(binding.name),\\n nctype=binding.nctype,\\n argument=binding.argument,\\n default=binding.default,\\n )\\n )\\n\\n return body, unpacked_bindings\\n\\n\\ndef get_base_name(f: NativeFunction) -> str:\\n return f.func.name.name.base # TODO: should be str(f.func.name.name)?\\n\\n\\ndef get_view_info(f: NativeFunction) -> str | None:\\n base_name = get_base_name(f)\\n view_info = VIEW_FUNCTIONS.get(base_name, None)\\n if view_info is None and base_name in RETURNS_VIEWS_OF_INPUT:\\n view_info = \\\"self\\\"\\n return view_info\\n\\n\\ndef emit_view_func(\\n f: NativeFunction, bindings: list[Binding], view_idx: str | None = None\\n) -> str:\\n \\\"\\\"\\\"Generate an additional lambda function to recover views in backward when as_strided is not supported.\\n See Note [View + Inplace update for base tensor] and [View + Inplace update for view tensor] for more details.\\n \\\"\\\"\\\"\\n # TODO: Clean this logic up if we get rid of reverse view funcs or reify them.\\n input_base = \\\"input_base\\\"\\n replay_view_func = \\\"\\\"\\n updated_args: list[str] = []\\n known_view_arg_simple_types: list[CType] = [\\n BaseCType(longT),\\n OptionalCType(BaseCType(longT)),\\n BaseCType(SymIntT),\\n OptionalCType(BaseCType(SymIntT)),\\n BaseCType(boolT),\\n BaseCType(intArrayRefT),\\n BaseCType(symIntArrayRefT),\\n ConstRefCType(BaseCType(tensorT)),\\n ConstRefCType(OptionalCType(BaseCType(tensorT))),\\n ]\\n for binding in bindings:\\n arg, arg_type = binding.name, binding.nctype.type\\n if arg == \\\"self\\\":\\n updated_args.append(input_base)\\n continue\\n if arg_type not in known_view_arg_simple_types:\\n known_types_str = \\\", \\\".join([str(t) for t in known_view_arg_simple_types])\\n raise TypeError(\\n f\\\"You are adding an {arg_type} {arg} argument to op {cpp.name(f.func)} in addition to known types: \\\"\\n f\\\"{known_types_str}. Please update the list or materialize it so that it can be closed \\\"\\n \\\"over by value, also add a test in pytorch/xla/test/test_operations.py where this code \\\"\\n \\\"is exercised.\\\"\\n )\\n if arg_type == BaseCType(intArrayRefT) or arg_type == BaseCType(\\n symIntArrayRefT\\n ):\\n # It's not safe to close over IntArrayRef by value, since this is a\\n # reference type, so materialize a vector to close over by value\\n arg_vec = arg + \\\"_vec\\\"\\n replay_view_func += ARRAYREF_TO_VEC.substitute(arg=arg, vec=arg_vec)\\n updated_args.append(arg_vec)\\n elif arg_type == OptionalCType(BaseCType(longT)):\\n # Materialize int64_t? to int64_t\\n arg_value = arg + \\\"_val\\\"\\n replay_view_func += OPTIONAL_TO_VAL.substitute(\\n arg=arg, val=arg_value, default=\\\"0\\\"\\n )\\n updated_args.append(arg_value)\\n elif arg_type == ConstRefCType(BaseCType(tensorT)) or arg_type == ConstRefCType(\\n OptionalCType(BaseCType(tensorT))\\n ):\\n # NB: Closing over a tensor. If a user modifies this tensor, this will be silently\\n # incorrect. The proper thing to do is to store the version counter and copy on write.\\n updated_args.append(arg)\\n else:\\n updated_args.append(arg)\\n\\n from .gen_view_funcs import view_func_name\\n\\n view_func_args = [b.name for b in bindings if b.name != \\\"self\\\"]\\n if view_idx is not None:\\n view_func_args.append(f\\\"{view_idx}\\\")\\n replay_view_func += REPLAY_VIEW_FUNC.substitute(\\n view_func_name=view_func_name(f, include_namespace=True),\\n view_func_args=view_func_args,\\n )\\n\\n input_view = \\\"input_view\\\"\\n reverse_unpacked_args = [\\n \\\"self\\\",\\n f\\\"{input_view}\\\",\\n # inverse_return_mode=\\n \\\"at::functionalization::InverseReturnMode::AlwaysView\\\",\\n *(() if view_idx is None else (f\\\"{view_idx}\\\",)),\\n # skip input_base arg\\n *updated_args[1:],\\n ]\\n\\n from torchgen.api.functionalization import reverse_name\\n\\n reverse_replay_view_call = REVERSE_VIEW_DISPATCH.substitute(\\n reverse_name=reverse_name(f, include_namespace=True),\\n unpacked_args=reverse_unpacked_args,\\n )\\n reverse_replay_view_func = REVERSE_REPLAY_VIEW_LAMBDA_FUNC.substitute(\\n input_view=input_view, reverse_replay_view_call=reverse_replay_view_call\\n )\\n\\n is_view_with_metadata_change = (\\n \\\"true\\\" if cpp.name(f.func) in VIEW_FUNCTIONS_WITH_METADATA_CHANGE else \\\"false\\\"\\n )\\n\\n return SETUP_REPLAY_VIEW_IF_NOT_SUPPORT_AS_STRIDED_OR_VIEW_WITH_METADATA_CHANGE.substitute(\\n is_view_with_metadata_change=is_view_with_metadata_change,\\n replay_view_func=replay_view_func,\\n reverse_replay_view_func=reverse_replay_view_func,\\n )\\n\\n\\ndef emit_view_body(\\n fn: NativeFunctionWithDifferentiabilityInfo, var: str\\n) -> tuple[str, str]:\\n # See NOTE [ Autograd View Variables ] in variable.h for details.\\n f = fn.func\\n base_name = get_base_name(f)\\n view_info = get_view_info(f)\\n call = \\\"\\\"\\n differentiable_outputs = gen_differentiable_outputs(fn)\\n differentiable_output_vars = {r.name for r in differentiable_outputs}\\n if not isinstance(view_info, str):\\n raise TypeError(\\n f\\\"The view info should be a string for {base_name}, but it is: {view_info}\\\"\\n )\\n if len(differentiable_output_vars) == 0:\\n # no output is differentiable (.indices() for SparseTensors for example)\\n rhs_value = (\\n f\\\"as_view({view_info}, {var}, \\\"\\n f\\\"/* is_bw_differentiable */ false, /* is_fw_differentiable */ false)\\\"\\n )\\n elif len(differentiable_output_vars) == 1:\\n # Single differentiable output (Tensor or Tensor[])\\n return_info = differentiable_outputs[0]\\n # We only support simple Tensor or a TensorList for functions that return views\\n if not is_tensor_type(return_info.type) and not is_tensor_list_type(\\n return_info.type\\n ):\\n raise RuntimeError(\\n f\\\"{base_name} that return differentiable views can only return Tensor or Tensor[]\\\"\\n )\\n\\n # See Note [ View + Inplace detection]\\n def get_creation_meta_in_mode(original: str) -> str:\\n creation_meta_with_grad_mode = f\\\"(at::GradMode::is_enabled() ? {original} : CreationMeta::NO_GRAD_MODE)\\\"\\n return f\\\"InferenceMode::is_enabled() ? CreationMeta::INFERENCE_MODE : {creation_meta_with_grad_mode}\\\"\\n\\n # Only allow rebasing of the history if we return a single Tensor\\n # If we are in a no grad block, raise a warning\\n # See NOTE [ View + Inplace detection ] for more details about this logic\\n if is_tensor_list_type(return_info.type):\\n creation_meta = get_creation_meta_in_mode(\\\"CreationMeta::MULTI_OUTPUT_NODE\\\")\\n view_idx = \\\"view_idx\\\"\\n view_func = emit_view_func(\\n f, extract_bindings(f), view_idx=view_idx\\n ).strip()\\n as_view_call = (\\n f\\\"as_view(/* base */ {view_info}, /* output */ {var}[{view_idx}], \\\"\\n \\\"/* is_bw_differentiable */ true, /* is_fw_differentiable */ true, \\\"\\n \\\"/* view_func */ std::move(func), /* rev_view_func */ rev_func, \\\"\\n f\\\"/* creation_meta */ {creation_meta});\\\"\\n )\\n call += MULTI_OUTPUT_VIEW_ITERATION.substitute(\\n var=var, view_idx=view_idx, body=f\\\"{view_func}\\\\n{as_view_call}\\\"\\n )\\n rhs_value = f\\\"std::move({var})\\\"\\n else:\\n call += emit_view_func(f, extract_bindings(f), view_idx=None)\\n creation_meta = get_creation_meta_in_mode(\\\"CreationMeta::DEFAULT\\\")\\n rhs_value = (\\n f\\\"as_view(/* base */ {view_info}, /* output */ {var}, /* is_bw_differentiable */ true, \\\"\\n \\\"/* is_fw_differentiable */ true, \\\"\\n f\\\"/* view_func */ std::move(func), /* rev_view_func */ rev_func, /* creation_meta */ {creation_meta})\\\"\\n )\\n else:\\n # This could be supported but we don't need it at the moment, so keeping things simple.\\n raise RuntimeError(\\n \\\"Function that return multiple differentiable output \\\"\\n \\\"when at least one of them is view is not supported.\\\"\\n )\\n return call, rhs_value\\n\\n\\ndef modifies_arguments(f: NativeFunction) -> bool:\\n return f.func.kind() in [SchemaKind.inplace, SchemaKind.out]\\n\\n\\n@with_native_function_with_differentiability_info\\ndef emit_inplace_or_view_body(fn: NativeFunctionWithDifferentiabilityInfo) -> list[str]:\\n f = fn.func\\n inplace_view_body: list[str] = []\\n\\n dispatcher_sig = DispatcherSignature.from_schema(f.func)\\n dispatcher_exprs = dispatcher_sig.exprs()\\n\\n # code-generated ADInplaceOrView kernels plumb and recompute dispatch keys directly through the kernel for performance.\\n # See Note [Plumbing Keys Through The Dispatcher] for details.\\n dispatch_key_set = \\\"ks & c10::after_ADInplaceOrView_keyset\\\"\\n redispatch_args = \\\", \\\".join([dispatch_key_set] + [a.expr for a in dispatcher_exprs])\\n\\n # Note that this calls the slow, dispatching variants of manual_cpp_binding ops.\\n # We could probably work harder to ensure that the fast variants are called instead, but the perf benefit would be minimal.\\n if modifies_arguments(f): # inplace op\\n inplace_view_body.append(\\n INPLACE_REDISPATCH.substitute(\\n unambiguous_name=f.func.name.unambiguous_name(),\\n unpacked_args=redispatch_args,\\n )\\n )\\n for r in cpp.return_names(f):\\n inplace_view_body.append(f\\\"increment_version({r});\\\")\\n else:\\n assert get_view_info(f) is not None\\n inplace_view_body.append(\\n VIEW_REDISPATCH.substitute(\\n assign_return_values=\\\"auto \\\" + TMP_VAR + \\\" = \\\",\\n unambiguous_name=f.func.name.unambiguous_name(),\\n unpacked_args=redispatch_args,\\n )\\n )\\n call, rhs_value = emit_view_body(fn, TMP_VAR)\\n inplace_view_body.append(call)\\n assert rhs_value is not None\\n inplace_view_body.append(\\n ASSIGN_RETURN_VALUE.substitute(\\n return_values=tie_return_values(f), rhs_value=rhs_value\\n )\\n )\\n if f.func.returns:\\n inplace_view_body.append(f\\\"return {get_return_value(f)};\\\")\\n return inplace_view_body\\n\\n\\n@with_native_function\\ndef gen_formals(f: NativeFunction) -> str:\\n return \\\", \\\".join(\\n # code-generated autograd kernels plumb and recompute dispatch keys directly through the kernel for performance.\\n # See Note [Plumbing Keys Through The Dispatcher] for details.\\n [\\\"c10::DispatchKeySet ks\\\"]\\n + [\\n f'{cpp.argument_type(a, binds=\\\"__placeholder__\\\", symint=True).cpp_type()} {a.name}'\\n for a in f.func.schema_order_arguments()\\n ]\\n )\\n\\n\\n@with_native_function_with_differentiability_info\\ndef inplace_or_view_method_definition(\\n fn: NativeFunctionWithDifferentiabilityInfo,\\n) -> str | None:\\n f = fn.func\\n if get_view_info(f) is None and (\\n # For functions that modify their inputs but don't return them,\\n # we can't give them autograd support.\\n # See https://github.com/pytorch/pytorch/issues/53796\\n not modifies_arguments(f)\\n or len(f.func.returns) == 0\\n ):\\n return None\\n return METHOD_DEFINITION.substitute(\\n return_type=cpp.returns_type(f.func.returns, symint=True).cpp_type(),\\n type_wrapper_name=type_wrapper_name(f),\\n formals=gen_formals(f),\\n type_definition_body=emit_inplace_or_view_body(fn),\\n )\\n\\n\\n@with_native_function_with_differentiability_info\\ndef inplace_or_view_method_registration(\\n fn: NativeFunctionWithDifferentiabilityInfo,\\n) -> str | None:\\n f = fn.func\\n if get_view_info(f) is None and (\\n not modifies_arguments(f) or len(f.func.returns) == 0\\n ):\\n return None\\n return WRAPPER_REGISTRATION.substitute(\\n unqual_operator_name_with_overload=f.func.name,\\n type_wrapper_name=type_wrapper_name(f),\\n class_type=\\\"ADInplaceOrView\\\",\\n )\\n\\n\\ndef use_derived(fn: NativeFunctionWithDifferentiabilityInfo) -> bool:\\n f = fn.func\\n name = cpp.name(f.func)\\n return name not in MANUAL_AUTOGRAD and dispatch_strategy(fn) == \\\"use_derived\\\"\\n\\n\\ndef gen_inplace_or_view_type_env(\\n fn: NativeFunctionWithDifferentiabilityInfo,\\n) -> dict[str, list[str]]:\\n definition = inplace_or_view_method_definition(fn)\\n registration = inplace_or_view_method_registration(fn)\\n\\n return {\\n \\\"ops_headers\\\": (\\n [f\\\"#include \\\"]\\n if definition is not None\\n else []\\n ),\\n \\\"inplace_or_view_method_definitions\\\": [definition]\\n if definition is not None\\n else [],\\n \\\"inplace_or_view_wrapper_registrations\\\": [registration]\\n if registration is not None\\n else [],\\n }\\n\\n\\ndef gen_inplace_or_view_type(\\n out: str,\\n native_yaml_path: str,\\n tags_yaml_path: str,\\n fns_with_infos: list[NativeFunctionWithDifferentiabilityInfo],\\n template_path: str,\\n) -> None:\\n # NOTE: see Note [Sharded File] at the top of the VariableType.cpp\\n # template regarding sharding of the generated files.\\n num_shards = 2\\n\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n fm.write_sharded(\\n \\\"ADInplaceOrViewType.cpp\\\",\\n [fn for fn in fns_with_infos if use_derived(fn)],\\n key_fn=lambda fn: fn.func.root_name,\\n base_env={\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/ADInplaceOrViewType.cpp\\\",\\n },\\n env_callable=gen_inplace_or_view_type_env,\\n num_shards=2,\\n sharded_keys={\\n \\\"ops_headers\\\",\\n \\\"inplace_or_view_method_definitions\\\",\\n \\\"inplace_or_view_wrapper_registrations\\\",\\n },\\n )\\n\\n\\n# Generates VariableType.h/cpp\\n#\\n# **If any changes are being made to the VariableType codegen please also check\\n# if updates are needed in torch/csrc/autograd/autograd_not_implemented_fallback.cpp\\n#\\n# VariableType is a subclass of at::Type that provides the binding code\\n# necessary to provide a differentiable version of ATen operators. There are a\\n# number of different things we could mean:\\n#\\n# - Given a non-differentiable forward implementation, we might\\n# directly associate it with a backward implementation to make\\n# it differentiable. This is the common case.\\n#\\n# - Some functions don't need a backwards implementation, because\\n# backpropagation will never propagate beyond them. There are a\\n# number of different reasons why this may be the case:\\n#\\n# - The function has no differentiable inputs\\n# - The function's output is not differentiable\\n# - The function has no data dependency on its input\\n#\\n# - Some function don't need a backwards implementation because they\\n# are implemented as a composition of other (differentiable) ATen\\n# functions. These are dispatched directly to the Type superclass,\\n# which will in turn dispatch back to VariableType for its\\n# differentiable subcomponents.\\n#\\n\\nfrom __future__ import annotations\\n\\nimport re\\nfrom typing import Callable, Sequence\\n\\nfrom torchgen.api import cpp\\nfrom torchgen.api.autograd import (\\n DifferentiableInput,\\n dispatch_strategy,\\n ForwardDerivative,\\n gen_differentiable_outputs,\\n is_differentiable,\\n NativeFunctionWithDifferentiabilityInfo,\\n SavedAttribute,\\n)\\nfrom torchgen.api.types import (\\n ArrayRefCType,\\n BaseCppType,\\n BaseCType,\\n Binding,\\n DispatcherSignature,\\n intArrayRefT,\\n iTensorListRefT,\\n ListCType,\\n MutRefCType,\\n OptionalCType,\\n scalarT,\\n SpecialArgName,\\n stringT,\\n symIntArrayRefT,\\n TENSOR_LIST_LIKE_CTYPES,\\n tensorListT,\\n tensorT,\\n TupleCType,\\n VectorCType,\\n)\\nfrom torchgen.code_template import CodeTemplate\\nfrom torchgen.context import (\\n native_function_manager,\\n with_native_function,\\n with_native_function_and,\\n)\\nfrom torchgen.model import (\\n Argument,\\n BaseType,\\n ListType,\\n NativeFunction,\\n SchemaKind,\\n SelfArgument,\\n TensorOptionsArguments,\\n)\\nfrom torchgen.utils import FileManager, mapMaybe\\n\\nfrom .context import with_native_function_with_differentiability_info_and_key\\nfrom .gen_inplace_or_view_type import (\\n ALL_VIEW_FUNCTIONS,\\n ASSIGN_RETURN_VALUE,\\n AUTOGRAD_NOT_IMPLEMENTED_REGISTRATION,\\n gen_formals,\\n get_base_name,\\n get_view_info,\\n is_tensor_list_type,\\n is_tensor_type,\\n METHOD_DEFINITION,\\n modifies_arguments,\\n TMP_VAR,\\n unpack_args,\\n unpacked_name,\\n use_derived,\\n WRAPPER_REGISTRATION,\\n)\\nfrom .gen_trace_type import (\\n get_return_value,\\n MANUAL_AUTOGRAD_AND_TRACER,\\n MANUAL_BACKEND,\\n tie_return_values,\\n type_wrapper_name,\\n)\\n\\n\\n# We don't set or modify grad_fn on these methods. Generally, they return\\n# tensors that have requires_grad=False. In-place functions listed here will\\n# not examine or modify requires_grad or grad_fn.\\n# NB: this does NOT include overload name\\nDONT_REQUIRE_DERIVATIVE = {\\n # These only depend on the input Tensor's shape and device, not the data\\n \\\"empty_like\\\",\\n \\\"ones_like\\\",\\n \\\"full_like\\\",\\n \\\"zeros_like\\\",\\n \\\"rand_like\\\",\\n \\\"randn_like\\\",\\n \\\"new_empty\\\",\\n \\\"new_empty_strided\\\",\\n \\\"new_full\\\",\\n \\\"new_zeros\\\",\\n \\\"new_ones\\\",\\n # These are only implemented on integral types\\n \\\"__and__\\\",\\n \\\"__iand__\\\",\\n \\\"__ilshift__\\\",\\n \\\"__ior__\\\",\\n \\\"__irshift__\\\",\\n \\\"__ixor__\\\",\\n \\\"__lshift__\\\",\\n \\\"__or__\\\",\\n \\\"__rshift__\\\",\\n \\\"__xor__\\\",\\n # These work on integral data types, and hence don't require derivative\\n \\\"_sobol_engine_draw\\\",\\n \\\"_sobol_engine_ff\\\",\\n \\\"_sobol_engine_scramble_\\\",\\n \\\"_sobol_engine_initialize_state_\\\",\\n # This is an unsafe method that is meant to be out of reach of autograd.\\n \\\"_coalesced_\\\",\\n # Quantize functions should not record gradients\\n \\\"quantize_per_tensor\\\",\\n \\\"quantize_per_channel\\\",\\n # Functions that return integers should not have output that require gradients\\n \\\"argmax\\\",\\n \\\"argmin\\\",\\n \\\"argsort\\\",\\n \\\"searchsorted\\\",\\n \\\"bucketize\\\",\\n # Functions that return booleans are not differentiable\\n \\\"isnan\\\",\\n \\\"isposinf\\\",\\n \\\"isneginf\\\",\\n \\\"isinf\\\",\\n \\\"signbit\\\",\\n \\\"isin\\\",\\n \\\"allclose\\\",\\n # Functions return none are not differentiable\\n \\\"record_stream\\\",\\n # These functions are not differentiable\\n \\\"logical_and\\\",\\n \\\"logical_xor\\\",\\n \\\"logical_not\\\",\\n \\\"logical_or\\\",\\n # This function returns nested_tensor shape as a tensor that is non-differentiable\\n \\\"_nested_tensor_size\\\",\\n \\\"_nested_tensor_strides\\\",\\n \\\"_nested_tensor_storage_offsets\\\",\\n}\\n\\n# The C -> R functions at the time of adding this are still being audited and tested\\n# but will not error out.\\n# C -> C, R -> C functions for which backward is correctly implemented and tested\\nGRADIENT_IMPLEMENTED_FOR_COMPLEX = {\\n \\\"fill\\\",\\n \\\"t\\\",\\n \\\"t_copy\\\",\\n \\\"view\\\",\\n \\\"reshape\\\",\\n \\\"reshape_as\\\",\\n \\\"view_as\\\",\\n \\\"view_copy\\\",\\n \\\"roll\\\",\\n \\\"clone\\\",\\n \\\"block_diag\\\",\\n \\\"diag_embed\\\",\\n \\\"repeat\\\",\\n \\\"expand\\\",\\n \\\"expand_copy\\\",\\n \\\"flip\\\",\\n \\\"fliplr\\\",\\n \\\"flipud\\\",\\n \\\"rot90\\\",\\n \\\"nanmean\\\",\\n \\\"nansum\\\",\\n \\\"transpose\\\",\\n \\\"permute\\\",\\n \\\"squeeze\\\",\\n \\\"unsqueeze\\\",\\n \\\"unsqueeze_copy\\\",\\n \\\"resize\\\",\\n \\\"resize_as\\\",\\n \\\"tril\\\",\\n \\\"triu\\\",\\n \\\"chunk\\\",\\n \\\"zero_\\\",\\n \\\"eq_\\\",\\n \\\"ne_\\\",\\n \\\"add\\\",\\n \\\"__radd__\\\",\\n \\\"sum\\\",\\n \\\"_conj\\\",\\n \\\"sin\\\",\\n \\\"cos\\\",\\n \\\"mul\\\",\\n \\\"sinc\\\",\\n \\\"sinh\\\",\\n \\\"cosh\\\",\\n \\\"__rmul__\\\",\\n \\\"sgn\\\",\\n \\\"asin\\\",\\n \\\"acos\\\",\\n \\\"sub\\\",\\n \\\"div\\\",\\n \\\"cat\\\",\\n \\\"view_as_complex\\\",\\n \\\"index_put\\\",\\n \\\"neg\\\",\\n \\\"complex\\\",\\n \\\"select\\\",\\n \\\"where\\\",\\n \\\"as_strided\\\",\\n \\\"as_strided_copy\\\",\\n \\\"as_strided_scatter\\\",\\n \\\"slice\\\",\\n \\\"constant_pad_nd\\\",\\n \\\"unbind\\\",\\n \\\"split\\\",\\n \\\"split_with_sizes\\\",\\n \\\"unsafe_split\\\",\\n \\\"split_with_sizes_backward\\\",\\n \\\"dot\\\",\\n \\\"vdot\\\",\\n \\\"cholesky\\\",\\n \\\"triangular_solve\\\",\\n \\\"mm\\\",\\n \\\"_unsafe_view\\\",\\n \\\"mv\\\",\\n \\\"outer\\\",\\n \\\"bmm\\\",\\n \\\"diagonal\\\",\\n \\\"alias\\\",\\n \\\"atan\\\",\\n \\\"log\\\",\\n \\\"log10\\\",\\n \\\"log1p\\\",\\n \\\"log2\\\",\\n \\\"logaddexp\\\",\\n \\\"logsumexp\\\",\\n \\\"logcumsumexp\\\",\\n \\\"reciprocal\\\",\\n \\\"tan\\\",\\n \\\"pow\\\",\\n \\\"rsqrt\\\",\\n \\\"tanh\\\",\\n \\\"tanh_backward\\\",\\n \\\"asinh\\\",\\n \\\"acosh\\\",\\n \\\"atanh\\\",\\n \\\"take\\\",\\n \\\"fill_\\\",\\n \\\"exp\\\",\\n \\\"exp2\\\",\\n \\\"expm1\\\",\\n \\\"nonzero\\\",\\n \\\"mean\\\",\\n \\\"std_mean\\\",\\n \\\"var_mean\\\",\\n \\\"inverse\\\",\\n \\\"solve\\\",\\n \\\"linalg_cholesky\\\",\\n \\\"addcmul\\\",\\n \\\"addcdiv\\\",\\n \\\"matrix_exp\\\",\\n \\\"linalg_matrix_exp\\\",\\n \\\"_linalg_eigh\\\",\\n \\\"cholesky_solve\\\",\\n \\\"linalg_qr\\\",\\n \\\"_linalg_svd\\\",\\n \\\"_fft_c2c\\\",\\n \\\"_fft_r2c\\\",\\n \\\"linalg_solve\\\",\\n \\\"sqrt\\\",\\n \\\"stack\\\",\\n \\\"gather\\\",\\n \\\"index_select\\\",\\n \\\"index_add_\\\",\\n \\\"linalg_inv\\\",\\n \\\"linalg_inv_ex\\\",\\n \\\"baddbmm\\\",\\n \\\"addbmm\\\",\\n \\\"addmm\\\",\\n \\\"addmv\\\",\\n \\\"addr\\\",\\n \\\"linalg_householder_product\\\",\\n \\\"ormqr\\\",\\n \\\"reflection_pad1d\\\",\\n \\\"reflection_pad2d\\\",\\n \\\"reflection_pad3d\\\",\\n \\\"linalg_cholesky_ex\\\",\\n \\\"linalg_eig\\\",\\n \\\"diagonal_copy\\\",\\n \\\"diagonal_scatter\\\",\\n \\\"alias_copy\\\",\\n \\\"select_backward\\\",\\n \\\"diagonal_backward\\\",\\n \\\"slice_backward\\\",\\n \\\"reflection_pad1d_backward\\\",\\n \\\"reflection_pad2d_backward\\\",\\n \\\"reflection_pad3d_backward\\\",\\n \\\"_sparse_sparse_matmul\\\",\\n \\\"replication_pad1d\\\",\\n \\\"replication_pad2d\\\",\\n \\\"replication_pad3d\\\",\\n \\\"put\\\",\\n \\\"put_\\\",\\n \\\"_to_copy\\\",\\n \\\"replication_pad1d_backward\\\",\\n \\\"replication_pad2d_backward\\\",\\n \\\"replication_pad3d_backward\\\",\\n \\\"diag\\\",\\n \\\"masked_scatter\\\",\\n \\\"masked_select\\\",\\n \\\"index_add\\\",\\n \\\"index_fill\\\",\\n \\\"trace\\\",\\n \\\"polar\\\",\\n \\\"cumsum\\\",\\n \\\"rsub\\\",\\n \\\"eig\\\",\\n \\\"lerp\\\",\\n \\\"linalg_vector_norm\\\",\\n \\\"cumprod\\\",\\n \\\"prod\\\",\\n \\\"index_copy\\\",\\n \\\"lu\\\",\\n \\\"unfold\\\",\\n \\\"unfold_backward\\\",\\n \\\"index\\\",\\n \\\"masked_fill\\\",\\n \\\"masked_scatter_backward\\\",\\n \\\"linalg_cross\\\",\\n \\\"lu_unpack\\\",\\n \\\"renorm\\\",\\n \\\"_conj_physical\\\",\\n \\\"linalg_lu_factor_ex\\\",\\n \\\"scatter\\\",\\n \\\"scatter_add\\\",\\n \\\"sigmoid\\\",\\n \\\"sigmoid_backward\\\",\\n \\\"sparse_mask\\\",\\n \\\"trapezoid\\\",\\n \\\"cumulative_trapezoid\\\",\\n \\\"conj_physical_\\\",\\n \\\"_neg_view\\\",\\n \\\"_reshape_alias\\\",\\n \\\"_reshape_copy\\\",\\n \\\"_linalg_det\\\",\\n \\\"lu_solve\\\",\\n \\\"linalg_solve_triangular\\\",\\n \\\"linalg_pinv\\\",\\n \\\"linalg_lstsq\\\",\\n \\\"unfold_copy\\\",\\n \\\"col2im\\\",\\n \\\"im2col\\\",\\n \\\"cholesky_inverse\\\",\\n \\\"to_sparse\\\",\\n \\\"sparse_sampled_addmm\\\",\\n \\\"linalg_lu\\\",\\n \\\"pixel_shuffle\\\",\\n \\\"pixel_unshuffle\\\",\\n \\\"channel_shuffle\\\",\\n \\\"linalg_lu_solve\\\",\\n \\\"_linalg_slogdet\\\",\\n \\\"_linalg_solve_ex\\\",\\n \\\"_unsafe_index\\\",\\n \\\"_unsafe_index_put\\\",\\n \\\"_unsafe_masked_index\\\",\\n \\\"_unsafe_masked_index_put_accumulate\\\",\\n}\\n\\nGRADIENT_IMPLEMENTED_FOR_SPARSE_COMPLEX = {\\n \\\"_to_dense\\\",\\n \\\"_coalesce\\\",\\n \\\"coalesce\\\",\\n \\\"values\\\",\\n \\\"_sparse_coo_tensor_with_dims_and_tensors\\\",\\n \\\"_sparse_addmm\\\",\\n}\\n\\nGRADIENT_IMPLEMENTED_FOR_COMPLEX.update(GRADIENT_IMPLEMENTED_FOR_SPARSE_COMPLEX)\\n\\n# Some operators invalidate the grad_accumulator. Let's reset it.\\nRESET_GRAD_ACCUMULATOR = {\\\"set_\\\", \\\"resize_\\\"}\\n\\n# NOTE [ TensorImpl and Storage Pointer Sanity Checks ]\\n#\\n# We check the following properties:\\n# 1) A function should never change the input tensors' underlying c10::TensorImpl\\n# pointers or c10::Storage pointers, even if it modifies its input tensors (via\\n# inplace or out-variants)\\n# If the function does not modify its arguments, we also check the following properties\\n# pertaining to its output:\\n# 2) Its TensorImpl has use_count of 1\\n# 3) If the function is a view function, it has the same StorageImpl as that of\\n# the input it is aliased with. Otherwise, its StorageImpl has use_count of 1\\n#\\n# The following code templates implement the checks for this invariant:\\nSAVE_TENSOR_STORAGE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nauto ${tensor_name}_storage_saved =\\n ${tensor_name}.has_storage() ? ::std::optional(${tensor_name}.storage()) : ::std::nullopt;\\n\\\"\\\"\\\"\\n)\\n\\n\\n# If tensor_name == out_tensor_name, used to enforce (1), otherwise used for (2)\\nENFORCE_SAME_TENSOR_STORAGE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (${tensor_name}_storage_saved.has_value() &&\\n !at::impl::dispatch_mode_enabled() &&\\n !at::impl::tensor_has_dispatch(${tensor_name}) &&\\n !at::impl::tensor_has_dispatch(${out_tensor_name}))\\n TORCH_INTERNAL_ASSERT(${tensor_name}_storage_saved.value().is_alias_of(${out_tensor_name}.storage()));\\n\\\"\\\"\\\"\\n)\\n\\nSAVE_TENSORLIST_STORAGE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstd::vector<::std::optional> ${tensorlist_name}_storage_saved(${tensorlist_name}.size());\\nfor (const Tensor& tensor : ${tensorlist_name})\\n ${tensorlist_name}_storage_saved.push_back(\\n tensor.has_storage() ? ::std::optional(tensor.storage()) : ::std::nullopt);\\n\\\"\\\"\\\"\\n)\\n\\nENFORCE_SAME_TENSORLIST_STORAGE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfor (size_t i=0; i<${tensorlist_name}.size() && !at::impl::dispatch_mode_enabled(); i++) {\\n if (${tensorlist_name}_storage_saved[i].has_value() && !at::impl::tensorlist_has_dispatch(${tensorlist_name}))\\n TORCH_INTERNAL_ASSERT(${tensorlist_name}_storage_saved[i].value().is_alias_of(${tensorlist_name}[i].storage()));\\n}\\n\\\"\\\"\\\"\\n)\\n\\nSAVE_OPTIONALTENSORLIST_STORAGE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstd::vector<::std::optional> ${tensorlist_name}_storage_saved(${tensorlist_name}.size());\\nfor (const ::std::optional& tensor : ${tensorlist_name})\\n ${tensorlist_name}_storage_saved.push_back(\\n tensor.has_value() && tensor->has_storage() ? ::std::optional(tensor->storage()) : ::std::nullopt);\\n\\\"\\\"\\\"\\n)\\n\\nENFORCE_SAME_OPTIONALTENSORLIST_STORAGE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfor (size_t i=0; i<${tensorlist_name}.size() && !at::impl::dispatch_mode_enabled(); i++) {\\n if (${tensorlist_name}_storage_saved[i].has_value() && !at::impl::tensorlist_has_dispatch(${tensorlist_name}))\\n TORCH_INTERNAL_ASSERT(${tensorlist_name}_storage_saved[i].value().is_alias_of(\\n static_cast<::std::optional>(${tensorlist_name}[i])->storage()));\\n}\\n\\\"\\\"\\\"\\n)\\n\\nSAVE_TENSOR_IMPL = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nc10::intrusive_ptr ${tensor_name}_impl_saved;\\nif (${tensor_name}.defined()) ${tensor_name}_impl_saved = ${tensor_name}.getIntrusivePtr();\\n\\\"\\\"\\\"\\n)\\n\\nENFORCE_SAME_TENSOR_IMPL = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (${tensor_name}_impl_saved && !at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(${tensor_name}))\\n TORCH_INTERNAL_ASSERT(${tensor_name}_impl_saved == ${tensor_name}.getIntrusivePtr());\\n\\\"\\\"\\\"\\n)\\n\\nENFORCE_TENSOR_IMPL_USE_COUNT_LT_OR_EQ_ONE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (!at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(${tensor_name}))\\n TORCH_INTERNAL_ASSERT(${tensor_name}.use_count() <= 1, \\\"function: ${fn_name}\\\");\\n\\\"\\\"\\\"\\n)\\n\\nENFORCE_TENSOR_STORAGE_USE_COUNT_EQUALS_ONE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (${tensor_name}.has_storage() && !at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(${tensor_name})) {\\n TORCH_INTERNAL_ASSERT(${tensor_name}.storage().use_count() == 1, \\\"function: ${fn_name}\\\");\\n}\\n\\\"\\\"\\\"\\n)\\n\\nSAVE_TENSORLIST_IMPL = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstd::vector> ${tensorlist_name}_impl_saved(${tensorlist_name}.size());\\nfor (size_t i=0; i<${tensorlist_name}.size(); i++)\\n if (${tensorlist_name}[i].defined()) ${tensorlist_name}_impl_saved[i] = ${tensorlist_name}[i].getIntrusivePtr();\\n\\\"\\\"\\\"\\n)\\n\\nENFORCE_SAME_TENSORLIST_IMPL = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfor (size_t i=0; i<${tensorlist_name}.size() && !at::impl::dispatch_mode_enabled(); i++) {\\n if (${tensorlist_name}_impl_saved[i] && !at::impl::tensorlist_has_dispatch(${tensorlist_name}))\\n TORCH_INTERNAL_ASSERT(${tensorlist_name}_impl_saved[i] == ${tensorlist_name}[i].getIntrusivePtr());\\n}\\n\\\"\\\"\\\"\\n)\\n\\nSAVE_OPTIONALTENSORLIST_IMPL = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstd::vector> ${tensorlist_name}_impl_saved(${tensorlist_name}.size());\\nfor (size_t i=0; i<${tensorlist_name}.size(); i++) {\\n ::std::optional t = ${tensorlist_name}[i];\\n if (t.has_value() && t->defined()) ${tensorlist_name}_impl_saved[i] = t->getIntrusivePtr();\\n}\\n\\\"\\\"\\\"\\n)\\n\\nENFORCE_SAME_OPTIONALTENSORLIST_IMPL = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfor (size_t i=0; i<${tensorlist_name}.size() && !at::impl::dispatch_mode_enabled(); i++) {\\n if (${tensorlist_name}_impl_saved[i])\\n TORCH_INTERNAL_ASSERT(\\n ${tensorlist_name}_impl_saved[i] == static_cast<::std::optional>(${tensorlist_name}[i])->getIntrusivePtr());\\n}\\n\\\"\\\"\\\"\\n)\\n\\n# The following list contains functions that we don't enforce the invariant on.\\nDONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE = {\\n # These functions are expected to change impl or storage of input tensors\\n \\\"set_\\\",\\n \\\"_cudnn_rnn_flatten_weight\\\",\\n \\\"_unsafe_masked_index\\\",\\n \\\"_unsafe_masked_index_put_accumulate\\\",\\n}\\nDONT_ENFORCE_TENSOR_IMPL_USE_COUNT = {\\n # These non-inplace, non-out functions return tensors with use_count > 1\\n # Therefore, they MAY (but not necessarily) return one of its inputs as-is\\n # See https://github.com/pytorch/pytorch/issues/60426 for more information\\n \\\"_embedding_bag\\\",\\n \\\"_embedding_bag_forward_only\\\",\\n \\\"q_per_channel_scales\\\",\\n \\\"q_per_channel_zero_points\\\",\\n \\\"lu_unpack\\\",\\n \\\"_cudnn_rnn_backward\\\",\\n # The below failed StorageImpl use_count check but we skip tensor_impl check\\n # just in case\\n \\\"_cudnn_rnn\\\",\\n \\\"dequantize_self\\\",\\n # lift() should never actually be called with a requires_grad=True tensor,\\n \\\"lift\\\",\\n \\\"lift_fresh\\\",\\n \\\"lift_fresh_copy\\\",\\n # Nested Tensors related functions\\n # _nested_tensor_size() should never actually be called with requires_grad=True tensor\\n \\\"_nested_tensor_size\\\",\\n \\\"_nested_tensor_strides\\\",\\n \\\"_nested_tensor_storage_offsets\\\",\\n}\\n\\nDONT_ENFORCE_STORAGE_IMPL_USE_COUNT = {\\n # These non-view functions return tensors with storage use_count != 1\\n \\\"_slow_conv2d_forward\\\",\\n \\\"slow_conv3d_forward\\\",\\n \\\"channel_shuffle\\\",\\n # If an input is returned as-is in output, we cannot guarantee its storage_impl\\n # use count to be 1 either.\\n *DONT_ENFORCE_TENSOR_IMPL_USE_COUNT,\\n}\\n# END CHECKS FOR [ TensorImpl and Storage Pointer Sanity Checks ]\\n\\nDECLARE_GRAD_FN = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstd::shared_ptr<${op}> grad_fn;\\n\\\"\\\"\\\"\\n)\\n\\nDECLARE_VECTOR_OF_GRAD_FN = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nstd::vector> grad_fns;\\n\\\"\\\"\\\"\\n)\\n\\nSETUP_ANY_REQUIRES_GRAD = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n[[maybe_unused]] auto _any_requires_grad = compute_requires_grad( ${args_with_derivatives} );\\n${extra_differentiability_conditions}\\n\\\"\\\"\\\"\\n)\\n\\nSETUP_DERIVATIVE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (_any_requires_grad) {\\n ${setup}\\n}\\n\\\"\\\"\\\"\\n)\\n\\nSETUP_NONE_REQUIRES_GRAD = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (compute_requires_grad( ${args_to_check} )) {\\n throw_error_out_requires_grad(\\\"${base_name}\\\");\\n}\\n\\\"\\\"\\\"\\n)\\n\\nASSIGN_GRAD_FN = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\ngrad_fn = std::shared_ptr<${op}>(new ${op}(${op_ctor}), deleteNode);\\ngrad_fn->set_next_edges(collect_next_edges( ${args_with_derivatives} ));\\n\\\"\\\"\\\"\\n)\\n\\n# note(crcrpar): `compute_requires_grad` in the template below is supplied with arguments indexed with `i`\\n# while the `SETUP_ANY_REQUIRES_GRAD` above takes whole tensors and scalars.\\nASSIGN_VECTOR_OF_GRAD_FN = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfor (const auto& i : c10::irange( ${irange} )) {\\n const auto ith_requires_grad = compute_requires_grad(${args_with_derivatives});\\n check_inplace(self[i], ith_requires_grad);\\n grad_fns.push_back([&]() -> std::shared_ptr<${op}> {\\n if (!ith_requires_grad) {\\n return nullptr;\\n } else {\\n auto grad_fn = std::shared_ptr<${op}>(new ${op}(${op_ctor}), deleteNode);\\n grad_fn->set_next_edges(collect_next_edges( ${args_with_derivatives} ));\\n return grad_fn;\\n }\\n }());\\n}\\n\\\"\\\"\\\"\\n)\\n\\nCALL_REDISPATCH = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nat::redispatch::${api_name}(${unpacked_args})\\\"\\\"\\\"\\n)\\n# If the non-variable operation has return values, we use the `tmp` variable to hold the\\n# values temporarily and pass the values to the return variables outside of the\\n# `at::AutoDispatchBelowAutograd` guard block.\\nDISPATCH_TO_NON_VAR_TYPE_WITH_TMP_RETURN_VALUES_JVP_DECOMP = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nauto ${tmp_var} = ([&]() {\\n if (${any_has_forward_grad}) {\\n static c10::OperatorName full_name(\\\"aten::${op_name}\\\", \\\"${op_overload}\\\");\\n static ::std::optional opt_op = c10::Dispatcher::singleton().findSchema(full_name);\\n return impl::run_jit_decomposition_with_args_for_jvp<${return_types}>(\\\"${op_name}\\\", *opt_op, ks, ${arg_names});\\n } else {\\n ${guard}\\n return ${base_type_call};\\n }\\n})();\\n\\\"\\\"\\\"\\n)\\n\\nDISPATCH_TO_NON_VAR_TYPE_WITH_TMP_RETURN_VALUES = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nauto ${tmp_var} = ([&]() {\\n ${guard}\\n return ${base_type_call};\\n})();\\n\\\"\\\"\\\"\\n)\\n\\nDISPATCH_TO_NON_VAR_TYPE_WITHOUT_RETURN_VALUES = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n{\\n ${guard}\\n ${base_type_call};\\n}\\n\\\"\\\"\\\"\\n)\\n\\nSET_HISTORY = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (grad_fn) {\\n ${fn}_history(${differentiable_outputs}, grad_fn);\\n}\\n\\\"\\\"\\\"\\n)\\n\\nLOOP_OVER_VECTOR_OF_GRAD_FNS = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (!grad_fns.empty()) {\\n ${preamble}\\n for (const auto& i : c10::irange(grad_fns.size())) {\\n auto grad_fn = grad_fns[i];\\n if (grad_fn != nullptr) {\\n ${statements}\\n }\\n }\\n}\\n\\\"\\\"\\\"\\n)\\n\\nCONDITIONAL = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (${cond}) {\\n ${statements}\\n}\\n\\\"\\\"\\\"\\n)\\n\\nRUN_ONLY_IN_DEBUG_MODE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n#ifndef NDEBUG\\n${statements}\\n#endif\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_CHECK_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nisFwGradDefined(${req_inp})\\\\\\n\\\"\\\"\\\"\\n)\\nFW_DERIVATIVE_SIZE_CHECK_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nTORCH_CHECK(\\n self.size() == ${inp_name}.size(),\\n \\\"Tensor lists must have the same number of tensors, got \\\",\\n self.size(),\\n \\\" and \\\",\\n ${inp_name}.size());\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_TENSORLIST_CHECK_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nisFwGradDefinedTensorList(${req_inp})\\\\\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_DEFINED_GRAD_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nauto ${inp_name}_t_raw = toNonOptFwGrad(${inp});\\nauto ${inp_name}_tensor = toNonOptTensor(${inp});\\nauto ${inp_name}_t = (${inp_name}_t_raw.defined() || !${inp_name}_tensor.defined())\\n ? ${inp_name}_t_raw : at::${zeros_fn}(${inp_name}_tensor.sym_sizes(), ${inp_name}_tensor.options());\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_DEFINED_PRIMAL_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nauto ${inp_name}_p = toNonOptPrimal(${inp});\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_SETTER_TENSOR = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (${out_arg}_new_fw_grad_opt.has_value() && ${out_arg}_new_fw_grad_opt.value().defined() && ${out_arg}.defined()) {\\n // The hardcoded 0 here will need to be updated once we support multiple levels.\\n ${out_arg}._set_fw_grad(${out_arg}_new_fw_grad_opt.value(), /* level */ 0, /* is_inplace_op */ ${is_inplace});\\n}\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_SETTER_TENSOR_FOREACH = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfor (const auto& i : c10::irange(${out_arg}_new_fw_grad_opts.size())) {\\n auto& ${out_arg}_new_fw_grad_opt = ${out_arg}_new_fw_grad_opts[i];\\n if (${out_arg}_new_fw_grad_opt.has_value() && ${out_arg}_new_fw_grad_opt.value().defined() && ${out_arg}[i].defined()) {\\n // The hardcoded 0 here will need to be updated once we support multiple levels.\\n ${out_arg}[i]._set_fw_grad(${out_arg}_new_fw_grad_opt.value(), /* level */ 0, /* is_inplace_op */ ${is_inplace});\\n }\\n}\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_SETTER_MULTI_OUTPUT = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (${all_res}_new_fw_grad_opt.has_value() && std::get<${idx}>(${all_res}_new_fw_grad_opt.value()).defined()\\n && ${out_arg}.defined()) {\\n ${out_arg}._set_fw_grad(std::get<${idx}>(${all_res}_new_fw_grad_opt.value()), /* level */ 0, /* is_inplace_op */ false);\\n}\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_SETTER_TENSOR_LIST = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nif (${out_arg}_new_fw_grad_opt.has_value()) {\\n auto ${out_arg}_new_fw_grad = ${out_arg}_new_fw_grad_opt.value();\\n TORCH_INTERNAL_ASSERT(${out_arg}.size() == ${out_arg}_new_fw_grad.size());\\n for (const auto i : c10::irange(${out_arg}.size())) {\\n if (${out_arg}_new_fw_grad[i].defined() && ${out_arg}[i].defined()) {\\n // The hardcoded 0 here will need to be updated once we support multiple levels.\\n ${out_arg}[i]._set_fw_grad(${out_arg}_new_fw_grad[i], /* level */ 0, /* is_inplace_op */ ${is_inplace});\\n }\\n }\\n}\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${fw_grad_opt_definition}\\nif (${requires_fw_grad}) {\\n ${unpacked_arguments}\\n ${out_arg}_new_fw_grad_opt = ${formula};\\n}\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_FOREACH_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n${fw_grad_opt_definition}\\nfor (const auto& i : c10::irange(${vector_of_optional_tensor}.size())) {\\n if (${any_has_forward_grad_for_current_index}) {\\n ${unpacked_arguments}\\n ${vector_of_optional_tensor}[i] = ${formula};\\n }\\n}\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_FORBID_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nTORCH_CHECK_NOT_IMPLEMENTED(!(${cond}), \\\"Trying to use forward AD with ${name} that does not support it ${msg}\\\");\\n\\\"\\\"\\\"\\n)\\n\\nFW_DERIVATIVE_FORBID_LIST_TEMPLATE = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nfor (const auto& _t: ${arg}) {\\n TORCH_CHECK_NOT_IMPLEMENTED(!(${cond}), \\\"Trying to use forward AD with ${name} that does not support it ${msg}\\\");\\n}\\n\\\"\\\"\\\"\\n)\\n\\n\\ndef gen_variable_type(\\n out: str,\\n native_yaml_path: str,\\n tags_yaml_path: str,\\n fns_with_diff_infos: list[NativeFunctionWithDifferentiabilityInfo],\\n template_path: str,\\n used_keys: set[str],\\n) -> None:\\n \\\"\\\"\\\"VariableType.h and VariableType.cpp body\\n\\n This is the at::Type subclass for differentiable tensors. The\\n implementation of each function dispatches to the base tensor type to\\n compute the output. The grad_fn is attached to differentiable functions.\\n \\\"\\\"\\\"\\n fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False)\\n fm.write(\\n \\\"VariableType.h\\\",\\n lambda: {\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/VariableType.h\\\"\\n },\\n )\\n\\n # helper that generates a TORCH_LIBRARY_IMPL macro for each\\n # dispatch key that appears in derivatives.yaml\\n def wrapper_registrations(used_keys: set[str]) -> str:\\n library_impl_macro_list: list[str] = []\\n for key in sorted(used_keys):\\n dispatch_key = key\\n if key == \\\"Default\\\":\\n dispatch_key = \\\"Autograd\\\"\\n library_impl_macro = (\\n f\\\"TORCH_LIBRARY_IMPL(aten, {dispatch_key}, m) \\\"\\n + \\\"{\\\\n\\\"\\n + \\\"${\\\"\\n + f\\\"wrapper_registrations_{key}\\\"\\n + \\\"}\\\\n}\\\"\\n )\\n library_impl_macro_list += [library_impl_macro]\\n return \\\"\\\\n\\\\n\\\".join(library_impl_macro_list)\\n\\n # Generate a new template from VariableType.cpp which replaces ${wrapper_registrations}\\n # with per key TORCH_LIBRARY_IMPL macros for each key that appears in derivatives.yaml\\n fm1 = FileManager(\\n install_dir=out + \\\"/templates\\\", template_dir=template_path, dry_run=False\\n )\\n fm1.write(\\n \\\"VariableType.cpp\\\",\\n lambda: {\\n \\\"type_derived_method_definitions\\\": \\\"\\\\n\\\\n\\\".join(\\n [\\n \\\"${\\\" + f\\\"type_derived_method_definitions_{key}\\\" + \\\"}\\\"\\n for key in sorted(used_keys)\\n ]\\n ),\\n \\\"wrapper_registrations\\\": wrapper_registrations(used_keys),\\n },\\n )\\n\\n # Generate final VariableType_*.cpp files from the generated template\\n fm2 = FileManager(install_dir=out, template_dir=out + \\\"/templates\\\", dry_run=False)\\n\\n sharded_keys = set(\\n [f\\\"type_derived_method_definitions_{key}\\\" for key in sorted(used_keys)]\\n + [f\\\"wrapper_registrations_{key}\\\" for key in sorted(used_keys)]\\n )\\n # NOTE: see Note [Sharded File] at the top of the VariableType.cpp\\n # template regarding sharding of the generated files.\\n fm2.write_sharded(\\n \\\"VariableType.cpp\\\",\\n [fn for fn in fns_with_diff_infos if use_derived(fn)],\\n key_fn=lambda fn: cpp.name(fn.func.func),\\n base_env={\\n \\\"generated_comment\\\": \\\"@\\\"\\n + f\\\"generated from {fm.template_dir_for_comments()}/VariableType.cpp\\\",\\n },\\n env_callable=gen_variable_type_func,\\n num_shards=5,\\n sharded_keys=sharded_keys,\\n )\\n\\n\\n@with_native_function_and\\ndef gen_wrapper_registration(f: NativeFunction, key: str = \\\"Default\\\") -> str:\\n return WRAPPER_REGISTRATION.substitute(\\n unqual_operator_name_with_overload=f.func.name,\\n type_wrapper_name=type_wrapper_name(f, key),\\n class_type=\\\"VariableType\\\",\\n )\\n\\n\\ndef gen_variable_type_func(\\n fn: NativeFunctionWithDifferentiabilityInfo,\\n) -> dict[str, list[str]]:\\n f = fn.func\\n result = {}\\n with native_function_manager(f):\\n name = cpp.name(f.func)\\n formals = gen_formals(f)\\n\\n if (\\n fn.info is None\\n and str(f.func.name.name) not in RESET_GRAD_ACCUMULATOR\\n and get_base_name(f) not in DONT_REQUIRE_DERIVATIVE\\n and len(gen_differentiable_outputs(fn)) > 0\\n and cpp.name(f.func) not in DONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE\\n and type_wrapper_name(f) not in DONT_ENFORCE_STORAGE_IMPL_USE_COUNT\\n and type_wrapper_name(f) not in DONT_ENFORCE_TENSOR_IMPL_USE_COUNT\\n ):\\n # NOTE: [ Registering AutogradNotImplemented boxed kernel ]\\n #\\n # When there is no derivatives.yaml entry, we register a generic boxed\\n # NotImplemented kernel to set grad_fn to be NotImplemented, so that forward\\n # proceeds as usual but an error is properly produced on backward.\\n # TODO: it would be nice to not have these special cases\\n #\\n # There are several cases where still let codegen handle it:\\n # 1) ops that need to reset grad accumulator (we let codegen handle this case\\n # because) the list is (currently) only accessible in Python.\\n # 2) User explicitly specifies DONT_REQUIRE_DERIVATIVE. This basically makes\\n # autograd a fallthrough with NDEBUG checks. This can be useful for when all\\n # outputs are integral.\\n # 3) When there are no differentiable outputs. This is similar to (2).\\n # 4) There are certain ops where we skip certain NDEBUG checks. this is similar\\n # to (1).\\n type_definition = \\\"\\\"\\n wrapper_registration = AUTOGRAD_NOT_IMPLEMENTED_REGISTRATION.substitute(\\n unqual_operator_name_with_overload=f.func.name\\n )\\n result[\\\"type_derived_method_definitions_Default\\\"] = [type_definition]\\n result[\\\"wrapper_registrations_Default\\\"] = [wrapper_registration]\\n else:\\n if not fn.info:\\n key = \\\"Default\\\"\\n type_definition = METHOD_DEFINITION.substitute(\\n return_type=cpp.returns_type(\\n f.func.returns, symint=True\\n ).cpp_type(),\\n type_wrapper_name=type_wrapper_name(f, key),\\n type_definition_body=emit_body(fn, key),\\n formals=formals,\\n )\\n wrapper_registration = gen_wrapper_registration(f, key)\\n result[f\\\"type_derived_method_definitions_{key}\\\"] = [type_definition]\\n result[f\\\"wrapper_registrations_{key}\\\"] = [wrapper_registration]\\n else:\\n for key in fn.info.keys():\\n type_definition = METHOD_DEFINITION.substitute(\\n return_type=cpp.returns_type(\\n f.func.returns, symint=True\\n ).cpp_type(),\\n type_wrapper_name=type_wrapper_name(f, key),\\n type_definition_body=emit_body(fn, key),\\n formals=formals,\\n )\\n wrapper_registration = gen_wrapper_registration(f, key)\\n result[f\\\"type_derived_method_definitions_{key}\\\"] = [type_definition]\\n result[f\\\"wrapper_registrations_{key}\\\"] = [wrapper_registration]\\n # See Note [Manual Backend kernels]\\n assert (name in MANUAL_BACKEND) == f.manual_kernel_registration\\n # If you want to register a kernel to Autograd, you must make the op abstract.\\n # In other words, this op must have dispatch section in native_functions.yaml.\\n if name in MANUAL_AUTOGRAD_AND_TRACER or (\\n fn.info and any(info.has_derivatives for info in fn.info.values())\\n ):\\n msg = (\\n f\\\"There's a formula for {name}(or its functional variant) in derivatives.yaml. \\\"\\n f\\\"It's required to add a dispatch section for it with explicit supported backends e.g CPU/CUDA \\\"\\n f\\\"or CompositeExplicitAutograd in native_functions.yaml. Please see \\\"\\n f\\\"https://github.com/pytorch/pytorch/tree/master/aten/src/ATen/native#choosing-the-right-dispatch-keyword \\\"\\n f\\\"for instructions to choose the right dispatch keyword.\\\"\\n )\\n assert f.is_abstract, msg\\n\\n return result\\n\\n\\n_foreach_ops_without_differentiability_info = {\\n # No reference backward available due to the lack of `{maximum, minimum}(tensor, scalar)`.\\n (\\\"_foreach_maximum\\\", \\\"Scalar\\\"),\\n (\\\"_foreach_maximum\\\", \\\"ScalarList\\\"),\\n (\\\"_foreach_minimum\\\", \\\"Scalar\\\"),\\n (\\\"_foreach_minimum\\\", \\\"ScalarList\\\"),\\n # No reference backward available as addcdiv/addcmul don't support Tensor as scaling factor.\\n (\\\"_foreach_addcdiv\\\", \\\"Tensor\\\"),\\n (\\\"_foreach_addcmul\\\", \\\"Tensor\\\"),\\n (\\\"_foreach_copy\\\", \\\"\\\"),\\n}\\n\\n_foreach_ops_with_different_arity = {\\n # These ops lack `alpha` of scaling factor to applied to the right hand side argument.\\n (\\\"_foreach_add\\\", \\\"Scalar\\\"),\\n (\\\"_foreach_add\\\", \\\"ScalarList\\\"),\\n (\\\"_foreach_sub\\\", \\\"Scalar\\\"),\\n (\\\"_foreach_sub\\\", \\\"ScalarList\\\"),\\n}\\n\\n\\n@with_native_function_with_differentiability_info_and_key\\ndef emit_body(\\n fn: NativeFunctionWithDifferentiabilityInfo, key: str = \\\"Default\\\"\\n) -> list[str]:\\n assert dispatch_strategy(fn) == \\\"use_derived\\\"\\n f = fn.func\\n info = fn.info[key] if fn.info else None\\n fw_derivatives = fn.fw_derivatives.get(key, []) if fn.fw_derivatives else []\\n\\n name = cpp.name(f.func)\\n inplace = f.func.kind() == SchemaKind.inplace\\n is_out_fn = f.func.kind() == SchemaKind.out\\n returns_void = len(f.func.returns) == 0\\n base_name = get_base_name(f)\\n view_info = get_view_info(f)\\n\\n is_foreach = name.startswith(\\\"_foreach\\\")\\n is_inplace_foreach = is_foreach and inplace\\n if is_inplace_foreach:\\n inplace_foreacharg2refarg: dict[Argument, Argument] = {}\\n refargname2inplace_foreacharg: dict[str, Argument] = {}\\n base_name_and_overload_name = (f.func.name.name.base, f.func.name.overload_name)\\n if info is None:\\n assert (\\n base_name_and_overload_name\\n in _foreach_ops_without_differentiability_info\\n ), f\\\"{'.'.join(base_name_and_overload_name)} should have a differentiability info\\\"\\n else:\\n assert (\\n len(f.func.arguments.flat_non_out)\\n == len(info.func.func.arguments.flat_non_out)\\n ) or (base_name_and_overload_name in _foreach_ops_with_different_arity), (\\n f\\\"{'.'.join(base_name_and_overload_name)} has {len(f.func.arguments.flat_non_out)} args \\\"\\n f\\\"but the reference has {len(info.func.func.arguments.flat_non_out)}\\\"\\n )\\n for foreach_arg, ref_arg in zip(\\n f.func.arguments.flat_non_out, info.func.func.arguments.flat_non_out\\n ):\\n foreach_arg_type = foreach_arg.type\\n if isinstance(foreach_arg_type, ListType):\\n foreach_arg_type = foreach_arg_type.elem\\n assert foreach_arg_type == ref_arg.type\\n inplace_foreacharg2refarg[foreach_arg] = ref_arg\\n refargname2inplace_foreacharg[ref_arg.name] = foreach_arg\\n\\n def gen_differentiable_input(\\n arg: Argument | SelfArgument | TensorOptionsArguments,\\n ) -> DifferentiableInput | None:\\n if isinstance(arg, TensorOptionsArguments):\\n return None\\n a: Argument = arg.argument if isinstance(arg, SelfArgument) else arg\\n\\n # TODO: `cpp_type` is only to keep it byte-for-byte compatible with the old codegen, should remove.\\n # NB: This is not a clone of cpp.argument() - TensorOptionsArguments / faithful / binds are\\n # not handled properly as they are irrelevant for this codegen.\\n cpp_type = cpp.argument_type(a, binds=a.name, symint=True).cpp_type()\\n\\n if not is_differentiable(a.name, a.type, info):\\n return None\\n return DifferentiableInput(\\n name=a.name,\\n type=a.type,\\n cpp_type=cpp_type,\\n )\\n\\n @with_native_function\\n def gen_differentiable_inputs(f: NativeFunction) -> list[DifferentiableInput]:\\n arguments = list(f.func.arguments.non_out)\\n if is_inplace_foreach and info is not None:\\n for i, arg in enumerate(f.func.arguments.flat_non_out):\\n if arg in inplace_foreacharg2refarg:\\n # note(crcrpar): From what I understand, what matters is only the name.\\n # Thus originally I only replace argument only when the names are different.\\n # TODO(crcrpar): Make it simpler.\\n mapped_arg = inplace_foreacharg2refarg[arg]\\n arguments[i] = Argument(\\n mapped_arg.name,\\n mapped_arg.type,\\n mapped_arg.default,\\n mapped_arg.annotation,\\n )\\n return list(mapMaybe(gen_differentiable_input, arguments))\\n\\n def find_args_with_derivatives(\\n differentiable_inputs: list[DifferentiableInput],\\n ) -> list[DifferentiableInput]:\\n \\\"\\\"\\\"Find arguments that have derivative definitions\\\"\\\"\\\"\\n if info is None or not info.has_derivatives:\\n return differentiable_inputs\\n names = {name for d in info.derivatives for name in d.var_names}\\n differentiable = [arg for arg in differentiable_inputs if arg.name in names]\\n if len(differentiable) != len(names):\\n missing = names - {arg.name for arg in differentiable}\\n raise RuntimeError(\\n f\\\"Missing arguments for derivatives: {missing} in {info.name}\\\"\\n )\\n return differentiable\\n\\n differentiable_inputs = gen_differentiable_inputs(f)\\n args_with_derivatives = find_args_with_derivatives(differentiable_inputs)\\n differentiable_outputs = gen_differentiable_outputs(fn, key)\\n\\n undifferentiable = (base_name in DONT_REQUIRE_DERIVATIVE) or (\\n name in DONT_REQUIRE_DERIVATIVE\\n )\\n\\n requires_derivative = (\\n (not undifferentiable)\\n and (len(differentiable_inputs) > 0)\\n and (\\n (len(differentiable_outputs) > 0)\\n # note(crcrpar): In-place foreach functions are a void function.\\n or is_inplace_foreach\\n )\\n )\\n\\n if (\\n info is not None\\n and info.has_derivatives\\n and not requires_derivative\\n # out= ops are allowed to have zero returns which cause requires_derivative to be False\\n # we shouldn't error out though (out= ops for autograd just redispatch)\\n and len(f.func.returns) > 0\\n ):\\n raise RuntimeError(\\n f\\\"ERROR: derivative ignored for {name} -- specified an autograd function without derivative\\\"\\n )\\n\\n # note(crcrpar): In-place foreach functions do not support forward AD\\n if requires_derivative and len(fw_derivatives) > 0 and not is_inplace_foreach:\\n assert sum(len(derivative.var_names) for derivative in fw_derivatives) == len(\\n differentiable_outputs\\n ), (\\n \\\"Expected the number of forward derivatives implemented to match the \\\"\\n \\\"number of differentiable outputs. NB: This only applies when at least \\\"\\n \\\"one forward derivative is implemented. Not implementing any forward \\\"\\n \\\"derivatives is also okay, and we would require inputs to the op to \\\"\\n \\\"not have associated tangents in that case.\\\"\\n )\\n\\n try_jit_decomposition = (\\n requires_derivative\\n and len(fw_derivatives) == 0\\n and (not modifies_arguments(f))\\n and (not returns_void)\\n )\\n\\n def emit_save_inputs() -> list[str]:\\n setup: list[str] = []\\n if info is None or not info.has_derivatives:\\n return setup\\n\\n has_tensorlist_arg = any(\\n is_tensor_list_type(arg.type) for arg in args_with_derivatives\\n )\\n\\n # We don't want to save tensors if we know that they will never be used\\n # when computing the derivative, so we add guards to those statements\\n def guard_for(arg: SavedAttribute) -> str | None:\\n assert info is not None\\n\\n # It's hard to determine the edge offset if we have TensorLists\\n # NOTE(crcrpar): in-place foreach functions' arguments include tensorlist\\n # but their derivatives don't use it, so let them bypass this check.\\n if has_tensorlist_arg and (not is_inplace_foreach):\\n return None\\n\\n # Empirical evaluation of the cases where we insert those guards in\\n # backward show that they are somewhat useless. E.g. there's no need\\n # to guard on some values captured from forward, because they had to\\n # require_grad if the backward function even gets executed. I don't\\n # have any good ideas for detecting those cases, so I simply disabled the\\n # checks.\\n if \\\"backward\\\" in info.name:\\n return None\\n\\n # If there's a single derivative we could compute, we already have\\n # a requires_grad check that is sufficient\\n if len(args_with_derivatives) <= 1:\\n return None\\n\\n # We really only care about trimming down the amount of tensors we save\\n if arg.nctype.type != BaseCType(tensorT):\\n return None\\n\\n # We want to emit simple guards, so we only allow that if checking one\\n # input is enough to determine whether we need that value\\n used_in = [d for d in info.derivatives if arg in d.saved_inputs]\\n assert len(used_in) > 0\\n if len(used_in) != 1:\\n return None\\n derivative = used_in[0]\\n\\n # Case with multioutput formulas\\n # TODO: process all derivative formulas!!!\\n if len(derivative.var_names) != 1:\\n wrap_opt_if_start = derivative.formula.find(\\n f\\\"wrap_opt_if({arg.nctype.name}\\\"\\n )\\n if wrap_opt_if_start == -1:\\n return None\\n\\n wrap_opt_if_match = re.match(\\n rf\\\"wrap_opt_if\\\\({arg.nctype.name},(.*?)\\\\)\\\",\\n derivative.formula[wrap_opt_if_start:],\\n )\\n assert wrap_opt_if_match is not None\\n\\n # Condition is between 'wrap_opt_if(var_name,' and ')'.\\n condition_slice = slice(len(rf\\\"wrap_opt_if\\\\({arg.nctype.name},\\\"), -1)\\n wrap_opt_if_condition = wrap_opt_if_match.group(0)[\\n condition_slice\\n ].strip()\\n # replace 'grad_input_mask[num]' with 'grad_fn->should_compute_output(num)'\\n wrap_opt_if_condition = re.sub(\\n r\\\"grad_input_mask\\\\[(\\\\d+)\\\\]\\\",\\n r\\\"grad_fn->should_compute_output(\\\\1)\\\",\\n wrap_opt_if_condition,\\n )\\n return f\\\"{wrap_opt_if_condition}\\\"\\n\\n # Figure out the offset of the edge that uses this variable\\n derivative_var_name = derivative.var_names[0]\\n for edge_off, a in enumerate(args_with_derivatives):\\n if a.name == derivative_var_name:\\n break\\n else:\\n raise AssertionError\\n return f\\\"grad_fn->should_compute_output({edge_off})\\\"\\n\\n if is_inplace_foreach:\\n save_input_stmts = save_variables(info.all_saved_inputs, False, guard_for)\\n if save_input_stmts:\\n setup.append(\\n LOOP_OVER_VECTOR_OF_GRAD_FNS.substitute(\\n preamble=\\\"\\\", statements=save_input_stmts\\n )\\n )\\n else:\\n setup.extend(save_variables(info.all_saved_inputs, False, guard_for))\\n for arg in args_with_derivatives:\\n if is_tensor_list_type(arg.type):\\n setup.append(f\\\"grad_fn->{arg.name}_size_ = {arg.name}.size();\\\")\\n return setup\\n\\n def setup_derivative(differentiable_inputs: list[DifferentiableInput]) -> list[str]:\\n body: list[str] = []\\n if is_out_fn:\\n # For out functions, ensure that no input or output requires grad\\n body.append(DECLARE_GRAD_FN.substitute(op=\\\"Node\\\"))\\n body.append(\\n SETUP_NONE_REQUIRES_GRAD.substitute(\\n base_name=base_name,\\n args_to_check=[arg.name for arg in differentiable_inputs],\\n )\\n )\\n body.append(\\n SETUP_NONE_REQUIRES_GRAD.substitute(\\n base_name=base_name,\\n args_to_check=[arg.name for arg in differentiable_outputs],\\n )\\n )\\n return body\\n\\n op = info.op if info is not None and info.has_derivatives else \\\"NotImplemented\\\"\\n setup = []\\n if not is_inplace_foreach:\\n setup.extend(\\n ASSIGN_GRAD_FN.substitute(\\n op=op,\\n op_ctor=\\\"\\\"\\n if info is not None and info.has_derivatives\\n else f'\\\"{cpp.name(f.func)}\\\"',\\n args_with_derivatives=[arg.name for arg in args_with_derivatives],\\n ).split(\\\"\\\\n\\\")\\n )\\n else:\\n # note(crcrpar): Assuming in-place foreach function's self_arg is always TensorList.\\n list_like_arg = \\\"self\\\"\\n args = [arg.name for arg in args_with_derivatives]\\n for i, arg in enumerate(args):\\n if is_inplace_foreach and info is not None:\\n if arg in refargname2inplace_foreacharg:\\n foreach_arg = refargname2inplace_foreacharg[arg]\\n args[i] = foreach_arg.name + (\\n \\\"[i]\\\" if isinstance(foreach_arg.type, ListType) else \\\"\\\"\\n )\\n else:\\n if arg == list_like_arg:\\n args[i] = arg + \\\"[i]\\\"\\n setup.extend(\\n ASSIGN_VECTOR_OF_GRAD_FN.substitute(\\n op=op,\\n op_ctor=\\\"\\\"\\n if info is not None and info.has_derivatives\\n else f'\\\"{cpp.name(f.func)}\\\"',\\n args_with_derivatives=args,\\n irange=f\\\"{list_like_arg}.size()\\\",\\n ).split(\\\"\\\\n\\\")\\n )\\n setup.extend(emit_save_inputs())\\n\\n body.extend(\\n emit_check_no_requires_grad(differentiable_inputs, args_with_derivatives)\\n )\\n declare_grad_fn_template = (\\n DECLARE_GRAD_FN if not is_inplace_foreach else DECLARE_VECTOR_OF_GRAD_FN\\n )\\n body.append(declare_grad_fn_template.substitute(op=op))\\n body.append(SETUP_DERIVATIVE.substitute(setup=setup))\\n return body\\n\\n def emit_check_if_in_complex_autograd_allowlist() -> list[str]:\\n body: list[str] = []\\n if base_name in GRADIENT_IMPLEMENTED_FOR_COMPLEX:\\n return body\\n for arg in differentiable_outputs:\\n name = arg.name\\n # TODO: should be `arg.type.is_tensor_like()`?\\n if arg.cpp_type == \\\"at::Tensor\\\" or arg.cpp_type in TENSOR_LIST_LIKE_CTYPES:\\n body.append(f'throw_error_for_complex_autograd({name}, \\\"{base_name}\\\");')\\n return body\\n\\n def emit_check_no_requires_grad(\\n tensor_args: list[DifferentiableInput],\\n args_with_derivatives: list[DifferentiableInput],\\n ) -> list[str]:\\n \\\"\\\"\\\"Checks that arguments without derivatives don't require grad\\\"\\\"\\\"\\n body: list[str] = []\\n for arg in tensor_args:\\n if arg in args_with_derivatives:\\n continue\\n arg_name = arg.name\\n if info and arg_name in info.non_differentiable_arg_names:\\n continue\\n if arg_name == \\\"output\\\":\\n # Double-backwards definitions sometimes take in 'input' and\\n # 'output', but only define the derivative for input.\\n continue\\n body.append(f'check_no_requires_grad({arg_name}, \\\"{arg_name}\\\", \\\"{name}\\\");')\\n return body\\n\\n def emit_original_self_definition() -> list[str]:\\n body: list[str] = []\\n if inplace:\\n if is_inplace_foreach:\\n body.append(\\n \\\"std::vector<::std::optional> original_selfs(self.size());\\\"\\n )\\n else:\\n body.append(\\\"::std::optional original_self;\\\")\\n\\n all_forward_grad_cond = []\\n for derivative in fw_derivatives:\\n if derivative.required_original_self_value:\\n all_forward_grad_cond.append(\\n get_any_has_forward_grad_name(derivative.var_names)\\n )\\n\\n if all_forward_grad_cond:\\n if not is_inplace_foreach:\\n body.append(f'if ({\\\" || \\\".join(all_forward_grad_cond)}) {{')\\n body.append(\\\" original_self = self.clone();\\\")\\n body.append(\\\"}\\\")\\n else:\\n current_all_forward_grad_cond = [\\n f\\\"{cond}[i]\\\" for cond in all_forward_grad_cond\\n ]\\n body.append(\\\"for (const auto& i : c10::irange(self.size())) {\\\")\\n body.append(\\n f\\\" if ({' || '.join(current_all_forward_grad_cond)}) {{\\\"\\n )\\n body.append(\\\" original_selfs[i] = self[i].clone();\\\")\\n body.append(\\\" }\\\")\\n body.append(\\\"}\\\")\\n\\n return body\\n\\n def save_variables(\\n saved_variables: Sequence[SavedAttribute],\\n is_output: bool,\\n guard_for: Callable[[SavedAttribute], str | None] = lambda name: None,\\n ) -> Sequence[str]:\\n # assign the saved variables to the generated grad_fn\\n stmts: list[str] = []\\n for arg in sorted(saved_variables, key=lambda sa: str(sa.nctype.name)):\\n name = (\\n arg.nctype.name.name\\n if isinstance(arg.nctype.name, SpecialArgName)\\n else arg.nctype.name\\n )\\n foreacharg: Argument | None = None\\n is_foreacharg_list_type: bool = False\\n type = arg.nctype.type\\n expr = arg.expr\\n stmts_prepend = None\\n if is_inplace_foreach and info is not None:\\n # todo(crcrpar): See if we can add some check e.g. `assert foreacharg is not None`.\\n # for now the example assert would fail.\\n name_to_query = name.split(\\\"_scalar_type\\\")[0]\\n if name_to_query in refargname2inplace_foreacharg:\\n foreacharg = refargname2inplace_foreacharg[name_to_query]\\n is_foreacharg_list_type = isinstance(foreacharg.type, ListType)\\n if foreacharg is not None:\\n name_in_expr = (\\n f\\\"{foreacharg.name}{'[i]' if is_foreacharg_list_type else ''}\\\"\\n )\\n src_name = name\\n if \\\"_scalar_type\\\" in src_name:\\n split_src_name = src_name.split(\\\"_scalar_type\\\")\\n assert len(split_src_name) == 2\\n src_name = split_src_name[0]\\n expr = expr.replace(src_name, name_in_expr)\\n if (\\n type == BaseCType(tensorT)\\n or type == OptionalCType(BaseCType(tensorT))\\n or type == MutRefCType(OptionalCType(BaseCType(tensorT)))\\n or (is_output and type == BaseCType(scalarT))\\n ):\\n # note(crcrpar): Here `expr` is generated from scratch, `arg.expr` is ignored.\\n var = name\\n name += \\\"_\\\"\\n if var == \\\"self\\\" and inplace:\\n original_self_var = (\\n \\\"original_self\\\"\\n if not is_inplace_foreach\\n else \\\"original_selfs[i]\\\"\\n )\\n self_var = var if not is_inplace_foreach else var + \\\"[i]\\\"\\n stmts_prepend = f\\\"if (!{original_self_var}.has_value()) {original_self_var} = {self_var}.clone()\\\"\\n var = f\\\"{original_self_var}.value()\\\"\\n assert not is_output\\n if inplace and is_output:\\n assert name == \\\"result_\\\"\\n var = (\\n \\\"self[i]\\\"\\n if is_inplace_foreach or is_foreacharg_list_type\\n else \\\"self\\\"\\n )\\n is_inplace_view = f\\\"{var}.is_view()\\\"\\n expr = f\\\"SavedVariable({var}, {str(is_output).lower()}, {is_inplace_view})\\\"\\n else:\\n expr = f\\\"SavedVariable({var}, {str(is_output).lower()})\\\"\\n if foreacharg is not None and \\\"original_selfs\\\" not in expr:\\n expr = expr.replace(src_name, name_in_expr)\\n elif (\\n type == BaseCType(tensorListT)\\n or type == ListCType(OptionalCType(BaseCType(tensorT)))\\n or type == BaseCType(iTensorListRefT)\\n or type == VectorCType(BaseCType(tensorT))\\n ):\\n # See Note [nuanced return type of out-of-place foreach functions]\\n if type == VectorCType(BaseCType(tensorT)):\\n assert is_foreach and is_output\\n expr = f\\\"make_saved_variable_list({name}, {str(is_foreach and is_output).lower()})\\\"\\n name += \\\"_\\\"\\n elif type == BaseCType(intArrayRefT):\\n expr = expr + \\\".vec()\\\"\\n elif type == BaseCType(symIntArrayRefT):\\n expr = expr + \\\".vec()\\\"\\n elif type == BaseCType(stringT):\\n expr = f\\\"std::string({expr})\\\"\\n elif type == OptionalCType(BaseCType(stringT)):\\n expr = f\\\"{expr}.has_value() ? ::std::optional(std::string({expr}.value())) : ::std::nullopt\\\"\\n elif type == ArrayRefCType(\\n elem=BaseCType(type=BaseCppType(ns=\\\"at\\\", name=\\\"Scalar\\\"))\\n ):\\n expr = expr + \\\".vec()\\\"\\n\\n guard = guard_for(arg)\\n if guard is None:\\n if stmts_prepend:\\n stmts.append(f\\\"{stmts_prepend};\\\")\\n stmts.append(f\\\"grad_fn->{name} = {expr};\\\")\\n else:\\n stmts.append(f\\\"if ({guard}) {{\\\")\\n if stmts_prepend:\\n stmts.append(f\\\" {stmts_prepend};\\\")\\n stmts.append(f\\\" grad_fn->{name} = {expr};\\\")\\n stmts.append(\\\"}\\\")\\n return stmts\\n\\n # Generates a Dispatcher::redispatch() call into the dispatcher. We do this mainly for performance reasons:\\n # - Pre-compute the full DispatchKeySet. This saves the dispatcher from having to read from TLS.\\n # - redispatch() avoids a redundant call to RecordFunction, which was already called right before\\n # we entered this autograd kernel.\\n def emit_dispatch_call(\\n f: NativeFunction, input_base: str, unpacked_args: Sequence[str]\\n ) -> str:\\n \\\"\\\"\\\"Dispatch call via function in a namespace or method on Tensor.\\\"\\\"\\\"\\n dispatcher_sig = DispatcherSignature.from_schema(f.func)\\n dispatcher_exprs = dispatcher_sig.exprs()\\n\\n # code-generated autograd kernels plumb and recompute dispatch keys directly through the kernel for performance.\\n # Ops also always have a function variant of the redispatch API.\\n # See Note [Plumbing Keys Through The Dispatcher] for details.\\n dispatch_key_set = \\\"ks & c10::after_autograd_keyset\\\"\\n call = CALL_REDISPATCH.substitute(\\n api_name=cpp.name(\\n f.func,\\n faithful_name_for_out_overloads=True,\\n symint_overload=f.func.has_symint(),\\n ),\\n unpacked_args=[dispatch_key_set] + list(unpacked_args),\\n )\\n return call\\n\\n def wrap_output(\\n f: NativeFunction, unpacked_bindings: list[Binding], var: str\\n ) -> str:\\n call = \\\"\\\"\\n rhs_value: str | None = None\\n if not any(r.type.is_tensor_like() for r in f.func.returns):\\n rhs_value = var\\n else:\\n rhs_value = f\\\"std::move({var})\\\"\\n assert rhs_value is not None\\n call += ASSIGN_RETURN_VALUE.substitute(\\n return_values=tie_return_values(f), rhs_value=rhs_value\\n )\\n return call\\n\\n def check_tensorimpl_and_storage(\\n call: str, unpacked_bindings: list[Binding]\\n ) -> str:\\n # See NOTE [ TensorImpl and Storage Pointer Sanity Checks ]\\n stmts_before_call: list[str] = []\\n stmts_after_call: list[str] = []\\n\\n if cpp.name(f.func) in DONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE:\\n return call\\n\\n # Check properties of inputs (enforce (1))\\n for unpacked_binding in unpacked_bindings:\\n arg = unpacked_binding.name\\n noref_cpp_type = unpacked_binding.nctype.type.remove_const_ref()\\n if noref_cpp_type == BaseCType(tensorListT) or noref_cpp_type == BaseCType(\\n iTensorListRefT\\n ):\\n stmts_before_call += [\\n SAVE_TENSORLIST_STORAGE.substitute(tensorlist_name=arg),\\n SAVE_TENSORLIST_IMPL.substitute(tensorlist_name=arg),\\n ]\\n stmts_after_call += [\\n ENFORCE_SAME_TENSORLIST_STORAGE.substitute(tensorlist_name=arg),\\n ENFORCE_SAME_TENSORLIST_IMPL.substitute(tensorlist_name=arg),\\n ]\\n elif noref_cpp_type == ListCType(OptionalCType(BaseCType(tensorT))):\\n stmts_before_call += [\\n SAVE_OPTIONALTENSORLIST_STORAGE.substitute(tensorlist_name=arg),\\n SAVE_OPTIONALTENSORLIST_IMPL.substitute(tensorlist_name=arg),\\n ]\\n stmts_after_call += [\\n ENFORCE_SAME_OPTIONALTENSORLIST_STORAGE.substitute(\\n tensorlist_name=arg\\n ),\\n ENFORCE_SAME_OPTIONALTENSORLIST_IMPL.substitute(\\n tensorlist_name=arg\\n ),\\n ]\\n elif noref_cpp_type == BaseCType(tensorT):\\n stmts_before_call += [\\n SAVE_TENSOR_STORAGE.substitute(tensor_name=arg),\\n SAVE_TENSOR_IMPL.substitute(tensor_name=arg),\\n ]\\n stmts_after_call += [\\n ENFORCE_SAME_TENSOR_STORAGE.substitute(\\n tensor_name=arg, out_tensor_name=arg\\n ),\\n ENFORCE_SAME_TENSOR_IMPL.substitute(tensor_name=arg),\\n ]\\n\\n assert (stmts_before_call and stmts_after_call) or (\\n not stmts_before_call and not stmts_after_call\\n )\\n\\n # Check properties of outputs (enforce (2), (3))\\n if f.func.kind() not in (SchemaKind.inplace, SchemaKind.out):\\n base_name = f.func.name.name.base # TODO: should be str(f.func.name.name)?\\n aliased_arg_name = ALL_VIEW_FUNCTIONS.get(base_name, None)\\n if aliased_arg_name is not None:\\n aliased_arg_name = unpacked_name(aliased_arg_name)\\n for i, (ret, ret_name) in enumerate(\\n zip(f.func.returns, cpp.return_names(f))\\n ):\\n noref_cpp_type = cpp.return_type(ret, symint=True).remove_const_ref()\\n if noref_cpp_type == BaseCType(tensorT):\\n if aliased_arg_name is not None:\\n assert (\\n i == 0\\n ), \\\"Expect non-CompositeImplicitAutograd view function {base} to return single output\\\"\\n stmts_after_call += [\\n ENFORCE_SAME_TENSOR_STORAGE.substitute(\\n tensor_name=aliased_arg_name, out_tensor_name=ret_name\\n )\\n ]\\n else:\\n if (\\n type_wrapper_name(f)\\n not in DONT_ENFORCE_STORAGE_IMPL_USE_COUNT\\n ):\\n stmts_after_call += [\\n ENFORCE_TENSOR_STORAGE_USE_COUNT_EQUALS_ONE.substitute(\\n tensor_name=ret_name, fn_name=type_wrapper_name(f)\\n )\\n ]\\n\\n if type_wrapper_name(f) not in DONT_ENFORCE_TENSOR_IMPL_USE_COUNT:\\n stmts_after_call += [\\n ENFORCE_TENSOR_IMPL_USE_COUNT_LT_OR_EQ_ONE.substitute(\\n tensor_name=ret_name, fn_name=type_wrapper_name(f)\\n )\\n ]\\n\\n # Currently we don't have any functions that return the following types, but\\n # we should update the checks once we do\\n elif noref_cpp_type == ListCType(OptionalCType(BaseCType(tensorT))):\\n raise AssertionError(\\n f\\\"Please add use_count checks for {noref_cpp_type}\\\"\\n )\\n elif noref_cpp_type == BaseCType(tensorListT):\\n raise AssertionError(\\n f\\\"Please add use_count checks for {noref_cpp_type}\\\"\\n )\\n\\n if stmts_before_call and stmts_after_call:\\n call = (\\n RUN_ONLY_IN_DEBUG_MODE.substitute(statements=stmts_before_call)\\n + call\\n + RUN_ONLY_IN_DEBUG_MODE.substitute(statements=stmts_after_call)\\n )\\n return call\\n\\n def emit_call(\\n f: NativeFunction, unpacked_bindings: list[Binding], try_jit_decomposition: bool\\n ) -> str:\\n # We only care about adding `at::AutoDispatchBelowAutograd` guard for non-variable dispatch\\n # (which corresponds to 'use_derived' strategy). The purpose of this guard is to make sure\\n # the baseType operations still dispatch to non-Variable type, even if the arguments passed\\n # in are now Variables.\\n # See NOTE [ Treating Variables as non-Variables in type dispatch ] for details.\\n unpacked_args = [b.name for b in unpacked_bindings]\\n base_type_call = emit_dispatch_call(f, \\\"self_\\\", unpacked_args)\\n\\n if get_view_info(f) is not None or modifies_arguments(f):\\n guard = \\\"at::AutoDispatchBelowAutograd guard;\\\"\\n else:\\n guard = \\\"at::AutoDispatchBelowADInplaceOrView guard;\\\"\\n\\n any_has_forward_grad = (\\n get_any_has_fw_grad_cond(derivative=None)\\n if requires_derivative\\n else \\\"false\\\"\\n )\\n return_types = \\\", \\\".join(\\n [cpp.return_type(a, symint=True).cpp_type() for a in f.func.returns]\\n )\\n if len(f.func.returns) > 1:\\n return_types = f\\\"std::tuple<{return_types}>\\\"\\n\\n arg_names = [\\n a.name\\n for a in cpp.arguments(\\n f.func.arguments,\\n faithful=True,\\n symint=True,\\n method=False,\\n cpp_no_default_args=set(),\\n )\\n ]\\n\\n if not modifies_arguments(f) and not returns_void:\\n if try_jit_decomposition:\\n call = DISPATCH_TO_NON_VAR_TYPE_WITH_TMP_RETURN_VALUES_JVP_DECOMP.substitute(\\n base_type_call=base_type_call,\\n tmp_var=TMP_VAR,\\n guard=guard,\\n any_has_forward_grad=any_has_forward_grad,\\n op_name=cpp.name(f.func),\\n op_overload=f.func.name.overload_name,\\n return_types=return_types,\\n arg_names=arg_names,\\n )\\n else:\\n call = DISPATCH_TO_NON_VAR_TYPE_WITH_TMP_RETURN_VALUES.substitute(\\n base_type_call=base_type_call,\\n tmp_var=TMP_VAR,\\n guard=guard,\\n )\\n\\n call += wrap_output(f, unpacked_bindings, TMP_VAR)\\n else:\\n assert not try_jit_decomposition\\n call = DISPATCH_TO_NON_VAR_TYPE_WITHOUT_RETURN_VALUES.substitute(\\n base_type_call=base_type_call, guard=guard\\n )\\n call = check_tensorimpl_and_storage(call, unpacked_bindings)\\n return call\\n\\n def emit_history() -> str:\\n fn = \\\"rebase\\\" if modifies_arguments(f) and view_info is None else \\\"set\\\"\\n output_names = [r.name for r in differentiable_outputs]\\n # TODO: flatten allocates a std::vector, which could be expensive\\n outs = CodeTemplate(\\\"flatten_tensor_args( ${outs} )\\\").substitute(\\n outs=output_names if not is_inplace_foreach else \\\"self\\\"\\n )\\n if not is_inplace_foreach:\\n return SET_HISTORY.substitute(fn=fn, differentiable_outputs=outs)\\n else:\\n return LOOP_OVER_VECTOR_OF_GRAD_FNS.substitute(\\n preamble=(\\n f\\\"auto differentiable_outputs = {outs};\\\\n\\\"\\n f\\\"TORCH_INTERNAL_ASSERT(differentiable_outputs.size() == grad_fns.size());\\\"\\n ),\\n statements=f\\\"{fn}_history(differentiable_outputs[i], grad_fns[i]);\\\",\\n )\\n\\n def emit_save_outputs() -> str:\\n if is_out_fn:\\n # out functions don't currently support differentiation\\n return \\\"\\\"\\n if info is not None and info.has_derivatives:\\n stmts = save_variables(info.all_saved_outputs, True)\\n if len(stmts) == 0:\\n return \\\"\\\"\\n if not is_inplace_foreach:\\n return CONDITIONAL.substitute(cond=\\\"grad_fn\\\", statements=stmts)\\n else:\\n return LOOP_OVER_VECTOR_OF_GRAD_FNS.substitute(\\n preamble=\\\"\\\", statements=stmts\\n )\\n return \\\"\\\"\\n\\n def emit_any_requires_grad() -> list[str]:\\n extra_condition = \\\"\\\"\\n if info and info.output_differentiability_conditions:\\n assert len(info.output_differentiability_conditions) == 1\\n extra_condition = f\\\"_any_requires_grad &= ({info.output_differentiability_conditions[0]});\\\"\\n names_of_args_with_derivatives = [arg.name for arg in args_with_derivatives]\\n if is_inplace_foreach and info is not None:\\n for i, arg in enumerate(names_of_args_with_derivatives):\\n for f_arg, r_arg in inplace_foreacharg2refarg.items():\\n if arg == r_arg.name:\\n names_of_args_with_derivatives[i] = f_arg.name\\n return [\\n SETUP_ANY_REQUIRES_GRAD.substitute(\\n args_with_derivatives=names_of_args_with_derivatives,\\n extra_differentiability_conditions=extra_condition,\\n )\\n ]\\n\\n def get_any_has_forward_grad_name(var_names: tuple[str, ...]) -> str:\\n if len(var_names) == 1:\\n return f\\\"_any_has_forward_grad_{var_names[0]}\\\"\\n else:\\n return f'_any_has_forward_grad_{\\\"_\\\".join(var_names)}'\\n\\n def emit_any_has_forward_grad() -> list[str]:\\n content: list[str] = []\\n if not is_foreach:\\n for derivative in fw_derivatives:\\n requires_fw_grad = get_any_has_fw_grad_cond(derivative=derivative)\\n if info and info.output_differentiability_conditions:\\n assert len(info.output_differentiability_conditions) == 1\\n requires_fw_grad = f\\\"({info.output_differentiability_conditions[0]}) && {requires_fw_grad}\\\"\\n content.append(\\n f\\\"[[maybe_unused]] auto {get_any_has_forward_grad_name(derivative.var_names)} = {requires_fw_grad};\\\"\\n )\\n else:\\n for derivative in fw_derivatives:\\n bool_vector_name = get_any_has_forward_grad_name(derivative.var_names)\\n cur_derivative_conditions = []\\n for inp in differentiable_inputs:\\n if derivative.required_inputs_fw_grad is None:\\n continue\\n if inp.name not in derivative.required_inputs_fw_grad:\\n continue\\n inp_name = (\\n inp.name\\n if not inplace\\n else refargname2inplace_foreacharg[inp.name].name\\n )\\n inp_type = (\\n inp.type\\n if not inplace\\n else refargname2inplace_foreacharg[inp.name].type\\n )\\n is_list_type = is_tensor_list_type(inp_type)\\n if is_list_type:\\n if inp_name != \\\"self\\\":\\n content.append(\\n FW_DERIVATIVE_SIZE_CHECK_TEMPLATE.substitute(\\n inp_name=inp_name\\n )\\n )\\n cur_derivative_conditions.append(\\n FW_DERIVATIVE_CHECK_TEMPLATE.substitute(\\n req_inp=inp_name + \\\"[i]\\\"\\n )\\n )\\n else:\\n cur_derivative_conditions.append(\\n FW_DERIVATIVE_CHECK_TEMPLATE.substitute(req_inp=inp_name)\\n )\\n\\n content.append(f\\\"std::vector {bool_vector_name}(self.size());\\\")\\n content.append(\\\"for (const auto& i : c10::irange(self.size())) {\\\")\\n content.append(\\n f\\\" {bool_vector_name}[i] = {' || '.join(cur_derivative_conditions)};\\\"\\n )\\n content.append(\\\"}\\\")\\n return content\\n\\n def emit_check_inplace() -> list[str]:\\n if not inplace:\\n return []\\n return [\\n f\\\"check_inplace({arg.name}, _any_requires_grad);\\\"\\n for arg in differentiable_outputs\\n ]\\n\\n def emit_fw_derivatives() -> list[str]:\\n content: list[str] = []\\n fw_grad_setters: list[str] = []\\n for derivative in fw_derivatives:\\n res = derivative.var_names\\n if f.func.name.name.inplace:\\n assert (\\n len(res) == 1\\n ), \\\"Expected number of outputs to be 1 if function is inplace\\\"\\n # TODO update this when inplace namings are unified\\n res = (\\\"self\\\",)\\n\\n assert derivative.required_inputs_fw_grad is not None\\n\\n unpacked_arguments = \\\"\\\"\\n for inp in differentiable_inputs:\\n inp_name = inp.name\\n is_input_tensorlist = is_foreach and is_tensor_list_type(\\n inp.type\\n if not inplace\\n else refargname2inplace_foreacharg[inp.name].type\\n )\\n input_suffix = \\\"[i]\\\" if is_input_tensorlist else \\\"\\\"\\n if is_inplace_foreach:\\n if inp.name in refargname2inplace_foreacharg:\\n inp_name = refargname2inplace_foreacharg[inp.name].name\\n zeros_fn = (\\n \\\"zeros_symint\\\"\\n if inplace and inp.name == \\\"self\\\"\\n else \\\"_efficientzerotensor_symint\\\"\\n )\\n if inp.name in derivative.required_inputs_fw_grad:\\n unpacked_arguments += (\\n FW_DERIVATIVE_DEFINED_GRAD_TEMPLATE.substitute(\\n inp_name=inp.name,\\n inp=inp_name + input_suffix,\\n zeros_fn=zeros_fn,\\n )\\n )\\n if inp.name in (derivative.required_inputs_primal or []):\\n unpacked_arguments += (\\n FW_DERIVATIVE_DEFINED_PRIMAL_TEMPLATE.substitute(\\n inp_name=inp.name,\\n inp=inp_name + input_suffix,\\n )\\n )\\n if derivative.required_original_self_value:\\n input_suffix = \\\"s[i]\\\" if is_inplace_foreach else \\\"\\\"\\n unpacked_arguments += FW_DERIVATIVE_DEFINED_GRAD_TEMPLATE.substitute(\\n inp_name=\\\"original_self\\\",\\n inp=\\\"original_self\\\" + input_suffix,\\n zeros_fn=zeros_fn,\\n )\\n unpacked_arguments += FW_DERIVATIVE_DEFINED_PRIMAL_TEMPLATE.substitute(\\n inp_name=\\\"original_self\\\",\\n inp=\\\"original_self\\\" + input_suffix,\\n )\\n elif inplace and derivative.is_reusing_outplace_formula:\\n # The gradient wasn't already cloned, do it if grad mode is enabled\\n unpacked_arguments += (\\n \\\"self_t = GradMode::is_enabled() ? self_t.clone() : self_t;\\\"\\n )\\n\\n if inplace:\\n is_inplace_str = \\\"true\\\"\\n else:\\n is_inplace_str = \\\"false\\\"\\n\\n requires_fw_grad = get_any_has_forward_grad_name(derivative.var_names)\\n\\n if all(\\n (isinstance(var_type, BaseType) and var_type.is_tensor_like())\\n for var_type in derivative.var_types\\n ):\\n # Is there a way to get from BaseType to BaseCType\\n if len(derivative.var_types) == 1:\\n opt_res_grad_type = OptionalCType(BaseCType(tensorT)).cpp_type()\\n if not is_foreach:\\n fw_grad_setters.append(\\n FW_DERIVATIVE_SETTER_TENSOR.substitute(\\n out_arg=res[0], is_inplace=is_inplace_str\\n )\\n )\\n else:\\n assert res[0] == (\\\"result\\\" if not inplace else \\\"self\\\")\\n fw_grad_setters.append(\\n FW_DERIVATIVE_SETTER_TENSOR_FOREACH.substitute(\\n out_arg=res[0], is_inplace=is_inplace_str\\n )\\n )\\n requires_fw_grad += f\\\" && ({derivative.var_names[0]}.defined())\\\"\\n else:\\n tuple_type = TupleCType(\\n [BaseCType(tensorT)] * len(derivative.var_types)\\n )\\n opt_res_grad_type = OptionalCType(tuple_type).cpp_type()\\n for idx, single_res in enumerate(res):\\n fw_grad_setters.append(\\n FW_DERIVATIVE_SETTER_MULTI_OUTPUT.substitute(\\n idx=idx, all_res=\\\"_\\\".join(res), out_arg=single_res\\n )\\n )\\n elif (\\n isinstance(derivative.var_types[0], ListType)\\n and derivative.var_types[0].is_tensor_like()\\n ):\\n assert (\\n len(derivative.var_types) == 1\\n ), \\\"Expected number of outputs to be 1 if function returns ListType\\\"\\n if not is_foreach:\\n opt_res_grad_type = OptionalCType(\\n VectorCType(BaseCType(tensorT))\\n ).cpp_type()\\n fw_grad_setters.append(\\n FW_DERIVATIVE_SETTER_TENSOR_LIST.substitute(\\n out_arg=res[0], is_inplace=is_inplace_str\\n )\\n )\\n else:\\n # TODO(crcrpar): Should this (= the foreach specific logic) be refactored somehow?\\n # Only out-place foreach functions that have entries in `tools/autograd/derivatives.yaml`\\n # can reach here.\\n opt_res_grad_type = OptionalCType(BaseCType(tensorT)).cpp_type()\\n fw_grad_setters.append(\\n FW_DERIVATIVE_SETTER_TENSOR_FOREACH.substitute(\\n out_arg=res[0], is_inplace=is_inplace_str\\n )\\n )\\n else:\\n raise RuntimeError(\\\"Unsupported output type for forward derivative\\\")\\n\\n if not is_foreach:\\n fw_grad_opt_definition = f\\\"{opt_res_grad_type} {'_'.join(res)}_new_fw_grad_opt = ::std::nullopt;\\\"\\n # View ops create fw_grad that already is a view of the base's fw_grad so just use that\\n content.append(\\n FW_DERIVATIVE_TEMPLATE.substitute(\\n fw_grad_opt_definition=fw_grad_opt_definition,\\n requires_fw_grad=requires_fw_grad,\\n formula=derivative.formula,\\n out_arg=\\\"_\\\".join(res),\\n unpacked_arguments=unpacked_arguments,\\n )\\n )\\n else:\\n # note(crcrpar): Assuming `self` is TensorList.\\n fw_grad_opt_definition = (\\n f\\\"std::vector<{opt_res_grad_type}> {'_'.join(res)}_new_fw_grad_opts\\\"\\n \\\"(self.size(), ::std::nullopt);\\\"\\n )\\n foreach_forward_grad_formula = derivative.formula\\n _foreach_arg: Argument | DifferentiableInput\\n if inplace:\\n for _foreach_arg, _ref_arg in inplace_foreacharg2refarg.items():\\n # note(crcrpar): Massage only Scalar and ArrayRef here.\\n if not (\\n is_tensor_type(_foreach_arg.type)\\n or is_tensor_list_type(_foreach_arg.type)\\n ):\\n pattern = _foreach_arg.name\\n if isinstance(_foreach_arg.type, ListType):\\n pattern += \\\"[i]\\\"\\n foreach_forward_grad_formula = (\\n foreach_forward_grad_formula.replace(\\n _ref_arg.name, pattern\\n )\\n )\\n else:\\n if (\\n \\\"result\\\" in foreach_forward_grad_formula\\n and \\\"result[i]\\\" not in foreach_forward_grad_formula\\n ):\\n foreach_forward_grad_formula = (\\n foreach_forward_grad_formula.replace(\\\"result\\\", \\\"result[i]\\\")\\n )\\n\\n content.append(\\n FW_DERIVATIVE_FOREACH_TEMPLATE.substitute(\\n fw_grad_opt_definition=fw_grad_opt_definition,\\n vector_of_optional_tensor=f\\\"{'_'.join(res)}_new_fw_grad_opts\\\",\\n any_has_forward_grad_for_current_index=\\\" || \\\".join(\\n get_any_has_forward_grad_name(derivative.var_names) + \\\"[i]\\\"\\n for derivative in fw_derivatives\\n ),\\n formula=foreach_forward_grad_formula,\\n unpacked_arguments=unpacked_arguments,\\n )\\n )\\n\\n # Set all the grads at the end to avoid: https://github.com/pytorch/pytorch/issues/67367\\n content.append(\\\"\\\\n\\\".join(fw_grad_setters))\\n return content\\n\\n def get_any_has_fw_grad_cond(derivative: ForwardDerivative | None) -> str:\\n #\\n # Produces a condition string (e.g, \\\"isFwGradDefined(grad_output) || isFwGradDefined(output)\\\")\\n #\\n if derivative is None:\\n # (1) If a derivative is NOT provided, cond will check fw_grad of ALL differentiable inputs\\n # - Used in the out_fn case when we want to forbid fw derivatives\\n # - Used in the case where the fw_derivative is not defined, but we want\\n # To check if there is a decomposition registered for jvp\\n to_check: list[str] = []\\n for inp in list(\\n mapMaybe(\\n gen_differentiable_input,\\n f.func.arguments.non_out + list(f.func.arguments.out), # type: ignore[operator]\\n )\\n ):\\n if is_tensor_type(inp.type):\\n to_check.append(\\n FW_DERIVATIVE_CHECK_TEMPLATE.substitute(req_inp=inp.name)\\n )\\n elif is_tensor_list_type(inp.type):\\n to_check.append(\\n FW_DERIVATIVE_TENSORLIST_CHECK_TEMPLATE.substitute(\\n req_inp=inp.name\\n )\\n )\\n else:\\n raise RuntimeError(\\n f'Unsupported input type for \\\"{name}\\\" when forbidding forward AD usage.'\\n )\\n return f'({\\\" || \\\".join(to_check)})'\\n else:\\n # (2) If derivative is provided, use that information to determine which inputs\\n # to check fw_grad for\\n assert derivative.required_inputs_fw_grad is not None\\n\\n if len(derivative.required_inputs_fw_grad) == 0:\\n # Handle functions like stack\\n # For these, we don't unpack anything and always call the user function\\n if not (\\n len(differentiable_inputs) == 1\\n and is_tensor_list_type(differentiable_inputs[0].type)\\n ):\\n raise RuntimeError(\\n f'No differentiable input to \\\"{name}\\\" is a differentiable Tensor (as the provided '\\n \\\"forward AD formula does not use any input tangent) even though a forward gradient \\\"\\n \\\"formula has been defined for it. This case should only happen for function that \\\"\\n \\\"take a single TensorList as input. All other cases are not supported right now.\\\"\\n )\\n any_has_fw_grad = \\\"true\\\"\\n else:\\n any_has_fw_grad = \\\" || \\\".join(\\n [\\n (\\n FW_DERIVATIVE_TENSORLIST_CHECK_TEMPLATE\\n if is_tensor_list_type(inp.type)\\n else FW_DERIVATIVE_CHECK_TEMPLATE\\n ).substitute(req_inp=inp.name)\\n for inp in differentiable_inputs\\n if inp.name in derivative.required_inputs_fw_grad\\n ]\\n )\\n any_has_fw_grad = f\\\"({any_has_fw_grad})\\\"\\n\\n return any_has_fw_grad\\n\\n def emit_forbid_fw_derivatives(is_out_fn: bool = False) -> str:\\n if is_out_fn:\\n msg = \\\"because it is an out= function\\\"\\n else:\\n msg = (\\n \\\"because it has not been implemented yet.\\\\\\\\nPlease file an issue \\\"\\n \\\"to PyTorch at https://github.com/pytorch/pytorch/issues/new?template=feature-request.yml \\\"\\n \\\"so that we can prioritize its implementation.\\\"\\n )\\n cond = get_any_has_fw_grad_cond(derivative=None)\\n return (\\n FW_DERIVATIVE_FORBID_TEMPLATE.substitute(cond=cond, name=name, msg=msg)\\n if cond != \\\"\\\"\\n else \\\"\\\"\\n )\\n\\n body: list[str] = []\\n unpack_args_stats, unpacked_bindings = unpack_args(f)\\n\\n body.extend(unpack_args_stats)\\n if requires_derivative:\\n body.extend(emit_any_requires_grad())\\n body.extend(emit_any_has_forward_grad())\\n body.extend(emit_check_inplace())\\n body.extend(emit_original_self_definition())\\n body.extend(setup_derivative(differentiable_inputs))\\n\\n body.append(emit_call(f, unpacked_bindings, try_jit_decomposition))\\n if requires_derivative:\\n # set_flags has to appear after version_counter, because rebase_history\\n # requires that the counter is incremented before it is called\\n body.append(emit_history())\\n body.extend(emit_check_if_in_complex_autograd_allowlist())\\n\\n if is_out_fn:\\n body.append(emit_forbid_fw_derivatives(is_out_fn=True))\\n else:\\n if requires_derivative and not try_jit_decomposition:\\n if len(fw_derivatives) > 0:\\n body.extend(emit_fw_derivatives())\\n else:\\n body.append(emit_forbid_fw_derivatives())\\n\\n if requires_derivative:\\n # Save only after the forward AD has been set up\\n body.append(emit_save_outputs())\\n\\n if str(f.func.name.name) in RESET_GRAD_ACCUMULATOR:\\n # `inplace` implies that there is exactly one output named `self`,\\n # so we can keep the generated code easy. If you need to\\n # `reset_grad_accumulator` in an operator that's not `inplace`, you can\\n # remove this assert but the code generation will get more elaborate\\n assert inplace\\n body.append(\\\"reset_grad_accumulator(self);\\\")\\n if not returns_void:\\n body.append(f\\\"return {get_return_value(f)};\\\")\\n return body\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n#include \\\"torch/csrc/autograd/VariableTypeUtils.h\\\"\\n#include \\\"torch/csrc/autograd/generated/ViewFuncs.h\\\"\\n\\n#include \\n#include \\n#include \\n\\n// ${generated_comment}\\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\nusing namespace at;\\nusing torch::autograd::CreationMeta;\\nusing torch::autograd::as_view;\\nusing torch::autograd::increment_version;\\n\\nnamespace torch {\\n\\nnamespace ADInplaceOrView {\\n\\nnamespace {\\n${inplace_or_view_method_definitions}\\n} // namespace\\n} // namespace ADInplaceOrView\\n\\nnamespace {\\n\\nTORCH_LIBRARY_IMPL(aten, ADInplaceOrView, m) {\\n ${inplace_or_view_wrapper_registrations};\\n}\\n\\n} // namespace\\n} // namespace torch\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\\"torch/csrc/Device.h\\\"\\n#include \\\"torch/csrc/DynamicTypes.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n#include \\\"torch/csrc/autograd/python_nn_functions.h\\\"\\n#include \\\"torch/csrc/autograd/generated/python_return_types.h\\\"\\n#include \\\"torch/csrc/autograd/python_variable.h\\\"\\n#include \\\"torch/csrc/autograd/utils/wrap_outputs.h\\\"\\n#include \\\"torch/csrc/autograd/utils/python_arg_parsing.h\\\"\\n#include \\\"torch/csrc/utils/pycfunction_helpers.h\\\"\\n#include \\\"torch/csrc/utils/python_arg_parser.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n#include \\\"torch/csrc/utils/tensor_memoryformats.h\\\"\\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\nusing at::Tensor;\\nusing at::Scalar;\\nusing at::MemoryFormat;\\nusing at::Generator;\\nusing at::IntArrayRef;\\nusing at::ArrayRef;\\n\\nusing namespace torch::autograd::utils;\\n\\nnamespace torch::autograd {\\n\\nstatic PyObject* THPNNVariableFunctionsModule = NULL;\\n\\nstatic PyObject * THPVariable__parse_to(PyObject* module, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"to(Device device=None, ScalarType dtype=None, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"to(ScalarType dtype, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"to(Tensor tensor, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)\\\",\\n });\\n ParsedArgs<5> parsed_args;\\n auto r = parser.parse(args, kwargs, parsed_args);\\n if (r.has_torch_function()) {\\n return handle_torch_function(r, args, kwargs, THPNNVariableFunctionsModule, \\\"torch.nn\\\", \\\"_parse_to\\\");\\n }\\n auto parsed = parse_to_conversion(r, /*allow_copy*/ false); // we don't want copy for nn.Module.to\\n auto& device = std::get<0>(parsed);\\n auto& scalarType = std::get<1>(parsed);\\n auto non_blocking = std::get<2>(parsed);\\n auto opt_memory_format = std::get<4>(parsed);\\n auto tuple = THPObjectPtr{PyTuple_New(4)};\\n if (!tuple) throw python_error();\\n if (device) {\\n PyTuple_SET_ITEM(tuple.get(), 0, THPDevice_New(*device));\\n } else {\\n Py_INCREF(Py_None);\\n PyTuple_SET_ITEM(tuple.get(), 0, Py_None);\\n }\\n if (scalarType) {\\n PyTuple_SET_ITEM(tuple.get(), 1, Py_NewRef(torch::getTHPDtype(*scalarType)));\\n } else {\\n Py_INCREF(Py_None);\\n PyTuple_SET_ITEM(tuple.get(), 1, Py_None);\\n }\\n PyTuple_SET_ITEM(tuple.get(), 2, torch::autograd::utils::wrap(non_blocking));\\n if (opt_memory_format.has_value()) {\\n PyTuple_SET_ITEM(tuple.get(), 3, Py_NewRef(torch::utils::getTHPMemoryFormat(opt_memory_format.value())));\\n } else {\\n Py_INCREF(Py_None);\\n PyTuple_SET_ITEM(tuple.get(), 3, Py_None);\\n }\\n return tuple.release();\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// generated forward declarations start here\\n\\n${py_forwards}\\n\\nstatic PyMethodDef nn_functions[] = {\\n {\\\"_parse_to\\\", castPyCFunctionWithKeywords(THPVariable__parse_to),\\n METH_VARARGS | METH_KEYWORDS, nullptr},\\n ${py_method_defs}\\n {NULL}\\n};\\n\\nvoid initNNFunctions(PyObject* module) {\\n static struct PyModuleDef def = {\\n PyModuleDef_HEAD_INIT,\\n \\\"torch._C._nn\\\",\\n NULL,\\n -1,\\n nn_functions\\n };\\n PyObject* nn = PyModule_Create(&def);\\n THPNNVariableFunctionsModule = nn;\\n if (!nn) {\\n throw python_error();\\n }\\n // steals a reference to nn\\n if (PyModule_AddObject(module, \\\"_nn\\\", nn) != 0) {\\n throw python_error();\\n }\\n}\\n\\n// generated methods start here\\n\\n${py_methods}\\n\\n} // namespace torch::autograd\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\\"torch/csrc/Device.h\\\"\\n#include \\\"torch/csrc/DynamicTypes.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n#include \\\"torch/csrc/autograd/python_linalg_functions.h\\\"\\n#include \\\"torch/csrc/autograd/generated/python_return_types.h\\\"\\n#include \\\"torch/csrc/autograd/python_variable.h\\\"\\n#include \\\"torch/csrc/autograd/utils/wrap_outputs.h\\\"\\n#include \\\"torch/csrc/autograd/utils/python_arg_parsing.h\\\"\\n#include \\\"torch/csrc/utils/pycfunction_helpers.h\\\"\\n#include \\\"torch/csrc/utils/python_arg_parser.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\nusing at::Tensor;\\nusing at::Scalar;\\nusing at::ScalarType;\\nusing at::MemoryFormat;\\nusing at::Generator;\\nusing at::IntArrayRef;\\nusing at::TensorList;\\n\\nusing namespace torch::autograd::utils;\\n\\nnamespace torch::autograd {\\n\\n// generated forward declarations start here\\n\\n${py_forwards}\\n\\nstatic PyMethodDef linalg_functions[] = {\\n ${py_method_defs}\\n {NULL}\\n};\\n\\nstatic PyObject* THPLinalgVariableFunctionsModule = NULL;\\n\\nvoid initLinalgFunctions(PyObject* module) {\\n static struct PyModuleDef def = {\\n PyModuleDef_HEAD_INIT,\\n \\\"torch._C._linalg\\\",\\n NULL,\\n -1,\\n linalg_functions\\n };\\n PyObject* linalg = PyModule_Create(&def);\\n THPLinalgVariableFunctionsModule = linalg;\\n if (!linalg) {\\n throw python_error();\\n }\\n // steals a reference to linalg\\n if (PyModule_AddObject(module, \\\"_linalg\\\", linalg) != 0) {\\n throw python_error();\\n }\\n}\\n\\n// generated methods start here\\n\\n${py_methods}\\n\\n} // namespace torch::autograd\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\\"torch/csrc/Device.h\\\"\\n#include \\\"torch/csrc/DynamicTypes.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n#include \\\"torch/csrc/autograd/python_nested_functions.h\\\"\\n#include \\\"torch/csrc/autograd/generated/python_return_types.h\\\"\\n#include \\\"torch/csrc/autograd/python_variable.h\\\"\\n#include \\\"torch/csrc/autograd/utils/wrap_outputs.h\\\"\\n#include \\\"torch/csrc/autograd/utils/python_arg_parsing.h\\\"\\n#include \\\"torch/csrc/autograd/generated/variable_factories.h\\\"\\n#include \\\"torch/csrc/utils/out_types.h\\\"\\n#include \\\"torch/csrc/utils/pycfunction_helpers.h\\\"\\n#include \\\"torch/csrc/utils/python_arg_parser.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n#include \\\"torch/csrc/utils/device_lazy_init.h\\\"\\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\nusing at::Tensor;\\nusing at::Device;\\nusing at::Layout;\\nusing at::Scalar;\\nusing at::ScalarType;\\nusing at::Backend;\\nusing at::OptionalDeviceGuard;\\nusing at::DeviceGuard;\\nusing at::TensorOptions;\\nusing at::IntArrayRef;\\nusing at::OptionalIntArrayRef;\\nusing at::Generator;\\nusing at::TensorList;\\nusing at::Dimname;\\nusing at::DimnameList;\\n\\nusing namespace torch::autograd::utils;\\n\\nnamespace torch::autograd {\\n\\n// generated forward declarations start here\\n\\n${py_forwards}\\n\\nstatic PyMethodDef nested_functions[] = {\\n {NULL, NULL, 0, NULL},\\n ${py_method_defs}\\n {NULL}\\n};\\n\\nstatic PyObject* THPNestedVariableFunctionsModule = NULL;\\n\\nvoid initNestedFunctions(PyObject* module) {\\n nested_functions[0] = get_nested_functions_manual()[0];\\n static struct PyModuleDef def = {\\n PyModuleDef_HEAD_INIT,\\n \\\"torch._C._nested\\\",\\n NULL,\\n -1,\\n nested_functions\\n };\\n PyObject* nested = PyModule_Create(&def);\\n THPNestedVariableFunctionsModule = nested;\\n if (!nested) {\\n throw python_error();\\n }\\n // steals a reference to nested\\n if (PyModule_AddObject(module, \\\"_nested\\\", nested) != 0) {\\n throw python_error();\\n }\\n}\\n\\n// generated methods start here\\n\\n${py_methods}\\n\\n} // namespace torch::autograd\\n\\n\\n#include \\n#include \\n#include \\n#include \\n\\nnamespace py = pybind11;\\nnamespace torch {\\n namespace autograd {\\n void initEnumTag(PyObject* module) {\\n auto m = py::handle(module).cast();\\n py::enum_(m, \\\"Tag\\\")\\n ${enum_of_valid_tags};\\n m.doc() = \\\"An Enum that contains tags that can be assigned to an operator registered in C++.\\\";\\n }\\n}}\\n\\n\\n#include \\n\\n#include \\n#include \\n#include \\n\\n#include \\\"torch/csrc/autograd/generated/python_return_types.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n\\nnamespace torch { namespace autograd { namespace generated {\\n\\n${py_return_types}\\n\\n}}}\\n\\nnamespace torch::autograd {\\n\\nstatic void addReturnType(\\n PyObject* module,\\n const char* name,\\n PyTypeObject* type) {\\n // hold onto the TypeObject for the unlikely case of user\\n // deleting or overriding it.\\n Py_INCREF(type);\\n if (PyModule_AddObject(\\n module,\\n name,\\n (PyObject*)type) != 0) {\\n Py_DECREF(type);\\n throw python_error();\\n }\\n}\\n\\nvoid initReturnTypes(PyObject* module) {\\n static struct PyModuleDef def = {\\n PyModuleDef_HEAD_INIT, \\\"torch._C._return_types\\\", nullptr, -1, {}};\\n PyObject* return_types_module = PyModule_Create(&def);\\n if (!return_types_module) {\\n throw python_error();\\n }\\n\\n ${py_return_types_registrations}\\n\\n // steals a reference to return_types on success\\n if (PyModule_AddObject(module, \\\"_return_types\\\", return_types_module) != 0) {\\n Py_DECREF(return_types_module);\\n throw python_error();\\n }\\n}\\n\\n} // namespace torch::autograd\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n#include \\\"torch/csrc/jit/frontend/tracer.h\\\"\\n\\n#include \\n\\n#include \\\"torch/csrc/autograd/function.h\\\"\\n\\n#include \\\"ATen/quantized/Quantizer.h\\\"\\n\\n// ${generated_comment}\\n\\n// See the `Tracer` section in `torch/csrc/jit/OVERVIEW.md`.\\n// NOTE See [Sharded File] comment in VariableType\\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\nusing namespace at;\\n\\nnamespace torch {\\n\\nnamespace TraceType {\\n\\nnamespace {\\n${trace_method_definitions}\\n} // namespace\\n} // namespace TraceType\\n\\nnamespace {\\n\\nTORCH_LIBRARY_IMPL(aten, Tracer, m) {\\n ${trace_wrapper_registrations};\\n}\\n\\n} // namespace\\n\\n} // namespace torch\\n\\n\\n#include \\\"torch/csrc/autograd/VariableTypeUtils.h\\\"\\n#include \\\"torch/csrc/autograd/generated/VariableType.h\\\"\\n#include \\\"torch/csrc/autograd/FunctionsManual.h\\\"\\n\\n#include \\n#include \\n#include \\n#include \\n\\n#include \\n\\n\\n// ${generated_comment}\\n\\n// NOTE [Sharded File]: on this file's split-into-shards state\\n//\\n// Back in the good old days, VariableType.cpp was generated as one\\n// file with every function in it, and everything was great and\\n// simple.\\n//\\n// However, this file was also very large (over 36,000 lines), and\\n// compiling it was very slow, and in fact was a significant\\n// bottleneck for incremental rebuilds. To address this, we now\\n// generate the file split across multiple shards, named\\n// VariableType_0.cpp and so on, which can be compiled in parallel.\\n//\\n// For ease of inspection and debugging, so that it's not necessary to\\n// go rooting around in multiple files, we also generate all the\\n// functions together in VariableTypeEverything.cpp. This generated\\n// file is only for convenience; it's not actually used in the\\n// build. If the file you're looking at now is one of the shards, you\\n// may want to switch over to the Everything variant to make you\\n// grepping smoother.\\n\\nusing namespace at;\\nusing namespace torch::autograd::generated;\\nusing namespace torch::autograd::generated::details;\\n\\n\\nnamespace torch::autograd {\\n\\nnamespace VariableType {\\nnamespace{\\n C10_UNUSED void reset_grad_accumulator(Variable & self) {\\n AutogradMeta* meta = torch::autograd::impl::get_autograd_meta(self);\\n if (meta != nullptr) {\\n meta->grad_accumulator_.reset();\\n }\\n }\\n}\\n\\nnamespace {\\n\\n\\n${type_derived_method_definitions}\\n}\\n}\\n\\nnamespace {\\n\\n${wrapper_registrations}\\n\\n}\\n\\n} // namespace torch::autograd\\n\\n\\n#include \\\"torch/csrc/autograd/FunctionsManual.h\\\"\\n#include \\\"torch/csrc/dynamo/compiled_autograd.h\\\"\\n\\n// ${generated_comment}\\n\\n// The manual function definitions that used to be here are now in torch/csrc/autograd/FunctionsManual.cpp\\n// This speeds up re-compilation and allow to share these implementations so that they can be\\n// used for forward mode AD formulas as well.\\n\\nusing namespace torch::autograd::generated::details;\\nusing at::Tensor;\\nusing at::Scalar;\\nusing at::IntArrayRef;\\nusing at::TensorList;\\n\\nnamespace torch::autograd::generated {\\n\\n${autograd_function_definitions}\\n\\n} // namespace torch::autograd::generated\\n\\n\\n#include \\n\\n// ${generated_comment}\\n\\n#include \\n#include \\n\\n#include \\n#include \\\"torch/csrc/autograd/generated/Functions.h\\\"\\n#include \\\"torch/csrc/autograd/python_cpp_function.h\\\"\\n#include \\n#include \\n#include \\n#include \\n#include \\n\\n// NOTE: See [Sharded File] comment in VariableType\\n\\nnamespace torch::autograd::generated {\\n\\ntemplate\\nstatic void addClass(PyObject* module, PyTypeObject& type, const char* name,\\n PyGetSetDef* function_properties=NULL, PyMethodDef* function_methods=NULL)\\n{\\n _initFunctionPyTypeObject(type, name, function_properties, function_methods);\\n Py_INCREF(&type);\\n PyModule_AddObject(module, name, (PyObject*)&type);\\n registerCppFunction(typeid(C), &type);\\n}\\n\\n${py_function_props_and_getters}\\n\\nvoid initialize_autogenerated_functions${shard_id}(PyObject* module) {\\n ${py_function_initializers}\\n}\\n\\n} // namespace torch::autograd::generated\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\\"torch/csrc/Device.h\\\"\\n#include \\\"torch/csrc/DynamicTypes.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n#include \\\"torch/csrc/autograd/python_special_functions.h\\\"\\n#include \\\"torch/csrc/autograd/generated/python_return_types.h\\\"\\n#include \\\"torch/csrc/autograd/python_variable.h\\\"\\n#include \\\"torch/csrc/autograd/utils/wrap_outputs.h\\\"\\n#include \\\"torch/csrc/autograd/utils/python_arg_parsing.h\\\"\\n#include \\\"torch/csrc/autograd/generated/variable_factories.h\\\"\\n#include \\\"torch/csrc/utils/out_types.h\\\"\\n#include \\\"torch/csrc/utils/pycfunction_helpers.h\\\"\\n#include \\\"torch/csrc/utils/python_arg_parser.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n#include \\\"torch/csrc/utils/device_lazy_init.h\\\"\\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\nusing at::Tensor;\\nusing at::Device;\\nusing at::Layout;\\nusing at::Scalar;\\nusing at::ScalarType;\\nusing at::Backend;\\nusing at::OptionalDeviceGuard;\\nusing at::DeviceGuard;\\nusing at::TensorOptions;\\nusing at::IntArrayRef;\\nusing at::Generator;\\nusing at::TensorList;\\nusing at::Dimname;\\nusing at::DimnameList;\\n\\nusing torch::utils::check_out_type_matches;\\nusing namespace torch::autograd::utils;\\n\\nnamespace torch::autograd {\\n\\n// generated forward declarations start here\\n\\n${py_forwards}\\n\\nstatic PyMethodDef special_functions[] = {\\n ${py_method_defs}\\n {NULL}\\n};\\n\\nstatic PyObject* THPSpecialVariableFunctionsModule = NULL;\\n\\nvoid initSpecialFunctions(PyObject* module) {\\n static struct PyModuleDef def = {\\n PyModuleDef_HEAD_INIT,\\n \\\"torch._C._special\\\",\\n NULL,\\n -1,\\n special_functions\\n };\\n PyObject* special = PyModule_Create(&def);\\n THPSpecialVariableFunctionsModule = special;\\n if (!special) {\\n throw python_error();\\n }\\n // steals a reference to special\\n if (PyModule_AddObject(module, \\\"_special\\\", special) != 0) {\\n throw python_error();\\n }\\n}\\n\\n// generated methods start here\\n\\n${py_methods}\\n\\n} // namespace torch::autograd\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\\"torch/csrc/Device.h\\\"\\n#include \\\"torch/csrc/DynamicTypes.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n#include \\\"torch/csrc/autograd/python_fft_functions.h\\\"\\n#include \\\"torch/csrc/autograd/generated/python_return_types.h\\\"\\n#include \\\"torch/csrc/autograd/python_variable.h\\\"\\n#include \\\"torch/csrc/autograd/utils/wrap_outputs.h\\\"\\n#include \\\"torch/csrc/autograd/utils/python_arg_parsing.h\\\"\\n#include \\\"torch/csrc/autograd/generated/variable_factories.h\\\"\\n#include \\\"torch/csrc/utils/out_types.h\\\"\\n#include \\\"torch/csrc/utils/pycfunction_helpers.h\\\"\\n#include \\\"torch/csrc/utils/python_arg_parser.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n#include \\\"torch/csrc/utils/device_lazy_init.h\\\"\\n\\n#include \\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\nusing at::Tensor;\\nusing at::Device;\\nusing at::Layout;\\nusing at::Scalar;\\nusing at::ScalarType;\\nusing at::Backend;\\nusing at::OptionalDeviceGuard;\\nusing at::DeviceGuard;\\nusing at::TensorOptions;\\nusing at::IntArrayRef;\\nusing at::Generator;\\nusing at::TensorList;\\nusing at::Dimname;\\nusing at::DimnameList;\\n\\nusing torch::utils::check_out_type_matches;\\nusing namespace torch::autograd::utils;\\n\\nnamespace torch::autograd {\\n\\n// generated forward declarations start here\\n\\n${py_forwards}\\n\\nstatic PyMethodDef fft_functions[] = {\\n ${py_method_defs}\\n {NULL}\\n};\\n\\nstatic PyObject* THPFFTVariableFunctionsModule = NULL;\\n\\nvoid initFFTFunctions(PyObject* module) {\\n static struct PyModuleDef def = {\\n PyModuleDef_HEAD_INIT,\\n \\\"torch._C._fft\\\",\\n NULL,\\n -1,\\n fft_functions\\n };\\n PyObject* fft = PyModule_Create(&def);\\n THPFFTVariableFunctionsModule = fft;\\n if (!fft) {\\n throw python_error();\\n }\\n // steals a reference to fft\\n if (PyModule_AddObject(module, \\\"_fft\\\", fft) != 0) {\\n throw python_error();\\n }\\n}\\n\\n// generated methods start here\\n\\n${py_methods}\\n\\n} // namespace torch::autograd\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\n\\n// Undefine the copysign macro so that at::copysign works as intended with MSVC\\n// https://github.com/python/cpython/blob/c60394c7fc9cc09b16e9675a3eeb5844b6d8523f/PC/pyconfig.h#L196\\n#ifdef _MSC_VER\\n#undef copysign\\n#endif // _MSC_VER\\n\\n#include \\\"torch/csrc/DynamicTypes.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n#include \\\"torch/csrc/Size.h\\\"\\n#include \\\"torch/csrc/autograd/generated/VariableType.h\\\"\\n#include \\\"torch/csrc/autograd/python_variable.h\\\"\\n#include \\\"torch/csrc/autograd/utils/python_arg_parsing.h\\\"\\n#include \\\"torch/csrc/autograd/utils/error_messages.h\\\"\\n#include \\\"torch/csrc/autograd/utils/wrap_outputs.h\\\"\\n#include \\\"torch/csrc/jit/frontend/tracer.h\\\"\\n#ifdef USE_CUDA\\n#include \\\"torch/csrc/cuda/Event.h\\\"\\n#endif\\n#include \\\"torch/csrc/utils/device_lazy_init.h\\\"\\n#include \\n#include \\\"torch/csrc/utils/object_ptr.h\\\"\\n#include \\\"torch/csrc/utils/pycfunction_helpers.h\\\"\\n#include \\\"torch/csrc/utils/python_arg_parser.h\\\"\\n#include \\\"torch/csrc/utils/python_numbers.h\\\"\\n#include \\\"torch/csrc/utils/python_strings.h\\\"\\n#include \\\"torch/csrc/utils/python_tuples.h\\\"\\n#include \\\"torch/csrc/utils/tensor_apply.h\\\"\\n#include \\\"torch/csrc/utils/tensor_list.h\\\"\\n#include \\\"torch/csrc/utils/tensor_new.h\\\"\\n#include \\\"torch/csrc/utils/tensor_numpy.h\\\"\\n#include \\\"torch/csrc/utils/tensor_types.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n#include \\\"torch/csrc/autograd/generated/python_return_types.h\\\"\\n\\n#include \\n#include \\n#include \\\"c10/util/Optional.h\\\"\\n#include \\\"c10/core/Stream.h\\\"\\n\\n#include \\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#include \\n#endif\\n\\nusing at::DeviceGuard;\\nusing at::device_of;\\nusing at::OptionalDeviceGuard;\\nusing at::Backend;\\nusing at::Scalar;\\nusing at::ScalarType;\\nusing at::Tensor;\\nusing c10::Stream;\\nusing namespace torch::autograd::utils;\\n\\nnamespace torch::autograd {\\n\\nstatic PyObject * THPVariable__is_view(PyObject *self, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"_is_view\\\", args);\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n if (self_.is_view()) {\\n Py_RETURN_TRUE;\\n } else {\\n Py_RETURN_FALSE;\\n }\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object bc no support for first-class functions in native_functions.yaml\\n// See: ATen/native/README.md for more context\\nstatic PyObject * THPVariable_apply_(PyObject* self, PyObject* arg)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n auto args = py::make_tuple(py::handle(arg));\\n return handle_torch_function(self, \\\"apply_\\\", args.ptr());\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n if (self_.requires_grad()) {\\n throw std::runtime_error(\\n \\\"Can't call apply_() on Variable that requires grad. Use \\\"\\n \\\"var.detach().apply_() instead.\\\");\\n }\\n return THPVariable_Wrap(torch::utils::apply_(self_, arg));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_size(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"size(int64_t? dim=None)\\\",\\n \\\"size(Dimname dim)\\\",\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<3> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n if (r.idx == 0) {\\n if (!r.toInt64Optional(0).has_value()) {\\n return THPSize_NewFromSymSizes(self_);\\n }\\n if (jit::tracer::isTracing()) {\\n // will error out if a tensor has symints\\n return wrap(jit::tracer::getSizeOf(self_, r.toInt64(0)));\\n } else {\\n return torch::toPyObject(self_.sym_size(r.toInt64(0)));\\n }\\n } else if (r.idx == 1) {\\n if (jit::tracer::isTracing()) {\\n TORCH_INTERNAL_ASSERT(false, \\\"NYI: Named tensors w/ JIT\\\");\\n }\\n return wrap(self_.size(r.dimname(0)));\\n }\\n Py_RETURN_NONE;\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_stride(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"stride(int64_t? dim=None)\\\",\\n \\\"stride(Dimname dim)\\\",\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<3> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n if (r.idx == 0) {\\n if (r.toInt64Optional(0).has_value()) {\\n return torch::toPyObject(self_.sym_stride(r.toInt64(0)));\\n }\\n // yes, this is called strides in ATen.\\n at::SymIntArrayRef strides = self_.sym_strides();\\n // we can't do the normal wrapping here because IntArrayRef maps to both\\n // torch.Size and tuple in python\\n // TODO: consider factoring this out\\n THPObjectPtr tuple(PyTuple_New(strides.size()));\\n if (!tuple) throw python_error();\\n for (size_t i = 0; i != strides.size(); i++) {\\n PyObject* s = torch::toPyObject(strides[i]);\\n if (!s) throw python_error();\\n PyTuple_SET_ITEM(tuple.get(), i, s);\\n }\\n return tuple.release();\\n } else if (r.idx == 1) {\\n return wrap(self_.stride(r.dimname(0)));\\n }\\n Py_RETURN_NONE;\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object to avoid dispatch overhead\\nstatic PyObject * THPVariable_get_device(PyObject* self_, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self_)) {\\n return handle_torch_function(self_, \\\"get_device\\\", args, nullptr);\\n }\\n auto& self = THPVariable_Unpack(self_);\\n return wrap(self.get_device());\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_has_names(PyObject* self_, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self_)) {\\n return handle_torch_function(self_, \\\"has_names\\\", args);\\n }\\n auto& self = THPVariable_Unpack(self_);\\n return wrap(self.has_names());\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object to avoid dispatch overhead\\nstatic PyObject * THPVariable_data_ptr(PyObject* self_, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self_)) {\\n return handle_torch_function(self_, \\\"data_ptr\\\", args);\\n }\\n auto& self = THPVariable_Unpack(self_);\\n return wrap(self.data_ptr());\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object to avoid dispatch overhead\\nstatic PyObject * THPVariable_storage_offset(PyObject* self_, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self_)) {\\n return handle_torch_function(self_, \\\"storage_offset\\\");\\n }\\n auto& self = THPVariable_Unpack(self_);\\n return py::cast(self.sym_storage_offset()).release().ptr();\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object to avoid dispatch overhead\\nstatic PyObject * THPVariable_dim(PyObject* self, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"dim\\\", args);\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n return THPUtils_packInt64(self_.dim());\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object to avoid dispatch overhead\\nstatic PyObject * THPVariable_numel(PyObject* self, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"numel\\\", args);\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n if (jit::tracer::isTracing()) {\\n return wrap(jit::tracer::getNumelOf(self_));\\n } else {\\n return py::cast(self_.sym_numel()).release().ptr();\\n }\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic Tensor dispatch_contiguous(const Tensor & self, at::MemoryFormat memory_format) {\\n pybind11::gil_scoped_release no_gil;\\n OptionalDeviceGuard device_guard(device_of(self));\\n return self.contiguous(memory_format);\\n}\\n\\nstatic PyObject * THPVariable_contiguous(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"contiguous(*, MemoryFormat memory_format=contiguous_format)\\\",\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto& self_ = THPVariable_Unpack(self);\\n auto memory_format = r.memoryformat(0);\\n // avoids touching the GIL or current device if self is already contiguous\\n if (self_.is_contiguous(memory_format)) {\\n // NOTE: this logic is duplicated from VariableType.cpp. Since we need to\\n // record this call to contiguous() in the trace regardless of whether\\n // we actually call contiguous here, we need to record this information\\n // manually.\\n if (jit::tracer::isTracing()) {\\n auto tracer_state = jit::tracer::getTracingState();\\n auto op_name = c10::Symbol::fromQualString(\\\"aten::contiguous\\\");\\n auto node = tracer_state->createNode(op_name, /*num_outputs=*/0);\\n jit::tracer::recordSourceLocation(node);\\n jit::tracer::addInputs(node, \\\"self\\\", self_);\\n jit::tracer::addInputs(node, \\\"memory_format\\\", memory_format);\\n tracer_state->insertNode(node);\\n jit::tracer::addOutput(node, self_);\\n }\\n Py_INCREF(self);\\n return self;\\n }\\n return THPVariable_Wrap(dispatch_contiguous(self_, memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic Tensor dispatch_copy_(const Tensor & self, const Tensor & other, bool non_blocking) {\\n pybind11::gil_scoped_release no_gil;\\n OptionalDeviceGuard device_guard(device_of(self));\\n return self.copy_(other, non_blocking);\\n}\\n\\n static PyObject * THPVariable_copy_(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"copy_(Tensor other, bool non_blocking=False)\\\",\\n \\\"copy_(Tensor other, bool async=False)|deprecated\\\"\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<2> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n return THPVariable_Wrap(dispatch_copy_(self_, r.tensor(0), r.toBool(1)));\\n END_HANDLE_TH_ERRORS\\n}\\n\\ntemplate\\nstatic T dispatch_to(const Tensor & self) {\\n pybind11::gil_scoped_release no_gil;\\n OptionalDeviceGuard device_guard(device_of(self));\\n TORCH_CHECK_VALUE(self.sym_numel() == 1, \\\"only one element tensors can be converted to Python scalars\\\");\\n return self.template item();\\n}\\n\\nstatic PyObject * THPVariable_float_scalar(PyObject* self, PyObject* args) {\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"__float__\\\", args);\\n }\\n jit::tracer::warn(\\\"Converting a tensor to a Python float\\\", jit::tracer::WARN_PYTHON_DATAFLOW);\\n auto& self_ = THPVariable_Unpack(self);\\n return wrap(dispatch_to(self_));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_complex_scalar(PyObject* self, PyObject* args) {\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"__complex__\\\", args);\\n }\\n jit::tracer::warn(\\\"Converting a tensor to a Python complex\\\", jit::tracer::WARN_PYTHON_DATAFLOW);\\n auto& self_ = THPVariable_Unpack(self);\\n return wrap(dispatch_to>(self_));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_integral_scalar(PyObject* self, PyObject* args) {\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"__int__\\\", args);\\n }\\n jit::tracer::warn(\\\"Converting a tensor to a Python integer\\\", jit::tracer::WARN_PYTHON_DATAFLOW);\\n auto& self_ = THPVariable_Unpack(self);\\n if (isFloatingType(self_.scalar_type())) {\\n // we can't dispatch to item here because we want to avoid ATen overflow checks;\\n // the python integral type (long in python2) can't overflow.\\n return THPUtils_packDoubleAsInt(dispatch_to(self_));\\n } else {\\n return wrap(dispatch_to(self_));\\n }\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// This is the __index__ function in Python which is similar to __int__, but\\n// called when used as a slice.\\nstatic PyObject * THPVariable_index_scalar(PyObject* self, PyObject* args) {\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"__index__\\\", args);\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n // TODO: change the condition to `self_.dim() != 0` once we expose scalars\\n // in PyTorch.\\n if (!isIntegralType(self_.scalar_type(), /*includeBool=*/true) || self_.sym_numel() != 1) {\\n throw TypeError(\\\"only integer tensors of a single element can be converted to an index\\\");\\n }\\n return wrap(dispatch_to(self_));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic Tensor dispatch_invert(const Tensor & self) {\\n pybind11::gil_scoped_release no_gil;\\n OptionalDeviceGuard device_guard(device_of(self));\\n return self.bitwise_not();\\n}\\n\\nstatic PyObject * THPVariable_invert(PyObject* self, PyObject* args) {\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"__invert__\\\", args);\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n if (!isIntegralType(self_.scalar_type(), /*includeBool=*/true)) {\\n throw TypeError(\\\"~ (operator.invert) is only implemented on integer and Boolean-type tensors\\\");\\n }\\n return THPVariable_Wrap(dispatch_invert(self_));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic Tensor dispatch_to(const Tensor & self, Device device, bool non_blocking, bool copy, std::optional optional_memory_format) {\\n pybind11::gil_scoped_release no_gil;\\n // NOTE: this is where we record aten::to in the graph during tracing. However, the behavior of aten::to\\n // is different with respect to TensorOptions fields that are not present: aten::to inherits fields that\\n // are missing from the self argument while the tracer assumes that they should be populated with the\\n // default values (eg. float for scalar type). By explicitly copying over the tensor options here we fully\\n // specify all tensor options and thus record the proper trace\\n return self.to(self.options().device(device).memory_format(optional_memory_format), non_blocking, copy);\\n}\\n\\nstatic Tensor dispatch_to(const Tensor & self, bool non_blocking, bool copy, std::optional optional_memory_format) {\\n pybind11::gil_scoped_release no_gil;\\n return self.to(self.options().memory_format(optional_memory_format), non_blocking, copy);\\n}\\n\\nstatic Tensor dispatch_to(const Tensor & self, ScalarType dtype, bool non_blocking, bool copy, std::optional optional_memory_format) {\\n pybind11::gil_scoped_release no_gil;\\n // TODO: Make this call the TensorOptions version, maybe?\\n return self.to(dtype, non_blocking, copy, optional_memory_format);\\n}\\n\\nstatic Tensor dispatch_to(const Tensor & self, Device device, ScalarType dtype, bool non_blocking, bool copy, std::optional optional_memory_format) {\\n pybind11::gil_scoped_release no_gil;\\n // TODO: Make this call the TensorOptions version, maybe?\\n return self.to(device, dtype, non_blocking, copy, optional_memory_format);\\n}\\n\\nstatic PyObject * THPVariable_cpu(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"cpu(*, MemoryFormat? memory_format=None)\\\"\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_Wrap(dispatch_to(self_, at::Device(at::DeviceType::CPU), false, false, opt_memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic Tensor dispatch_nonzero(const Tensor & self) {\\n pybind11::gil_scoped_release no_gil;\\n OptionalDeviceGuard device_guard(device_of(self));\\n return self.nonzero();\\n}\\n\\nstatic std::vector dispatch_nonzero_numpy(const Tensor & self) {\\n pybind11::gil_scoped_release no_gil;\\n OptionalDeviceGuard device_guard(device_of(self));\\n return self.nonzero_numpy();\\n}\\n\\nstatic PyObject * THPVariable_nonzero(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"nonzero()\\\",\\n \\\"nonzero(*, bool as_tuple)\\\",\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<2> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n if (r.idx == 0 || (r.idx == 1 && !r.toBool(0))) {\\n return wrap(dispatch_nonzero(self_));\\n } else {\\n return wrap(dispatch_nonzero_numpy(self_));\\n }\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_cuda(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"cuda(Device? device=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"cuda(Device? device=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated\\\"\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<3> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto device = r.isNone(0) ? at::Device(at::DeviceType::CUDA) : r.device(0);\\n auto opt_memory_format = r.memoryformatOptional(2);\\n TORCH_CHECK(device.is_cuda(), \\\"Invalid device, must be cuda device\\\");\\n torch::utils::device_lazy_init(at::kCUDA);\\n return THPVariable_Wrap(dispatch_to(self_, device, r.toBool(1), false, opt_memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_mtia(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"mtia(Device? device=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"mtia(Device? device=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated\\\"\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<3> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if (r.has_torch_function()) {\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto device = r.isNone(0) ? at::Device(at::DeviceType::MTIA) : r.device(0);\\n auto opt_memory_format = r.memoryformatOptional(2);\\n TORCH_CHECK(device.is_mtia(), \\\"Invalid device, must be MTIA device\\\");\\n torch::utils::device_lazy_init(at::kMTIA);\\n return THPVariable_Wrap(dispatch_to(self_, device, r.toBool(1), false, opt_memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_xpu(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"xpu(Device? device=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"xpu(Device? device=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated\\\"\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<3> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if (r.has_torch_function()) {\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto device = r.isNone(0) ? at::Device(at::DeviceType::XPU) : r.device(0);\\n auto opt_memory_format = r.memoryformatOptional(2);\\n TORCH_CHECK(device.is_xpu(), \\\"Invalid device, must be xpu device\\\");\\n torch::utils::device_lazy_init(at::kXPU);\\n return THPVariable_Wrap(dispatch_to(self_, device, r.toBool(1), false, opt_memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_ipu(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"ipu(Device? device=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"ipu(Device? device=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated\\\"\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<3> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if (r.has_torch_function()) {\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto device = r.isNone(0) ? at::Device(at::DeviceType::IPU) : r.device(0);\\n auto opt_memory_format = r.memoryformatOptional(2);\\n TORCH_CHECK(device.is_ipu(), \\\"Invalid device, must be ipu device\\\");\\n return THPVariable_Wrap(dispatch_to(self_, device, r.toBool(1), false, opt_memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_to_type(PyObject* self, ScalarType scalarType, std::optional optional_memory_format) {\\n HANDLE_TH_ERRORS\\n auto& self_ = THPVariable_Unpack(self);\\n return THPVariable_Wrap(dispatch_to(self_, scalarType, false, false, optional_memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_byte(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"byte(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Byte, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_char(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"char(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Char, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_double(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"double(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Double, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_float(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"float(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Float, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_cdouble(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"cdouble(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::ComplexDouble, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_cfloat(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"cfloat(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::ComplexFloat, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_half(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"half(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Half, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_int(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"int(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Int, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_long(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"long(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Long, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_short(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"short(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Short, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_bool(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"bool(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::Bool, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_bfloat16(PyObject* self, PyObject* args, PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"bfloat16(*, MemoryFormat? memory_format=None)\\\"\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n auto opt_memory_format = r.memoryformatOptional(0);\\n return THPVariable_to_type(self, ScalarType::BFloat16, opt_memory_format);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_element_size(PyObject* self, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"element_size\\\", args);\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n return THPUtils_packInt64(self_.element_size());\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object bc PyObjects not declarable in native_functions.yaml\\n// See: ATen/native/README.md for more context\\nstatic PyObject * THPVariable_numpy(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"numpy(*, bool force=False)\\\"\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if (r.has_torch_function()) {\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n jit::tracer::warn(\\\"Converting a tensor to a NumPy array\\\", jit::tracer::WARN_PYTHON_DATAFLOW);\\n return torch::utils::tensor_to_numpy(self_, r.toBool(0));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_requires_grad_(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"requires_grad_(bool requires_grad=True)\\\",\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n // temporary hack to improve functorch UX.\\n const auto& functorch_tls = at::functorch::functorchTLSAccessor();\\n if (functorch_tls) {\\n functorch_tls->checkSupportsInplaceRequiresGrad();\\n }\\n\\n auto requires_grad = r.toBool(0);\\n // should we throw if requires_grad is true? var.requires_grad = True throws here\\n // but it's nice to let this be a no-op.\\n if (!self_.is_leaf() && !requires_grad) {\\n throw std::runtime_error(autograd::utils::requires_grad_leaf_error(requires_grad));\\n }\\n if (requires_grad && ! isDifferentiableType(at::typeMetaToScalarType(self_.dtype()))) {\\n throw std::runtime_error(\\\"only Tensors of floating point dtype can require gradients\\\");\\n }\\n self_.set_requires_grad(requires_grad);\\n return THPVariable_Wrap(self_);\\n END_HANDLE_TH_ERRORS\\n}\\n\\ninline bool dispatch_is_contiguous(const Tensor & self, MemoryFormat memory_format) {\\n return self.is_contiguous(memory_format);\\n}\\n\\n// implemented on the python object to avoid dispatch overhead\\nstatic PyObject * THPVariable_is_contiguous(PyObject* self_, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"is_contiguous(*, MemoryFormat memory_format=contiguous_format)\\\",\\n });\\n ParsedArgs<1> parsed_args;\\n auto r = parser.parse(self_, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self_, args, kwargs, PyObject_Type(self_), \\\"torch.Tensor\\\");\\n }\\n\\n auto memory_format = r.memoryformat(0);\\n auto& self = THPVariable_Unpack(self_);\\n return wrap(dispatch_is_contiguous(self, memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object to avoid dispatch overhead\\nstatic PyObject * THPVariable_item(PyObject* self, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"item\\\", args);\\n }\\n jit::tracer::warn(\\\"Converting a tensor to a Python number\\\", jit::tracer::WARN_PYTHON_DATAFLOW);\\n auto& self_ = THPVariable_Unpack(self);\\n auto dispatch_item_ = [](const Tensor& self) -> at::Scalar {\\n pybind11::gil_scoped_release no_gil;\\n return self.item();\\n };\\n return py::cast(dispatch_item_(self_)).release().ptr();\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object bc no support for first class functions in native_functions.yaml\\n// See: ATen/native/README.md for more context\\nstatic PyObject * THPVariable_map_(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({ \\\"map_(Tensor other, PyObject* callable)\\\" });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<2> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n Variable other = r.tensor(0);\\n if (self_.requires_grad() || other.requires_grad()) {\\n throw std::runtime_error(\\n \\\"Can't call map_() on Variable that requires grad. Use \\\"\\n \\\"var.detach().map_() instead.\\\");\\n }\\n TORCH_CHECK(\\n !self_.unsafeGetTensorImpl()->is_python_dispatch() && !other.unsafeGetTensorImpl()->is_python_dispatch(),\\n \\\".map_ is not supported for tensor subclasses.\\\");\\n\\n return THPVariable_Wrap(torch::utils::map_(self_, other, r.pyobject(1)));\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object bc no support for first class functions in native_functions.yaml\\n// See: ATen/native/README.md for more context\\nstatic PyObject * THPVariable_map2_(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({ \\\"map2_(Tensor x, Tensor y, PyObject* callable)\\\" });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<3> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n Variable x = r.tensor(0);\\n Variable y = r.tensor(1);\\n if (self_.requires_grad() || x.requires_grad() || y.requires_grad()) {\\n throw std::runtime_error(\\n \\\"Can't call map2_() on Variable that requires grad. Use \\\"\\n \\\"var.detach().map2_() instead.\\\");\\n }\\n TORCH_CHECK(\\n !x.unsafeGetTensorImpl()->is_python_dispatch() && !y.unsafeGetTensorImpl()->is_python_dispatch(),\\n \\\".map2_ is not supported for tensor subclasses.\\\");\\n return THPVariable_Wrap(torch::utils::map2_(self_, x, y, r.pyobject(2)));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_new(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"new\\\", args, kwargs);\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n OptionalDeviceGuard device_guard(device_of(self_));\\n return THPVariable_Wrap(torch::utils::legacy_tensor_new(legacyExtractDispatchKey(self_), self_.scalar_type(), args, kwargs));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_new_tensor(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"new_tensor\\\", args, kwargs);\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n OptionalDeviceGuard device_guard(device_of(self_));\\n return THPVariable_Wrap(torch::utils::new_tensor(legacyExtractDispatchKey(self_), self_.scalar_type(), args, kwargs));\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_storage(PyObject* self, PyObject* arg)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"untyped_storage\\\");\\n }\\n auto& self_ = THPVariable_Unpack(self);\\n return createPyObject(self_.storage());\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_to(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"to(Device device=None, ScalarType dtype=None, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"to(ScalarType dtype, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"to(Tensor tensor, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)\\\",\\n });\\n ParsedArgs<5> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n if (r.has_torch_function()) {\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n auto parsed = parse_to_conversion(r, /*allow_copy*/ true);\\n auto& device = std::get<0>(parsed);\\n auto& scalarType = std::get<1>(parsed);\\n auto non_blocking = std::get<2>(parsed);\\n auto copy = std::get<3>(parsed);\\n auto opt_memory_format = std::get<4>(parsed);\\n auto& self_ = THPVariable_Unpack(self);\\n torch::utils::maybe_initialize_device(device);\\n if (!device && !scalarType && !copy && !opt_memory_format.has_value()) {\\n Py_INCREF(self);\\n return self;\\n } else if (!device && !scalarType) {\\n return THPVariable_Wrap(\\n dispatch_to(self_, non_blocking, copy, opt_memory_format));\\n } else if (!device) {\\n return THPVariable_Wrap(dispatch_to(self_, *scalarType, non_blocking, copy, opt_memory_format));\\n } else if (!scalarType) {\\n return THPVariable_Wrap(dispatch_to(self_, *device, non_blocking, copy, opt_memory_format));\\n } else {\\n return THPVariable_Wrap(dispatch_to(self_, *device, *scalarType, non_blocking, copy, opt_memory_format));\\n }\\n Py_RETURN_NONE;\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// implemented on the python object b/c arbitrarily nested list not declarable in native_functions.yaml\\n// See: ATen/native/README.md for more context\\nstatic PyObject * THPVariable_tolist(PyObject* self, PyObject* args)\\n{\\n HANDLE_TH_ERRORS\\n if (check_has_torch_function(self)) {\\n return handle_torch_function(self, \\\"tolist\\\", args);\\n }\\n jit::tracer::warn(\\\"Converting a tensor to a Python list\\\", jit::tracer::WARN_PYTHON_DATAFLOW);\\n auto self_ = THPVariable_Unpack(self);\\n return torch::utils::tensor_to_list(self_);\\n END_HANDLE_TH_ERRORS\\n}\\n\\nstatic PyObject * THPVariable_type(PyObject* self, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n static PythonArgParser parser({\\n \\\"type(PyObject* dtype=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)\\\",\\n \\\"type(PyObject* dtype=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated\\\"\\n });\\n auto& self_ = THPVariable_Unpack(self);\\n ParsedArgs<3> parsed_args;\\n auto r = parser.parse(self, args, kwargs, parsed_args);\\n\\n if(r.has_torch_function()){\\n return handle_torch_function(r, self, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n\\n if (r.isNone(0)) {\\n return THPUtils_packString(torch::utils::options_to_string(self_.options()));\\n }\\n auto obj = r.pyobject(0);\\n auto opt_memory_format = r.memoryformatOptional(2);\\n std::string type_name;\\n bool is_dtype = false;\\n if (PyType_Check(obj)) {\\n if (obj == THPVariableClass) {\\n type_name = \\\"torch.Tensor\\\";\\n } else {\\n type_name = ((PyTypeObject*)obj)->tp_name;\\n }\\n } else if (THPUtils_checkString(obj)) {\\n type_name = THPUtils_unpackString(obj);\\n } else if (THPDtype_Check(obj)) {\\n is_dtype = true;\\n } else {\\n throw TypeError(\\\"dtype must be a type, str, or dtype object\\\");\\n }\\n ScalarType scalar_type;\\n Device device = self_.device();\\n if (is_dtype) {\\n scalar_type = r.scalartype(0);\\n return THPVariable_Wrap(dispatch_to(self_, scalar_type, /*non_blocking=*/ r.toBool(1), /*copy=*/ false, opt_memory_format));\\n }\\n at::TensorOptions options = torch::utils::options_from_string(type_name);\\n scalar_type = at::typeMetaToScalarType(options.dtype());\\n auto device_type = options.device().type();\\n if (device_type != device.type()) {\\n device = at::Device(device_type);\\n }\\n torch::utils::maybe_initialize_device(device);\\n return THPVariable_Wrap(dispatch_to(self_, device, scalar_type, /*non_blocking=*/ r.toBool(1), /*copy=*/ false, opt_memory_format));\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// generated methods start here\\n\\n${py_methods}\\n\\nstatic PyObject * THPVariable_bool_scalar(PyObject* self, PyObject* args) {\\n if (check_has_torch_function(self)) {\\n HANDLE_TH_ERRORS\\n return handle_torch_function(self, \\\"__bool__\\\", args);\\n END_HANDLE_TH_ERRORS\\n }\\n jit::tracer::warn(\\\"Converting a tensor to a Python boolean\\\", jit::tracer::WARN_PYTHON_DATAFLOW);\\n return THPVariable_is_nonzero(self, args);\\n}\\n\\nstatic PyObject * THPVariable___eq__(PyObject* self_, PyObject* args, PyObject* kwargs)\\n{\\n HANDLE_TH_ERRORS\\n#ifdef USE_NUMPY\\n if (torch::utils::is_numpy_available()) {\\n static PythonArgParser parser({\\n \\\"__eq__(PyObject* other)\\\",\\n }, /*traceable=*/true);\\n\\n ParsedArgs<1> parsed_args;\\n auto _r = parser.parse(self_, args, kwargs, parsed_args);\\n if(_r.has_torch_function()) {\\n return handle_torch_function(_r, self_, args, kwargs, THPVariableClass, \\\"torch.Tensor\\\");\\n }\\n switch (_r.idx) {\\n case 0: {\\n auto other = _r.pyobject(0);\\n if (PyArray_Check(other)) {\\n auto other_tensor = torch::utils::tensor_from_numpy(other);\\n auto dispatch_eq = [](const at::Tensor & self, const at::Tensor & other) -> at::Tensor {\\n pybind11::gil_scoped_release no_gil;\\n return self.eq(other);\\n };\\n const Tensor& self = THPVariable_Unpack(self_);\\n return wrap(dispatch_eq(self, other_tensor));\\n }\\n }\\n }\\n }\\n#endif\\n return THPVariable_eq(self_, args, kwargs);\\n Py_RETURN_NONE;\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// Wrapper converts a raised TypeError into returning NotImplemented\\n// Used to implement binary arithmetic operators\\ntemplate \\nstatic PyObject * TypeError_to_NotImplemented_(PyObject* self, PyObject* args, PyObject* kwargs) {\\n\\n PyObject* ret = Func(self, args, kwargs);\\n if (!ret && PyErr_ExceptionMatches(PyExc_TypeError)) {\\n PyErr_Clear();\\n Py_INCREF(Py_NotImplemented);\\n ret = Py_NotImplemented;\\n }\\n return ret;\\n}\\n\\n// set_ has to be defined in the template because the c10::Storage object\\n// does not have a type, and we need to make sure the Python storage object's\\n// type matches the tensor's type\\nstatic PyObject* THPVariable_set_(\\n PyObject* self_,\\n PyObject* args,\\n PyObject* kwargs) {\\n HANDLE_TH_ERRORS\\n const Tensor& self = THPVariable_Unpack(self_);\\n static PythonArgParser parser(\\n {\\n \\\"set_()\\\",\\n \\\"set_(Storage source)\\\",\\n \\\"set_(Storage source, SymInt storage_offset, SymIntArrayRef size, SymIntArrayRef stride=None)\\\",\\n \\\"set_(Tensor source)\\\",\\n \\\"set_(Tensor source, SymInt storage_offset, SymIntArrayRef size, SymIntArrayRef stride=None)\\\",\\n },\\n /*traceable=*/false);\\n\\n ParsedArgs<4> parsed_args;\\n auto _r = parser.parse(args, kwargs, parsed_args);\\n\\n switch (_r.idx) {\\n case 0: {\\n // aten::set_(Tensor(a!) self) -> Tensor(a!)\\n auto dispatch_set_ = [](const Tensor& self) -> Tensor {\\n pybind11::gil_scoped_release no_gil;\\n return self.set_();\\n };\\n return wrap(dispatch_set_(self));\\n }\\n case 1: {\\n // aten::set_.source_Storage(Tensor(a!) self, Storage source) ->\\n // Tensor(a!)\\n at::ScalarType storage_scalar_type;\\n bool is_typed_storage = true;\\n at::Storage storage = _r.storage(0, storage_scalar_type, is_typed_storage);\\n TORCH_CHECK(storage_scalar_type == self.dtype() || !is_typed_storage,\\n \\\"Expected a Storage of type \\\", self.dtype(),\\n \\\" or an UntypedStorage, but got type \\\", storage_scalar_type,\\n \\\" for argument 1 'storage'\\\");\\n auto dispatch_set_ = [](const Tensor& self, Storage source) -> Tensor {\\n pybind11::gil_scoped_release no_gil;\\n return self.set_(source);\\n };\\n return wrap(dispatch_set_(self, storage));\\n }\\n case 2: {\\n // aten::set_.source_Storage_storage_offset(Tensor(a!) self, Storage\\n // source, int storage_offset, int[] size, int[] stride=[]) -> Tensor(a!)\\n at::ScalarType storage_scalar_type;\\n bool is_typed_storage = true;\\n at::Storage storage = _r.storage(0, storage_scalar_type, is_typed_storage);\\n TORCH_CHECK(storage_scalar_type == self.dtype() || !is_typed_storage,\\n \\\"Expected a Storage of type \\\", self.dtype(),\\n \\\" or an UntypedStorage, but got type \\\", storage_scalar_type,\\n \\\" for argument 1 'storage'\\\");\\n auto dispatch_set_ = [](const Tensor& self,\\n Storage source,\\n c10::SymInt storage_offset,\\n c10::SymIntArrayRef size,\\n c10::SymIntArrayRef stride) -> Tensor {\\n pybind11::gil_scoped_release no_gil;\\n return self.set__symint(source, storage_offset, size, stride);\\n };\\n return wrap(dispatch_set_(\\n self, storage, _r.toSymInt(1), _r.symintlist(2), _r.symintlist(3)));\\n }\\n case 3: {\\n // aten::set_.source_Tensor(Tensor(a!) self, Tensor source) -> Tensor(a!)\\n auto dispatch_set_ = [](const Tensor& self, const Tensor& source) -> Tensor {\\n TORCH_CHECK(source.dtype() == self.dtype(), \\\"Could not set tensor of type \\\", source.dtype(), \\\" to a tensor of type \\\", self.dtype());\\n pybind11::gil_scoped_release no_gil;\\n return self.set_(source);\\n };\\n return wrap(dispatch_set_(self, _r.tensor(0)));\\n }\\n case 4: {\\n // aten::set_.source_Tensor_storage_offset(Tensor(a!) self, Tensor\\n // source, int storage_offset, int[] size, int[] stride=[]) -> Tensor(a!)\\n at::Tensor storage = _r.tensor(0);\\n auto dispatch_set_ = [](const Tensor& self,\\n const Tensor& source,\\n c10::SymInt storage_offset,\\n c10::SymIntArrayRef size,\\n c10::SymIntArrayRef stride) -> Tensor {\\n pybind11::gil_scoped_release no_gil;\\n return self.set__symint(source, storage_offset, size, stride);\\n };\\n return wrap(dispatch_set_(\\n self, storage, _r.toSymInt(1), _r.symintlist(2), _r.symintlist(3)));\\n }\\n }\\n Py_RETURN_NONE;\\n END_HANDLE_TH_ERRORS\\n}\\n\\n// XXX: ops that are bound here are not exposed to the C++ api nor the JIT.\\n// Any new ops added here should be accompanied with a comment why they are not\\n// being registered through native_functions.yaml, and be tagged cpp / JIT\\nPyMethodDef variable_methods[] = {\\n // These magic methods are all implemented on python object to wrap NotImplementedError\\n {\\\"__add__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__radd__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__iadd__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__rmul__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__mul__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__imul__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__sub__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__isub__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__div__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__truediv__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__floordiv__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__idiv__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__ifloordiv__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__mod__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__imod__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__eq__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__ne__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__lt__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__le__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__gt__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__ge__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__rand__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__ror__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__rxor__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"__bool__\\\", THPVariable_bool_scalar, METH_NOARGS, NULL},\\n {\\\"__float__\\\", THPVariable_float_scalar, METH_NOARGS, NULL},\\n {\\\"__complex__\\\", THPVariable_complex_scalar, METH_NOARGS, NULL},\\n {\\\"__int__\\\", THPVariable_integral_scalar, METH_NOARGS, NULL},\\n {\\\"__long__\\\", THPVariable_integral_scalar, METH_NOARGS, NULL},\\n {\\\"__index__\\\", THPVariable_index_scalar, METH_NOARGS, NULL},\\n {\\\"__nonzero__\\\", THPVariable_bool_scalar, METH_NOARGS, NULL},\\n {\\\"__invert__\\\", THPVariable_invert, METH_NOARGS, NULL},\\n {\\\"__matmul__\\\", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"_is_view\\\", THPVariable__is_view, METH_NOARGS, NULL},\\n {\\\"apply_\\\", THPVariable_apply_, METH_O, NULL},\\n {\\\"bfloat16\\\", castPyCFunctionWithKeywords(THPVariable_bfloat16), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"byte\\\", castPyCFunctionWithKeywords(THPVariable_byte), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"char\\\", castPyCFunctionWithKeywords(THPVariable_char), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"contiguous\\\", castPyCFunctionWithKeywords(THPVariable_contiguous), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"copy_\\\", castPyCFunctionWithKeywords(THPVariable_copy_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"cpu\\\", castPyCFunctionWithKeywords(THPVariable_cpu), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"cuda\\\", castPyCFunctionWithKeywords(THPVariable_cuda), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"mtia\\\", castPyCFunctionWithKeywords(THPVariable_mtia), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"xpu\\\", castPyCFunctionWithKeywords(THPVariable_xpu), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"ipu\\\", castPyCFunctionWithKeywords(THPVariable_ipu), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"data_ptr\\\", THPVariable_data_ptr, METH_NOARGS, NULL},\\n {\\\"dim\\\", THPVariable_dim, METH_NOARGS, NULL},\\n {\\\"has_names\\\", THPVariable_has_names, METH_NOARGS, NULL},\\n {\\\"double\\\", castPyCFunctionWithKeywords(THPVariable_double), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"cdouble\\\", castPyCFunctionWithKeywords(THPVariable_cdouble), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"element_size\\\", THPVariable_element_size, METH_NOARGS, NULL},\\n {\\\"float\\\", castPyCFunctionWithKeywords(THPVariable_float), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"cfloat\\\", castPyCFunctionWithKeywords(THPVariable_cfloat), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"get_device\\\", THPVariable_get_device, METH_NOARGS, NULL},\\n {\\\"bool\\\", castPyCFunctionWithKeywords(THPVariable_bool), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"half\\\", castPyCFunctionWithKeywords(THPVariable_half), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"int\\\", castPyCFunctionWithKeywords(THPVariable_int), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"is_contiguous\\\", castPyCFunctionWithKeywords(THPVariable_is_contiguous), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"item\\\", THPVariable_item, METH_NOARGS, NULL},\\n {\\\"long\\\", castPyCFunctionWithKeywords(THPVariable_long), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"map_\\\", castPyCFunctionWithKeywords(THPVariable_map_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"map2_\\\", castPyCFunctionWithKeywords(THPVariable_map2_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"ndimension\\\", THPVariable_dim, METH_NOARGS, NULL},\\n {\\\"nelement\\\", THPVariable_numel, METH_NOARGS, NULL},\\n {\\\"new\\\", castPyCFunctionWithKeywords(THPVariable_new), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"new_tensor\\\", castPyCFunctionWithKeywords(THPVariable_new_tensor), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"nonzero\\\", castPyCFunctionWithKeywords(THPVariable_nonzero), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"numel\\\", THPVariable_numel, METH_NOARGS, NULL},\\n {\\\"numpy\\\", castPyCFunctionWithKeywords(THPVariable_numpy), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"requires_grad_\\\", castPyCFunctionWithKeywords(THPVariable_requires_grad_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"set_\\\", castPyCFunctionWithKeywords(THPVariable_set_), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"short\\\", castPyCFunctionWithKeywords(THPVariable_short), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"size\\\", castPyCFunctionWithKeywords(THPVariable_size), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"untyped_storage\\\", THPVariable_storage, METH_NOARGS, NULL},\\n {\\\"storage_offset\\\", THPVariable_storage_offset, METH_NOARGS, NULL},\\n {\\\"stride\\\", castPyCFunctionWithKeywords(THPVariable_stride), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"to\\\", castPyCFunctionWithKeywords(THPVariable_to), METH_VARARGS | METH_KEYWORDS, NULL},\\n {\\\"tolist\\\", THPVariable_tolist, METH_NOARGS, NULL},\\n {\\\"type\\\", castPyCFunctionWithKeywords(THPVariable_type), METH_VARARGS | METH_KEYWORDS, NULL},\\n ${py_method_defs}\\n {NULL}\\n};\\n\\n} // namespace torch::autograd\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\\"torch/csrc/Device.h\\\"\\n#include \\\"torch/csrc/DynamicTypes.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n#include \\\"torch/csrc/autograd/python_sparse_functions.h\\\"\\n#include \\\"torch/csrc/autograd/python_variable.h\\\"\\n#include \\\"torch/csrc/autograd/utils/wrap_outputs.h\\\"\\n#include \\\"torch/csrc/autograd/utils/python_arg_parsing.h\\\"\\n#include \\\"torch/csrc/utils/pycfunction_helpers.h\\\"\\n#include \\\"torch/csrc/utils/python_arg_parser.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\nusing at::Tensor;\\nusing at::Scalar;\\nusing at::ScalarType;\\nusing at::MemoryFormat;\\nusing at::Generator;\\nusing at::IntArrayRef;\\nusing at::TensorList;\\n\\nusing namespace torch::autograd::utils;\\n\\nnamespace torch::autograd {\\n\\n// generated forward declarations start here\\n\\n${py_forwards}\\n\\nstatic PyMethodDef sparse_functions[] = {\\n ${py_method_defs}\\n {NULL}\\n};\\n\\nstatic PyObject* THPSparseVariableFunctionsModule = NULL;\\n\\nvoid initSparseFunctions(PyObject* module) {\\n static struct PyModuleDef def = {\\n PyModuleDef_HEAD_INIT,\\n \\\"torch._C._sparse\\\",\\n NULL,\\n -1,\\n sparse_functions\\n };\\n PyObject* sparse = PyModule_Create(&def);\\n THPSparseVariableFunctionsModule = sparse;\\n if (!sparse) {\\n throw python_error();\\n }\\n // steals a reference to sparse\\n if (PyModule_AddObject(module, \\\"_sparse\\\", sparse) != 0) {\\n throw python_error();\\n }\\n}\\n\\n// generated methods start here\\n\\n${py_methods}\\n\\n} // namespace torch::autograd\\n\\n\\n#include \\n\\n// ${generated_comment}\\n\\nusing at::Tensor;\\nusing at::Scalar;\\nusing at::IntArrayRef;\\nusing at::TensorList;\\n\\nnamespace torch::autograd::generated {\\n\\n${view_func_definitions}\\n\\n} // namespace torch::autograd::generated\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n// Python bindings for torch.* functions implemented through ATen.\\n//\\n// The functions are bound as static methods on a class\\n// torch._C._VariableFunctions which is also aliased as Variable._torch\\n// and also copied into 'torch' module.\\n\\n#include \\n\\n// Undefine the copysign macro so that at::copysign works as intended with MSVC\\n// https://github.com/python/cpython/blob/c60394c7fc9cc09b16e9675a3eeb5844b6d8523f/PC/pyconfig.h#L196\\n#ifdef _MSC_VER\\n#undef copysign\\n#endif // _MSC_VER\\n\\n#include \\\"torch/csrc/autograd/python_torch_functions.h\\\"\\n#include \\\"torch/csrc/autograd/python_variable.h\\\"\\n#include \\\"torch/csrc/autograd/utils/wrap_outputs.h\\\"\\n#include \\\"torch/csrc/Dtype.h\\\"\\n#include \\\"torch/csrc/DynamicTypes.h\\\"\\n#include \\\"torch/csrc/Exceptions.h\\\"\\n#include \\\"torch/csrc/utils/out_types.h\\\"\\n#include \\\"torch/csrc/utils/pybind.h\\\"\\n#include \\\"torch/csrc/utils/pycfunction_helpers.h\\\"\\n#include \\\"torch/csrc/utils/python_arg_parser.h\\\"\\n#include \\\"torch/csrc/utils/tensor_layouts.h\\\"\\n#include \\\"torch/csrc/utils/tensor_new.h\\\"\\n#include \\\"torch/csrc/utils/tensor_numpy.h\\\"\\n#include \\\"torch/csrc/jit/frontend/tracer.h\\\"\\n#include \\\"torch/csrc/autograd/generated/variable_factories.h\\\"\\n#include \\\"torch/csrc/utils/structseq.h\\\"\\n#include \\\"torch/csrc/utils/device_lazy_init.h\\\"\\n#include \\\"torch/csrc/autograd/generated/python_return_types.h\\\"\\n\\n#include \\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n$ops_headers\\n#endif\\n\\n#include \\n#include \\n#include \\n#include \\n\\nusing at::Tensor;\\nusing at::Device;\\nusing at::Layout;\\nusing at::Scalar;\\nusing at::ScalarType;\\nusing at::Backend;\\nusing at::OptionalDeviceGuard;\\nusing at::DeviceGuard;\\nusing at::TensorOptions;\\nusing at::IntArrayRef;\\nusing at::Generator;\\nusing at::TensorList;\\nusing at::Dimname;\\nusing at::DimnameList;\\nusing at::ArrayRef;\\n\\nusing torch::utils::check_out_type_matches;\\nusing namespace torch::autograd::utils;\\n\\n// NOTE: See [Sharded File] comment in VariableType\\n\\nnamespace torch::autograd {\\n\\n// generated forward declarations start here\\n\\n${py_forwards}\\n\\nstatic PyMethodDef torch_functions_shard[] = {\\n ${py_method_defs}\\n};\\n\\nvoid gatherTorchFunctions${shard_id}(std::vector &torch_functions) {\\n constexpr size_t num_functions = sizeof(torch_functions_shard) / sizeof(torch_functions_shard[0]);\\n torch_functions.insert(\\n torch_functions.end(),\\n torch_functions_shard,\\n torch_functions_shard + num_functions);\\n}\\n\\n// generated methods start here\\n\\n${py_methods}\\n\\n} // namespace torch::autograd\\n\\n\\n#define TORCH_ASSERT_NO_OPERATORS\\n\\n#include \\n#include \\n#include \\n\\nnamespace at {\\n\\n// NB: this is explicitly copied here (via codegen) rather than\\n// included via NativeFunctions.h to avoid recompiling this file when\\n// NativeFunctions.h changes\\nnamespace meta {\\n${meta_declaration}\\n}\\n\\nnamespace native {\\n${native_declaration}\\n${native_definitions}\\n}} // namespace at::native\\n\\n\\n#include \\n#include \\n\\n#include \\n\\nnamespace at {\\n\\nnamespace {\\n\\n// Verifies the requested type is the same as the Tensor's type.\\nvoid check_type(const TensorBase& tensor, ScalarType type, c10::string_view type_name) {\\n TORCH_CHECK(\\n tensor.scalar_type() == type\\n || (isQIntType(tensor.scalar_type())\\n && toUnderlying(tensor.scalar_type()) == type),\\n \\\"expected scalar type \\\", type_name, \\\" but found \\\", tensor.scalar_type());\\n}\\n\\n} // namespace\\n\\n#define DEFINE_CAST(T, name) \\\\\\n template <> \\\\\\n TORCH_API const T* TensorBase::const_data_ptr() const { \\\\\\n check_type(*this, ScalarType::name, #name); \\\\\\n return this->unsafeGetTensorImpl()->data_ptr_impl(); \\\\\\n } \\\\\\n \\\\\\n template <> \\\\\\n TORCH_API const T* TensorBase::const_data_ptr() const { \\\\\\n check_type(*this, ScalarType::name, #name); \\\\\\n return this->unsafeGetTensorImpl()->data_ptr_impl>(); \\\\\\n } \\\\\\n \\\\\\n template <> \\\\\\n TORCH_API T* TensorBase::mutable_data_ptr() const { \\\\\\n check_type(*this, ScalarType::name, #name); \\\\\\n return this->unsafeGetTensorImpl()->mutable_data_ptr_impl(); \\\\\\n } \\\\\\n \\\\\\n template <> \\\\\\n TORCH_API T* TensorBase::data_ptr() const { \\\\\\n return mutable_data_ptr(); \\\\\\n } \\\\\\n\\n AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CAST)\\n AT_FORALL_QINT_TYPES(DEFINE_CAST)\\n DEFINE_CAST(uint16_t, UInt16)\\n DEFINE_CAST(uint32_t, UInt32)\\n DEFINE_CAST(uint64_t, UInt64)\\n #undef DEFINE_CAST\\n\\n #define DEFINE_ITEM(T, name) \\\\\\n template <> \\\\\\n TORCH_API T Tensor::item() const { \\\\\\n return item().to##name(); \\\\\\n }\\n\\n AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_ITEM)\\n #undef DEFINE_ITEM\\n\\n } //namespace at\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#include \\n#else\\n// needed for the meta tensor calls to get stride info in functionalization\\n#include \\n// needed for special handling of copy_().\\n// See Note [functionalizating copy_() and not preserving strides]\\n#include \\n#include \\n\\n$ops_headers\\n#endif\\n\\nnamespace at {\\nnamespace functionalization {\\n\\n// This keyset is used by functionalization when it calls into meta kernels\\n// to accurately propagate stride metadata.\\n// Exclude any modes: the purpose of calling into meta kernels is only as an implementation\\n// detail to perform shape inference, and we don't want any modal keys to run.\\n// Specifically, we want to prevent functionalization and Python modes from running.\\nconstexpr auto exclude_keys_for_meta_dispatch =\\n c10::functorch_transforms_ks |\\n c10::DispatchKeySet({\\n c10::DispatchKey::FuncTorchDynamicLayerBackMode,\\n c10::DispatchKey::FuncTorchDynamicLayerFrontMode,\\n c10::DispatchKey::Python,\\n c10::DispatchKey::PreDispatch,\\n\\n });\\n\\n// Helper around at::has_internal_overlap.\\n// The ATen util is used in hot-path eager mode: it's always fast,\\n// but might return TOO_HARD sometimes.\\n// During functionalization, we're ok taking a bit longer\\n// to detect memory overlap.\\ninline bool has_internal_overlap_helper(const at::Tensor t) {\\n auto has_overlap = at::has_internal_overlap(t);\\n if (has_overlap == at::MemOverlap::Yes) return true;\\n if (has_overlap == at::MemOverlap::No) return false;\\n return false;\\n}\\n\\n\\ninline Tensor to_meta(const Tensor& t) {\\n if (!t.defined()) return t;\\n return at::native::empty_strided_meta_symint(t.sym_sizes(), t.sym_strides(),\\n/*dtype=*/std::make_optional(t.scalar_type()), /*layout=*/std::make_optional(t.layout()),\\n/*device=*/std::make_optional(c10::Device(kMeta)), /*pin_memory=*/std::nullopt);\\n}\\n\\ninline std::optional to_meta(const std::optional& t) {\\n if (t.has_value()) {\\n return std::make_optional(to_meta(*t));\\n }\\n return std::nullopt;\\n}\\n\\ninline std::vector to_meta(at::ITensorListRef t_list) {\\n std::vector outputs;\\n outputs.reserve(t_list.size());\\n for (const auto& tensor : t_list) {\\n outputs.push_back(to_meta(tensor));\\n }\\n return outputs;\\n}\\n\\ninline c10::List to_meta(const c10::List& t_list) {\\n c10::List outputs;\\n outputs.reserve(t_list.size());\\n for (const auto i : c10::irange(t_list.size())) {\\n outputs.push_back(to_meta(t_list[i]));\\n }\\n return outputs;\\n}\\n\\ninline c10::List<::std::optional> to_meta(const c10::List<::std::optional>& t_list) {\\n c10::List<::std::optional> outputs;\\n outputs.reserve(t_list.size());\\n for (const auto i : c10::irange(t_list.size())) {\\n outputs.push_back(to_meta(t_list[i]));\\n }\\n return outputs;\\n}\\n\\n\\n${func_definitions}\\n\\n} // namespace functionalization\\n\\nnamespace {\\n\\nTORCH_LIBRARY_IMPL(aten, Functionalize, m) {\\n ${func_registrations};\\n}\\n\\n} // namespace\\n\\n} // namespace at\\n\\n\\n#include \\n#include \\n#include \\n\\n#include \\n\\n// ${generated_comment}\\n\\n// NOTE [Sharded File]: This file is generated in a sharded fashion to speed up\\n// incremental rebuilds. See the comment at the top of\\n// templates/VariableType.cpp for an analogous, in-depth discussion.\\n//\\n// Generated by tools/jit/gen_unboxing.py. This file registers all ATen ops into JIT op registry instead of c10\\n// dispatcher. JIT op registry only takes boxed kernels, so we are calling unboxing functions in UnboxingFunctions.h\\n// to cast arguments into C++ types (instead of IValue) and delegate to unboxed kernels.\\n\\nnamespace torch { namespace jit {\\n\\nusing autograd::Variable;\\nusing autograd::variable_list;\\nusing at::Scalar;\\nusing at::ScalarType;\\nusing at::Tensor;\\nusing at::TensorOptions;\\nusing at::DeviceGuard;\\n\\nusing ::c10::fmap;\\nusing ::c10::filter;\\n\\nnamespace {\\n\\nRegisterOperators reg({\\n\\n // Generated operators\\n ${unboxed_ops}\\n});\\n\\n} // anon namespace\\n\\n\\n}} // namespace torch::jit\\n\\n\\n// ${generated_comment}\\n${includes}\\n${native_functions_include}\\n\\nnamespace {\\n${helper_fns}\\n} // namespace\\n\\n${namespace_prologue}\\n\\n${native_function_definitions}\\n\\n${namespace_epilogue}\\n\\n\\n#include \\n#include \\n\\n// ${generated_comment}\\n// NOTE See [Sharded File] comment in VariableType\\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n${operator_headers}\\n#endif\\n\\n${static_dispatch_extra_headers}\\n\\nnamespace at { namespace _ops {\\n\\n${definitions}\\n\\n}} // namespace at::_ops\\n\\n\\n// ${generated_comment}\\n\\n#include \\n#include \\n\\n#include \\n#include \\n\\nnamespace at {\\n\\nnamespace redispatch {\\n ${function_redispatch_definitions}\\n} // namespace redispatch\\n\\n} // namespace at\\n\\n\\n#include \\n\\n#include \\n#include \\n#include \\n#include \\n#include \\n\\n// ${generated_comment}\\n\\nnamespace at {\\n\\nnamespace {\\nstruct OpNameEquals final {\\n bool operator()(const std::pair& lhs, const std::pair& rhs) const {\\n return 0 == strcmp(lhs.first, rhs.first) && 0 == strcmp(lhs.second, rhs.second);\\n }\\n};\\n\\nstruct OpNameHash final {\\n size_t operator()(const std::pair& p) const {\\n // use std::hash because std::hash would hash pointers and not pointed-to strings\\n return std::hash()(p.first) ^ (~ std::hash()(p.second));\\n }\\n};\\n}\\n\\nbool is_custom_op(const c10::OperatorName& opName) {\\n static std::unordered_set, OpNameHash, OpNameEquals> ops {\\n ${aten_ops}\\n {\\\"\\\", \\\"\\\"}\\n };\\n return ops.count(std::make_pair(\\n opName.name.c_str(), opName.overload_name.c_str())) == 0;\\n}\\n}\\n\\n\\n// required for old g++ to compile PRId64 macros, see\\n// https://github.com/pytorch/pytorch/issues/3571\\n// for context\\n#ifndef __STDC_FORMAT_MACROS\\n#define __STDC_FORMAT_MACROS\\n#endif\\n\\n// an external backend might generate file within its code tree\\n// and check all the source files within the tree with clang-format.\\n// so, disable it since the backend might have a different config.\\n// clang-format off\\n\\n// NOTE: This condition is true for all PyTorch internal libraries, it\\n// just excludes external projects such as torch_xla which\\n// re-use some of the PyTorch codegen machinery.\\n#if defined(CAFFE2_BUILD_MAIN_LIB) || \\\\\\n defined(TORCH_CUDA_BUILD_MAIN_LIB) || \\\\\\n defined(TORCH_HIP_BUILD_MAIN_LIB) || \\\\\\n defined(TORCH_CUDA_CU_BUILD_MAIN_LIB) || \\\\\\n defined(TORCH_CUDA_CPP_BUILD_MAIN_LIB)\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n#endif\\n\\n// ${generated_comment}\\n\\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n\\n#include \\n#include \\n#include \\n#include \\n\\n#include \\n#include \\n#include \\n$extra_cuda_headers\\n$external_backend_headers\\n$dispatch_headers\\n$ops_headers\\n\\n// See template file RegisterDispatchDefinitions.ini\\n$dispatch_definitions\\n\\n\\n#define TORCH_ASSERT_NO_OPERATORS\\n\\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n\\nnamespace at {\\nnamespace native {\\n${native_definitions}\\n}} // namespace at::native\\n\\n\\n// We register ops with a higher priority dispatch key (BackendSelect) than the usual backend-specific keys (e.g. CPU)\\n// which makes calls to the factory functions dispatch to here.\\n// We then 'manually' compute a lower-priority to re-dispatch to (e.g. CPU) to get to the eventually correct backend.\\n// ${generated_comment}\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n#include \\n#include \\n#include \\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n\\n${ops_headers}\\n#endif\\n\\nnamespace at {\\n\\nnamespace {\\n\\n${backend_select_method_definitions}\\n\\nTORCH_LIBRARY_IMPL(aten, BackendSelect, m) {\\n ${backend_select_function_registrations};\\n}\\n\\n} // namespace\\n} // at\\n\\n\\n// ${generated_comment}\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n#include \\n\\nnamespace at {\\nTORCH_LIBRARY(aten, m) {\\n ${aten_schema_registrations};\\n // Distributed Ops\\n // Implementations located in torch/csrc/jit/runtime/register_distributed_ops.cpp\\n m.def(\\\"get_gradients(int context_id) -> Dict(Tensor, Tensor)\\\");\\n}\\n${schema_registrations}\\n} // namespace at\\n\\n\\n#define TORCH_ASSERT_ONLY_METHOD_OPERATORS\\n// ${generated_comment}\\n\\n#include \\n#include \\n#include \\n\\n#ifndef AT_PER_OPERATOR_HEADERS\\n#include \\n#else\\n#include \\n$ops_headers\\n#endif\\n\\nnamespace at {\\nnamespace native {\\n\\n// This file contains a number of kernels for aten functions that are fully code-generated.\\n// TODO: rename this file to something more generic.\\n\\nnamespace {\\nat::Tensor clone_arg(const at::Tensor& t) {\\n return t.clone();\\n}\\n\\nstd::vector clone_arg(const at::TensorList& t_list) {\\n std::vector out(t_list.size());\\n for (const auto& i : c10::irange(t_list.size())) {\\n out[i] = t_list[i].clone();\\n }\\n return out;\\n}\\n\\n// duped with gen_resize_out_helper from structured kernels\\nvoid copy_arg(const at::Tensor& dst, const at::Tensor& src) {\\n TORCH_CHECK(src.dtype() == dst.dtype(),\\n \\\"Expected out tensor to have dtype \\\", src.dtype(), \\\", but got \\\", dst.dtype(), \\\" instead\\\");\\n TORCH_CHECK(src.device() == dst.device(),\\n \\\"Expected out tensor to have device \\\", src.device(), \\\", but got \\\", dst.device(), \\\" instead\\\");\\n dst.copy_(src);\\n}\\n\\nvoid copy_arg(const at::TensorList& dst, const at::TensorList& src) {\\n TORCH_INTERNAL_ASSERT(dst.size() == src.size());\\n for (const auto& i : c10::irange(dst.size())) {\\n copy_arg(dst[i], src[i]);\\n }\\n}\\n\\n// TODO: this doesn't handle restriding empty tensors correctly; see\\n// gen_resize_out_helper for the correct algorithm\\n\\nvoid resize_out_helper(const at::Tensor& dst, const at::Tensor& src) {\\n at::native::resize_output(dst, src.sizes());\\n}\\n\\nvoid resize_out_helper(const at::TensorList& dst, const at::TensorList& src) {\\n TORCH_INTERNAL_ASSERT(dst.size() == src.size());\\n for (const auto& i : c10::irange(dst.size())) {\\n at::native::resize_output(dst[i], src[i].sizes());\\n }\\n}\\n}\\n\\n\\n${CompositeViewCopyKernel_Definitions}\\n\\n${GeneratedCompositeFunctional_Definitions}\\n\\n${GeneratedCompositeOut_Definitions}\\n\\n} // namespace native\\n} // namespace at\\n\\n\\n#include \\n#include \\n\\n#include \\n#include \\n#include \\n#include \\n#include \\n\\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\nnamespace at {\\nnamespace unboxing {\\n\\nusing ::c10::fmap;\\nusing ::c10::filter;\\nusing torch::jit::peek;\\nusing torch::jit::drop;\\nusing torch::jit::pack;\\nusing torch::jit::pop;\\n\\n// Generated function declaration\\n${definitions}\\n\\n} // namespace unboxing\\n} // namespace at\\n\\n\\n#include \\n\\n#include \\n#include \\n#include \\n\\nnamespace at {\\n\\nTensor TensorMaker::make_tensor() {\\n AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove.\\n tracer::impl::NoTracerDispatchMode tracer_guard{};\\n\\n check_size_nonnegative(sizes_);\\n\\n TORCH_CHECK_VALUE(\\n !deleter_ || !ctx_,\\n \\\"The deleter and context arguments are mutually exclusive.\\\");\\n\\n if (device_ == std::nullopt) {\\n device_ = globalContext().getDeviceFromPtr(data_, opts_.device().type());\\n }\\n\\n if (opts_.device().has_index()) {\\n // clang-format off\\n TORCH_CHECK_VALUE(\\n opts_.device() == *device_,\\n \\\"Specified device \\\", opts_.device(), \\\" does not match device of data \\\", *device_);\\n // clang-format on\\n }\\n\\n std::size_t size_bytes = computeStorageSize();\\n\\n DataPtr data_ptr{};\\n if (deleter_) {\\n data_ptr = makeDataPtrFromDeleter();\\n } else {\\n data_ptr = makeDataPtrFromContext();\\n }\\n\\n TORCH_CHECK(!resizeable_ || allocator_ != nullptr, \\\"Must specify an allocator with allocator() if you want to use resizeable_storage()\\\");\\n Storage storage{Storage::use_byte_size_t{}, size_bytes, std::move(data_ptr), /*allocator=*/allocator_, /*resizable=*/resizeable_};\\n\\n Tensor tensor = detail::make_tensor(\\n std::move(storage), opts_.computeDispatchKey(), opts_.dtype());\\n\\n TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl();\\n if (strides_) {\\n tensor_impl->set_sizes_and_strides(sizes_, *strides_);\\n } else {\\n tensor_impl->set_sizes_contiguous(sizes_);\\n }\\n if (storage_offset_) {\\n tensor_impl->set_storage_offset(*storage_offset_);\\n }\\n\\n return tensor;\\n }\\n\\n std::size_t TensorMaker::computeStorageSize() const noexcept {\\n std::size_t itemsize = opts_.dtype().itemsize();\\n\\n if (strides_) {\\n auto storage_size = detail::computeStorageNbytes(sizes_, *strides_, itemsize);\\n if (storage_offset_) {\\n storage_size += storage_offset_.value();\\n }\\n return storage_size;\\n }\\n\\n std::size_t size = 1;\\n for (std::int64_t s : sizes_) {\\n size *= static_cast(s);\\n }\\n auto storage_size = size * itemsize;\\n if (storage_offset_) {\\n storage_size += storage_offset_.value();\\n }\\n return storage_size;\\n }\\n\\n inline DataPtr TensorMaker::makeDataPtrFromDeleter() noexcept {\\n return InefficientStdFunctionContext::makeDataPtr(data_, std::move(deleter_), *device_);\\n }\\n\\n inline DataPtr TensorMaker::makeDataPtrFromContext() noexcept {\\n return DataPtr{data_, ctx_.release(), ctx_.get_deleter(), *device_};\\n }\\n\\n IntArrayRef TensorMaker::makeTempSizes() const noexcept {\\n static std::int64_t zeros[5] = {0, 0, 0, 0, 0};\\n if (opts_.has_memory_format()) {\\n MemoryFormat format = *opts_.memory_format_opt();\\n if (format == MemoryFormat::ChannelsLast) {\\n return IntArrayRef(zeros, 4);\\n }\\n if (format == MemoryFormat::ChannelsLast3d) {\\n return IntArrayRef(zeros, 5);\\n }\\n }\\n return IntArrayRef(zeros, 1);\\n }\\n\\n} // namespace at\\n\\n\\n#!/usr/bin/env python3\\n\\nfrom __future__ import annotations\\n\\nimport os\\nfrom enum import Enum\\nfrom operator import itemgetter\\nfrom pathlib import Path\\nfrom typing import Any\\n\\nimport torch\\nfrom torch.jit.generate_bytecode import generate_upgraders_bytecode\\nfrom torchgen.code_template import CodeTemplate\\nfrom torchgen.operator_versions.gen_mobile_upgraders_constant import (\\n MOBILE_UPGRADERS_HEADER_DESCRIPTION,\\n)\\n\\n\\nclass ByteCode(Enum):\\n instructions = 1\\n constants = 2\\n types = 3\\n operators = 4\\n register_size = 5\\n\\n\\nEXCLUDED_OP_SET = [\\n \\\"aten::full.names\\\",\\n \\\"aten::full.out\\\",\\n \\\"aten::full\\\",\\n]\\n\\nEXCLUE_UPGRADER_SET = [\\\"full_0_4\\\", \\\"full_out_0_4\\\"]\\n\\nONE_INSTRUCTION = CodeTemplate(\\n \\\"\\\"\\\"\\n Instruction{OpCode::${operator_name}, ${X}, ${N}},\\\"\\\"\\\"\\n)\\n\\nINSTRUCTION_LIST = CodeTemplate(\\n \\\"\\\"\\\"std::vector({\\n ${instruction_list}\\n }), // instructions list\\\"\\\"\\\"\\n)\\n\\nONE_CONSTANT = CodeTemplate(\\n \\\"\\\"\\\"\\n c10::IValue(${constant}),\\\"\\\"\\\"\\n)\\n\\nCONSTANT_LIST = CodeTemplate(\\n \\\"\\\"\\\"std::vector({\\n ${constant_list}\\n }), // constants list\\\"\\\"\\\"\\n)\\n\\nCONSTANTS_LIST_EMPTY = \\\"\\\"\\\"std::vector(), // constants list\\\"\\\"\\\"\\n\\nONE_TYPE = CodeTemplate(\\\"\\\"\\\"c10::parseType(\\\"${type_str}\\\"),\\\"\\\"\\\")\\n\\nTYPE_LIST = CodeTemplate(\\n \\\"\\\"\\\"std::vector({\\n ${type_list}\\n }), // types list\\\"\\\"\\\"\\n)\\n\\nTYPE_LIST_EMPTY = \\\"\\\"\\\"std::vector(), // types list\\\"\\\"\\\"\\n\\nONE_OPERATOTR_STRING = CodeTemplate(\\n \\\"\\\"\\\"\\n OperatorString({\\\"${operator_name}\\\", \\\"${overload_name}\\\", ${num_of_args}}),\\\"\\\"\\\"\\n)\\n\\nOPERATOR_STRING_LIST = CodeTemplate(\\n \\\"\\\"\\\"\\n std::vector({\\n ${operator_string_list}\\n }), // operators list\\\"\\\"\\\"\\n)\\n\\nONE_UPGRADER_FUNCTION = CodeTemplate(\\n \\\"\\\"\\\"\\n mobile::Function::registerFunc(\\n \\\"${upgrader_name}\\\",\\n ${instruction_list},\\n ${constant_list},\\n ${type_list},\\n ${register_size}\\n )\\\"\\\"\\\"\\n)\\n\\nONE_UPGRADER_SRC = CodeTemplate(\\n \\\"\\\"\\\"\\n ByteCodeFunctionWithOperator({\\n ${bytecode_function},\\n ${operator_string_list}\\n }),\\\"\\\"\\\"\\n)\\n\\n\\nONE_UPGRADER_IN_VERSION_MAP = CodeTemplate(\\n \\\"\\\"\\\"Upgrader({${upgrader_min_version}, ${upgrader_max_version}, \\\"${upgrader_name}\\\", ${bytecode_func_index}})\\\"\\\"\\\"\\n) # noqa: E501\\n\\nONE_OPERATOR_IN_VERSION_MAP = CodeTemplate(\\n \\\"\\\"\\\"\\n {std::string(\\\"${operator_name}\\\"),\\n std::vector({\\n ${upgrader_list_in_version_map}\\n })},\\\"\\\"\\\"\\n)\\n\\n\\nOPERATOR_VERSION_MAP = CodeTemplate(\\n \\\"\\\"\\\"\\nconst std::unordered_map>\\ngetOperatorVersionMapForMobile() {\\n static std::unordered_map>\\n operatorVersionMapForMobile({\\n ${operator_list_in_version_map}\\n });\\n return operatorVersionMapForMobile;\\n}\\n\\\"\\\"\\\"\\n)\\n\\n\\nUPGRADER_CPP_SRC = CodeTemplate(\\n MOBILE_UPGRADERS_HEADER_DESCRIPTION\\n + \\\"\\\"\\\"\\n#include \\n#include \\n\\nnamespace c10 {\\nTypePtr parseType(const std::string& pythonStr);\\n} // namespace c10\\n\\nnamespace torch {\\nnamespace jit {\\n\\n// clang-format off\\n\\n// From operator_versions_map\\n${operator_version_map}\\n\\nconst std::vector& getUpgraderBytecodeList() {\\n auto generate_upgrader_bytecode_list = []() {\\n std::vector upgrader_function_list({\\n ${upgrader_bytecode}\\n });\\n for (const auto& upgrader_function : upgrader_function_list) {\\n for (const auto& op : upgrader_function.operators) {\\n upgrader_function.function.append_operator(\\n op.name,\\n op.overload_name,\\n op.num_specified_args);\\n }\\n }\\n return upgrader_function_list;\\n };\\n static std::vector upgraderBytecodeList =\\n generate_upgrader_bytecode_list();\\n return upgraderBytecodeList;\\n}\\n\\n// clang-format on\\n\\n} // namespace jit\\n} // namespace torch\\n\\\"\\\"\\\"\\n)\\n\\nUPGRADER_MOBILE_FILE_NAME = \\\"upgrader_mobile.cpp\\\"\\n\\nUPGRADER_ELEMENT = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\nUpgrader({${min_version}, ${max_version}, ${operator_name}, ${index}}),\\n\\\"\\\"\\\"\\n)\\n\\nPER_OPERATOR_UPGRADER_LIST = CodeTemplate(\\n \\\"\\\"\\\"\\\\\\n{\\n std::string(${operator_name}),\\n std::vector({${upgrader_list}});\\n}\\n\\\"\\\"\\\"\\n)\\n\\n\\ndef construct_instruction(instruction_list_from_yaml: list[Any]) -> str:\\n instruction_list_part = []\\n for instruction in instruction_list_from_yaml:\\n instruction_list_part.append(\\n ONE_INSTRUCTION.substitute(\\n operator_name=instruction[0],\\n X=instruction[1],\\n N=instruction[2],\\n )\\n )\\n return INSTRUCTION_LIST.substitute(\\n instruction_list=\\\"\\\".join(instruction_list_part).lstrip(\\\"\\\\n\\\")\\n )\\n\\n\\ndef construct_constants(constants_list_from_yaml: list[Any]) -> str:\\n constants_list_part = []\\n for constant_from_yaml in constants_list_from_yaml:\\n convert_constant = None\\n if isinstance(constant_from_yaml, str):\\n # Add quotes if it's string\\n convert_constant = f'\\\"{constant_from_yaml}\\\"'\\n elif isinstance(constant_from_yaml, bool):\\n convert_constant = \\\"true\\\" if constant_from_yaml else \\\"false\\\"\\n elif constant_from_yaml is None:\\n convert_constant = \\\"\\\"\\n elif isinstance(constant_from_yaml, int):\\n convert_constant = str(constant_from_yaml)\\n else:\\n raise ValueError(\\n f\\\"The type of {constant_from_yaml} is {type(constant_from_yaml)}. \\\"\\n \\\"Please add change in construct_constants function in gen_mobile_upgraders.py.\\\"\\n )\\n constants_list_part.append(ONE_CONSTANT.substitute(constant=convert_constant))\\n if len(constants_list_part) == 0:\\n return CONSTANTS_LIST_EMPTY\\n return CONSTANT_LIST.substitute(\\n constant_list=\\\"\\\".join(constants_list_part).lstrip(\\\"\\\\n\\\")\\n )\\n\\n\\ndef construct_operators(operator_list_from_yaml: list[Any]) -> str:\\n operator_list_part = []\\n for operator in operator_list_from_yaml:\\n operator_list_part.append(\\n ONE_OPERATOTR_STRING.substitute(\\n operator_name=operator[0],\\n overload_name=operator[1],\\n num_of_args=operator[2],\\n )\\n )\\n return OPERATOR_STRING_LIST.substitute(\\n operator_string_list=\\\"\\\".join(operator_list_part).lstrip(\\\"\\\\n\\\")\\n )\\n\\n\\ndef construct_types(types_tr_list_from_yaml: list[Any]) -> str:\\n types_tr_list_part = []\\n for types_tr in types_tr_list_from_yaml:\\n types_tr_list_part.append(ONE_TYPE.substitute(type_str=types_tr))\\n if len(types_tr_list_part) == 0:\\n return TYPE_LIST_EMPTY\\n return TYPE_LIST.substitute(type_list=\\\"\\\".join(types_tr_list_part).lstrip(\\\"\\\\n\\\"))\\n\\n\\ndef construct_register_size(register_size_from_yaml: int) -> str:\\n if not isinstance(register_size_from_yaml, int):\\n raise ValueError(\\n f\\\"Input register size is {register_size_from_yaml} and\\\"\\n \\\"it's type is {type(register_size_from_yaml)}. An int type is expected.\\\"\\n )\\n return str(register_size_from_yaml)\\n\\n\\ndef construct_version_maps(\\n upgrader_bytecode_function_to_index_map: dict[str, Any]\\n) -> str:\\n version_map = torch._C._get_operator_version_map()\\n sorted_version_map_ = sorted(version_map.items(), key=itemgetter(0)) # type: ignore[no-any-return]\\n sorted_version_map = dict(sorted_version_map_)\\n\\n operator_list_in_version_map_part = []\\n for op_name in sorted_version_map:\\n upgraders_in_version_map_part = []\\n # TODO: remove the skip after these two operators schemas are fixed\\n if op_name in EXCLUDED_OP_SET:\\n continue\\n upgrader_ranges = torch._C._get_upgrader_ranges(op_name)\\n upgrader_entries = sorted_version_map[op_name]\\n assert len(upgrader_ranges) == len(upgrader_entries)\\n for idx, upgrader_entry in enumerate(upgrader_entries):\\n upgrader_name = upgrader_entry.upgrader_name\\n bytecode_function_index = upgrader_bytecode_function_to_index_map[\\n upgrader_name\\n ]\\n upgraders_in_version_map_part.append(\\n ONE_UPGRADER_IN_VERSION_MAP.substitute(\\n upgrader_min_version=upgrader_ranges[idx].min_version,\\n upgrader_max_version=upgrader_ranges[idx].max_version,\\n upgrader_name=upgrader_name,\\n bytecode_func_index=bytecode_function_index,\\n )\\n )\\n operator_list_in_version_map_part.append(\\n ONE_OPERATOR_IN_VERSION_MAP.substitute(\\n operator_name=op_name,\\n upgrader_list_in_version_map=\\\"\\\".join(upgraders_in_version_map_part),\\n )\\n )\\n return OPERATOR_VERSION_MAP.substitute(\\n operator_list_in_version_map=\\\"\\\".join(operator_list_in_version_map_part).lstrip(\\n \\\"\\\\n\\\"\\n )\\n )\\n\\n\\ndef get_upgrader_bytecode_function_to_index_map(\\n upgrader_dict: list[dict[str, Any]]\\n) -> dict[str, Any]:\\n upgrader_bytecode_function_to_index_map = {}\\n index = 0\\n for upgrader_bytecode in upgrader_dict:\\n for upgrader_name in upgrader_bytecode.keys():\\n if upgrader_name in EXCLUE_UPGRADER_SET:\\n continue\\n upgrader_bytecode_function_to_index_map[upgrader_name] = index\\n index += 1\\n return upgrader_bytecode_function_to_index_map\\n\\n\\ndef write_cpp(cpp_path: str, upgrader_dict: list[dict[str, Any]]) -> None:\\n body_parts = []\\n upgrader_bytecode_function_to_index_map = (\\n get_upgrader_bytecode_function_to_index_map(upgrader_dict)\\n )\\n version_map_src = construct_version_maps(upgrader_bytecode_function_to_index_map)\\n all_upgrader_src_string = []\\n for upgrader_bytecode in upgrader_dict:\\n for upgrader_name, bytecode in upgrader_bytecode.items():\\n # TODO: remove the skip after these two operators schemas are fixed\\n if upgrader_name in EXCLUE_UPGRADER_SET:\\n continue\\n instruction_list_str = \\\"\\\"\\n constant_list_str = \\\"\\\"\\n type_list_str = \\\"\\\"\\n register_size_str = \\\"\\\"\\n operator_list_str = \\\"\\\"\\n for table_name, contents in bytecode.items():\\n element = ByteCode[table_name]\\n body_string = \\\"\\\"\\n if element is ByteCode.instructions:\\n instruction_list_str = construct_instruction(contents)\\n elif element is ByteCode.constants:\\n constant_list_str = construct_constants(contents)\\n elif element is ByteCode.operators:\\n operator_list_str = construct_operators(contents)\\n elif element is ByteCode.types:\\n type_list_str = construct_types(contents)\\n elif element is ByteCode.register_size:\\n register_size_str = construct_register_size(contents)\\n\\n one_upgrader_function_string = ONE_UPGRADER_FUNCTION.substitute(\\n upgrader_name=upgrader_name,\\n instruction_list=instruction_list_str,\\n constant_list=constant_list_str,\\n type_list=type_list_str,\\n register_size=register_size_str,\\n )\\n one_upgrader_src_string = ONE_UPGRADER_SRC.substitute(\\n bytecode_function=one_upgrader_function_string.lstrip(\\\"\\\\n\\\"),\\n operator_string_list=operator_list_str.lstrip(\\\"\\\\n\\\"),\\n )\\n all_upgrader_src_string.append(one_upgrader_src_string)\\n\\n upgrader_file_content = UPGRADER_CPP_SRC.substitute(\\n operator_version_map=version_map_src,\\n upgrader_bytecode=\\\"\\\".join(all_upgrader_src_string).lstrip(\\\"\\\\n\\\"),\\n )\\n body_parts.append(upgrader_file_content)\\n print(\\\"writing file to : \\\", cpp_path + \\\"/\\\" + UPGRADER_MOBILE_FILE_NAME)\\n with open(os.path.join(cpp_path, UPGRADER_MOBILE_FILE_NAME), \\\"wb\\\") as out_file:\\n final_output = \\\"\\\".join(body_parts)\\n out_file.write(upgrader_file_content.encode(\\\"utf-8\\\"))\\n\\n\\ndef sort_upgrader(upgrader_list: list[dict[str, Any]]) -> list[dict[str, Any]]:\\n sorted_upgrader_list = sorted(\\n upgrader_list, key=lambda one_upgrader: next(iter(one_upgrader))\\n )\\n return sorted_upgrader_list\\n\\n\\ndef main() -> None:\\n upgrader_list = generate_upgraders_bytecode()\\n sorted_upgrader_list = sort_upgrader(upgrader_list)\\n for up in sorted_upgrader_list:\\n print(\\\"after sort upgrader : \\\", next(iter(up)))\\n\\n pytorch_dir = Path(__file__).resolve().parents[2]\\n upgrader_path = pytorch_dir / \\\"torch\\\" / \\\"csrc\\\" / \\\"jit\\\" / \\\"mobile\\\"\\n write_cpp(str(upgrader_path), sorted_upgrader_list)\\n\\n\\nif __name__ == \\\"__main__\\\":\\n main()\\n\\n\\nMOBILE_UPGRADERS_HEADER_DESCRIPTION = \\\"\\\"\\\"/**\\n * @generated\\n * This is an auto-generated file. Please do not modify it by hand.\\n * To re-generate, please run:\\n * cd ~/pytorch && python torchgen/operator_versions/gen_mobile_upgraders.py\\n */\\n\\\"\\\"\\\"\\n\\n\\n\\n\\nfrom __future__ import annotations\\n\\nimport json\\nimport logging\\nimport math\\nfrom typing import Sequence\\n\\nimport torchgen.api.cpp as cpp\\nfrom torchgen.context import native_function_manager\\nfrom torchgen.model import (\\n Argument,\\n BackendIndex,\\n BaseTy,\\n BaseType,\\n FunctionSchema,\\n NativeFunctionsGroup,\\n NativeFunctionsViewGroup,\\n OptionalType,\\n SelfArgument,\\n TensorOptionsArguments,\\n Type,\\n)\\nfrom torchgen.static_runtime import config\\n\\n\\nlogger: logging.Logger = logging.getLogger()\\n\\n\\ndef has_alias(\\n arguments: Sequence[Argument | SelfArgument | TensorOptionsArguments],\\n) -> bool:\\n for arg in arguments:\\n annotation = getattr(arg, \\\"annotation\\\", None)\\n if not annotation:\\n continue\\n alias_set = getattr(annotation, \\\"alias_set\\\", ())\\n if alias_set:\\n return True\\n return False\\n\\n\\nBLOCKED_OPS = frozenset(\\n (\\n # non cpu ops\\n \\\"sparse_sampled_addmm\\\",\\n \\\"hspmm\\\",\\n \\\"linalg_svdvals\\\",\\n # sparse ops\\n \\\"sspaddmm\\\",\\n \\\"coalesce\\\",\\n \\\"_indices\\\",\\n \\\"indices\\\",\\n \\\"_values\\\",\\n \\\"values\\\",\\n \\\"crow_indices\\\",\\n \\\"col_indices\\\",\\n # deprecated ops\\n \\\"floor_divide\\\",\\n \\\"ger\\\",\\n # buggy ops\\n \\\"conj_physical\\\", # P495807361\\n \\\"binary_cross_entropy\\\", # P496394764\\n \\\"arccosh\\\",\\n # uncommon ops\\n \\\"cholesky\\\",\\n \\\"lu_solve\\\",\\n \\\"linalg_cholesky\\\",\\n \\\"linalg_householder_product\\\",\\n \\\"linalg_ldl_solve\\\",\\n \\\"_compute_linear_combination\\\",\\n # training related ops\\n \\\"_make_dual\\\",\\n # cannot call directly\\n \\\"_fw_primal\\\",\\n # no documentation\\n \\\"_index_reduce\\\",\\n # TODO: these ones got added recently and need manual inspection\\n \\\"_new_zeros_with_same_feature_meta\\\",\\n \\\"_conj_physical\\\",\\n \\\"binary_cross_entropy_with_logits\\\",\\n \\\"bincount\\\",\\n \\\"conv_tbc\\\",\\n \\\"copy\\\",\\n \\\"_copy_from\\\",\\n \\\"_copy_from_and_resize\\\",\\n \\\"count_nonzero\\\",\\n \\\"cudnn_affine_grid_generator\\\",\\n \\\"cudnn_affine_grid_generator_backward\\\",\\n \\\"cudnn_grid_sampler\\\",\\n \\\"diag_embed\\\",\\n \\\"embedding\\\",\\n \\\"embedding_dense_backward\\\",\\n \\\"_embedding_bag_dense_backward\\\",\\n \\\"_embedding_bag_per_sample_weights_backward\\\",\\n \\\"grid_sampler_2d\\\",\\n \\\"_grid_sampler_2d_cpu_fallback\\\",\\n \\\"grid_sampler_3d\\\",\\n \\\"isnan\\\",\\n \\\"mkldnn_linear\\\",\\n \\\"median\\\",\\n \\\"nanmedian\\\",\\n \\\"_sparse_sparse_matmul\\\",\\n \\\"batch_norm_backward_elemt\\\",\\n \\\"_euclidean_dist\\\",\\n \\\"pixel_shuffle\\\",\\n \\\"pixel_unshuffle\\\",\\n \\\"channel_shuffle\\\",\\n \\\"_reshape_nested_backward\\\",\\n \\\"relu\\\",\\n \\\"prelu\\\",\\n \\\"celu\\\",\\n \\\"slice_scatter\\\",\\n \\\"select_scatter\\\",\\n \\\"diagonal_scatter\\\",\\n \\\"sum\\\",\\n \\\"_mkldnn_transpose\\\",\\n \\\"_nested_tensor_from_mask\\\",\\n \\\"_nested_from_padded\\\",\\n \\\"_nested_tensor_size\\\",\\n \\\"_nested_from_padded_and_nested_example\\\",\\n \\\"_standard_gamma_grad\\\",\\n \\\"_dirichlet_grad\\\",\\n \\\"native_norm\\\",\\n \\\"_sparse_softmax\\\",\\n \\\"_sparse_softmax_backward_data\\\",\\n \\\"_sparse_log_softmax\\\",\\n \\\"_sparse_log_softmax_backward_data\\\",\\n \\\"zero\\\",\\n \\\"_sparse_addmm\\\",\\n \\\"sparse_mask\\\",\\n \\\"_sparse_mask_projection\\\",\\n \\\"_to_dense\\\",\\n \\\"_coalesce\\\",\\n \\\"_coalesced\\\",\\n \\\"copy_sparse_to_sparse\\\",\\n \\\"to_sparse\\\",\\n \\\"to_sparse_csr\\\",\\n \\\"to_sparse_csc\\\",\\n \\\"to_mkldnn\\\",\\n \\\"quantize_per_tensor_dynamic\\\",\\n \\\"quantize_per_channel\\\",\\n \\\"q_per_channel_scales\\\",\\n \\\"q_per_channel_zero_points\\\",\\n \\\"int_repr\\\",\\n \\\"_make_per_channel_quantized_tensor\\\",\\n \\\"set\\\",\\n \\\"lift\\\",\\n \\\"lift_fresh\\\",\\n \\\"lift_fresh_copy\\\",\\n \\\"masked_scatter\\\",\\n \\\"_masked_softmax\\\",\\n \\\"_masked_softmax_backward\\\",\\n \\\"put\\\",\\n \\\"index_reduce\\\",\\n \\\"trace\\\",\\n \\\"_cholesky_solve_helper\\\",\\n \\\"dist\\\",\\n \\\"max\\\",\\n \\\"_torch_cuda_cu_linker_symbol_op\\\",\\n \\\"glu_jvp\\\",\\n \\\"glu_backward_jvp\\\",\\n \\\"hardswish_backward\\\",\\n \\\"rrelu_with_noise_backward\\\",\\n \\\"mkldnn_adaptive_avg_pool2d_backward\\\",\\n \\\"_adaptive_avg_pool2d_backward\\\",\\n \\\"_adaptive_avg_pool3d_backward\\\",\\n \\\"isinf\\\",\\n \\\"linalg_lu_solve\\\",\\n \\\"linalg_vecdot\\\",\\n \\\"linalg_matrix_exp\\\",\\n \\\"linalg_eigvalsh\\\",\\n \\\"_test_warn_in_autograd\\\",\\n \\\"_test_autograd_multiple_dispatch_view\\\",\\n \\\"_test_autograd_multiple_dispatch_view_copy\\\",\\n \\\"_segment_reduce\\\",\\n \\\"_segment_reduce_backward\\\",\\n \\\"_fw_primal_copy\\\",\\n \\\"_make_dual_copy\\\",\\n \\\"view_as_real_copy\\\",\\n \\\"view_as_complex_copy\\\",\\n \\\"_conj_copy\\\",\\n \\\"_neg_view_copy\\\",\\n \\\"diagonal_copy\\\",\\n \\\"detach_copy\\\",\\n \\\"squeeze_copy\\\",\\n \\\"t_copy\\\",\\n \\\"unsqueeze_copy\\\",\\n \\\"_indices_copy\\\",\\n \\\"_values_copy\\\",\\n \\\"indices_copy\\\",\\n \\\"values_copy\\\",\\n \\\"crow_indices_copy\\\",\\n \\\"col_indices_copy\\\",\\n \\\"ccol_indices\\\",\\n \\\"ccol_indices_copy\\\",\\n \\\"row_indices\\\",\\n \\\"row_indices_copy\\\",\\n \\\"unfold_copy\\\",\\n \\\"alias_copy\\\",\\n \\\"_triton_multi_head_attention\\\",\\n \\\"special_airy_ai\\\",\\n \\\"special_bessel_j0\\\",\\n \\\"special_bessel_j1\\\",\\n \\\"special_bessel_y0\\\",\\n \\\"special_bessel_y1\\\",\\n \\\"special_chebyshev_polynomial_t\\\",\\n \\\"special_chebyshev_polynomial_u\\\",\\n \\\"special_chebyshev_polynomial_v\\\",\\n \\\"special_chebyshev_polynomial_w\\\",\\n \\\"special_hermite_polynomial_h\\\",\\n \\\"special_hermite_polynomial_he\\\",\\n \\\"special_laguerre_polynomial_l\\\",\\n \\\"special_legendre_polynomial_p\\\",\\n \\\"special_modified_bessel_i0\\\",\\n \\\"special_modified_bessel_i1\\\",\\n \\\"special_modified_bessel_k0\\\",\\n \\\"special_modified_bessel_k1\\\",\\n \\\"special_scaled_modified_bessel_k0\\\",\\n \\\"special_scaled_modified_bessel_k1\\\",\\n \\\"special_shifted_chebyshev_polynomial_t\\\",\\n \\\"special_shifted_chebyshev_polynomial_u\\\",\\n \\\"special_shifted_chebyshev_polynomial_v\\\",\\n \\\"special_shifted_chebyshev_polynomial_w\\\",\\n \\\"special_spherical_bessel_j0\\\",\\n \\\"_foobar\\\",\\n \\\"_nested_tensor_strides\\\",\\n \\\"_nested_tensor_storage_offsets\\\",\\n \\\"_nested_get_values\\\", # no CPU backend\\n \\\"_nested_get_values_copy\\\", # no CPU backend\\n \\\"_nested_view_from_jagged\\\", # testing needs to be patched\\n \\\"_nested_view_from_jagged_copy\\\", # testing needs to be patched\\n \\\"_nested_view_from_buffer\\\", # testing needs to be patched\\n \\\"_nested_view_from_buffer_copy\\\", # testing needs to be patched\\n \\\"_int_mm\\\", # testing needs to be patched\\n \\\"_to_sparse_csc\\\", # testing needs to be patched\\n \\\"_to_sparse_csr\\\", # testing needs to be patched\\n \\\"segment_reduce\\\", # testing needs to be patched\\n )\\n)\\n\\n\\ndef is_supported(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> bool:\\n base_op_name = \\\"\\\"\\n func = None\\n if isinstance(g, NativeFunctionsViewGroup):\\n base_op_name = g.view.root_name\\n func = g.view.func\\n else:\\n base_op_name = g.out.func.name.name.base\\n func = g.out.func\\n if config.is_hand_written(g):\\n logger.info(\\\"HAND WRITTEN: %s\\\", base_op_name)\\n return False\\n if base_op_name in BLOCKED_OPS:\\n logger.info(\\\"BLOCKED: %s\\\", base_op_name)\\n return False\\n for arg in func.schema_order_arguments():\\n maybe_method = ivalue_type_conversion_method(arg.type)\\n if not maybe_method:\\n # Type converting is unsupported yet.\\n logger.info(\\\"NOT SUPPORTED TYPE CONVERTING: %s\\\", func)\\n return False\\n\\n if isinstance(g, NativeFunctionsViewGroup):\\n # TODO: stop doing type tests by converting to C++ and then testing\\n # the string, just test the dang thing directly\\n if \\\"at::Tensor\\\" != cpp.returns_type(func.returns, symint=False).cpp_type():\\n # Returns a non-Tensor value.\\n logger.info(\\\"NON-TENSOR RET TYPE: %s\\\", str(func))\\n return False\\n return True\\n\\n # For out variant ops, we need to check the arguments of its functional func.\\n for arg in g.functional.func.schema_order_arguments():\\n maybe_method = ivalue_type_conversion_method(arg.type)\\n if not maybe_method:\\n # Type converting is unsupported yet.\\n logger.info(\\\"NOT SUPPORTED TYPE CONVERTING: %s\\\", g.functional.func)\\n return False\\n\\n if not g.structured:\\n # In case of unstructured op, we check if it has out variant implementation.\\n # The out variant implementation satisfies the minimum requirement that it has the output tensor as the last\\n # parameter.\\n if (\\n not hasattr(g, \\\"out\\\")\\n or not str(func).endswith(\\\"Tensor(a!) out) -> Tensor(a!)\\\")\\n or not str(func.name).endswith(\\\".out\\\")\\n ):\\n return False\\n # TODO: stop type testing by converting to C++\\n if \\\"at::Tensor &\\\" != cpp.returns_type(func.returns, symint=False).cpp_type():\\n logger.info(\\\"NON_TENSOR RET TYPE: %s\\\", func)\\n return False\\n if has_alias(func.arguments.non_out):\\n # This op may create an alias of inputs.\\n logger.info(\\\"INPUTS ALIAS: %s\\\", base_op_name)\\n return False\\n return True\\n\\n\\ndef ivalue_type_conversion_method(\\n arg_type: BaseType | OptionalType | Type,\\n) -> tuple[bool, str] | None:\\n \\\"\\\"\\\"\\n Return the method call expression of `c10::ivalue' to convert its contained value to\\n the expected value of `arg_type` type. For example, for `arg_type` == BaseTy.Tensor,\\n this function returns \\\".toTensor()\\\", so that it can be appended to the ivalue's\\n variable name to get the value of the expected type.\\n \\\"\\\"\\\"\\n type_conversion_methods = {\\n BaseTy.Tensor: ((True, \\\"toTensor()\\\"), (False, \\\"toOptional()\\\")),\\n BaseTy.int: ((False, \\\"toInt()\\\"), (False, \\\"toOptional()\\\")),\\n BaseTy.bool: ((False, \\\"toBool()\\\"), (False, \\\"toOptional()\\\")),\\n BaseTy.Scalar: ((False, \\\"toScalar()\\\"), (False, \\\"toOptional()\\\")),\\n BaseTy.ScalarType: (\\n (False, \\\"toScalarType()\\\"),\\n (False, \\\"toOptional()\\\"),\\n ),\\n BaseTy.str: (\\n (False, \\\"toStringView()\\\"),\\n (False, \\\"toOptional()\\\"),\\n ),\\n }\\n\\n base_ty_object = None\\n if isinstance(arg_type, BaseType):\\n base_ty_object = arg_type.name\\n elif isinstance(arg_type, OptionalType):\\n if not isinstance(arg_type.elem, BaseType):\\n # ListType is currently unsupported.\\n return None\\n base_ty_object = arg_type.elem.name\\n else:\\n return None\\n\\n if base_ty_object not in type_conversion_methods:\\n return None\\n methods = type_conversion_methods[base_ty_object]\\n if isinstance(arg_type, BaseType):\\n return methods[0]\\n return methods[1]\\n\\n\\nshould_use_int_tensor_ops_ = frozenset(\\n (\\n \\\"bitwise_not\\\",\\n \\\"bitwise_and\\\",\\n \\\"bitwise_or\\\",\\n \\\"bitwise_xor\\\",\\n \\\"bitwise_left_shift\\\",\\n \\\"bitwise_right_shift\\\",\\n \\\"gcd\\\",\\n \\\"lcm\\\",\\n \\\"scatter\\\",\\n \\\"gather\\\",\\n \\\"_convert_indices_from_coo_to_csr\\\",\\n \\\"_convert_indices_from_csr_to_coo\\\",\\n )\\n)\\nshould_use_complex_tensor_ops_ = frozenset((\\\"view_as_real\\\", \\\"imag\\\", \\\"_conj\\\"))\\n\\n\\ndef should_use_int_tensor(op_name: str) -> bool:\\n return op_name in should_use_int_tensor_ops_\\n\\n\\ndef should_use_complex_tensor(op_name: str) -> bool:\\n return op_name in should_use_complex_tensor_ops_\\n\\n\\ntest_tensor_dim_ops_1_ = frozenset(\\n (\\n \\\"addmv\\\",\\n \\\"index_add\\\",\\n \\\"_convert_indices_from_coo_to_csr\\\",\\n \\\"_convert_indices_from_csr_to_coo\\\",\\n \\\"nll_loss_backward\\\",\\n \\\"dot\\\",\\n \\\"vdot\\\",\\n \\\"outer\\\",\\n \\\"ger\\\",\\n )\\n)\\ntest_tensor_dim_ops_2_ = frozenset(\\n (\\\"addmm\\\", \\\"mm\\\", \\\"nuclear_norm\\\", \\\"diag\\\", \\\"_addmm_activation\\\", \\\"matrix_H\\\", \\\"t\\\")\\n)\\n\\n\\ndef test_tensor_dim(op_name: str) -> int:\\n if op_name in test_tensor_dim_ops_1_:\\n return 1\\n if op_name in test_tensor_dim_ops_2_:\\n return 2\\n return 3\\n\\n\\ntest_tensor_shapes_string = '{\\\"view_as_complex\\\": \\\"{2, 2}\\\"}'\\ntest_tensor_shape_json: dict[str, str] = json.loads(test_tensor_shapes_string)\\n\\n\\ndef test_tensor_shape(op_name: str) -> str:\\n if op_name in test_tensor_shape_json:\\n return test_tensor_shape_json[op_name]\\n else:\\n return \\\"\\\"\\n\\n\\ndef test_value_expression(\\n arg_type: BaseType | OptionalType | Type, index: int, op_name: str\\n) -> str:\\n tensor_size_ex = test_tensor_shape(op_name)\\n if tensor_size_ex == \\\"\\\":\\n num_tensors = 16 if index == 0 else 64\\n num_dim = test_tensor_dim(op_name)\\n size_per_dim = math.ceil(num_tensors / float(num_dim))\\n size_per_dim += size_per_dim % 2\\n tensor_size_ex = \\\"{{{}}}\\\".format(\\\",\\\".join([f\\\"{size_per_dim}\\\"] * num_dim))\\n if should_use_int_tensor(op_name):\\n tensor_expression = f\\\"at::randint(1, 100, {tensor_size_ex}, at::kInt)\\\"\\n elif should_use_complex_tensor(op_name):\\n tensor_expression = f\\\"at::randn({tensor_size_ex}, at::kComplexFloat)\\\"\\n else:\\n tensor_expression = f\\\"at::rand({tensor_size_ex})\\\"\\n\\n value_expressions = {\\n BaseTy.Tensor: tensor_expression,\\n BaseTy.int: \\\"1\\\",\\n BaseTy.bool: \\\"false\\\",\\n BaseTy.Scalar: \\\"2\\\",\\n BaseTy.ScalarType: \\\"at::ScalarType::Float\\\",\\n BaseTy.str: '\\\"floor\\\"',\\n }\\n\\n base_ty_object = None\\n if isinstance(arg_type, BaseType):\\n base_ty_object = arg_type.name\\n else:\\n assert isinstance(arg_type, OptionalType) and isinstance(\\n arg_type.elem, BaseType\\n )\\n base_ty_object = arg_type.elem.name\\n assert base_ty_object in value_expressions, \\\"not expected type\\\"\\n value_expression = value_expressions[base_ty_object]\\n return value_expression\\n\\n\\ndef generate_test_value_definitions(schema: FunctionSchema, index: int) -> str:\\n assert not schema.is_out_fn()\\n schema_name = schema.name.name.base\\n arg_map = {}\\n for arg in schema.schema_order_arguments():\\n test_value_exp = test_value_expression(arg.type, index, schema_name)\\n arg_map[arg.name] = test_value_exp\\n config.override_test_values(arg_map, schema_name, index)\\n arg_populations = []\\n for arg_name, arg_value in arg_map.items():\\n arg_populations.append(f\\\"auto {arg_name}{index} = {arg_value}\\\")\\n return \\\";\\\\n \\\".join(arg_populations) + \\\";\\\"\\n\\n\\ndef generate_test_value_names(schema: FunctionSchema, index: int) -> str:\\n assert not schema.is_out_fn()\\n return \\\",\\\".join(f\\\"{arg.name}{index}\\\" for arg in schema.schema_order_arguments())\\n\\n\\ngenerate_test_ir_arguments_base_ty_to_type_str_ = {\\n BaseTy.Tensor: \\\"Tensor\\\",\\n BaseTy.int: \\\"int\\\",\\n BaseTy.float: \\\"float\\\",\\n BaseTy.str: \\\"str\\\",\\n BaseTy.Scalar: \\\"int\\\",\\n BaseTy.ScalarType: \\\"int\\\",\\n BaseTy.bool: \\\"bool\\\",\\n}\\n\\n\\ndef generate_test_ir_arguments(\\n schema: FunctionSchema,\\n) -> list[tuple[str, str | None]]:\\n def ir_argument(arg: Argument) -> tuple[str, str | None]:\\n t = arg.type\\n add_optional = False\\n if isinstance(t, OptionalType):\\n t = t.elem\\n add_optional = True\\n assert isinstance(t, BaseType)\\n type_str = None\\n if t.name in generate_test_ir_arguments_base_ty_to_type_str_:\\n type_str = generate_test_ir_arguments_base_ty_to_type_str_[t.name]\\n if type_str and add_optional:\\n type_str = f\\\"{type_str}?\\\"\\n return (\\\"%\\\" + arg.name, type_str)\\n\\n return [ir_argument(arg) for arg in schema.schema_order_arguments()]\\n\\n\\ndef generate_arg_extraction(schema: FunctionSchema) -> str:\\n arg_populations = []\\n for i, arg in enumerate(schema.schema_order_arguments()):\\n maybe_method = ivalue_type_conversion_method(arg.type)\\n assert maybe_method\\n is_reference, type_conversion_method = maybe_method\\n reference = \\\"&\\\" if is_reference else \\\"\\\"\\n arg_populations.append(\\n f\\\"const auto{reference} {arg.name} = p_node->Input({i}).{type_conversion_method}\\\"\\n )\\n return \\\";\\\\n \\\".join(arg_populations) + \\\";\\\"\\n\\n\\ndef get_kernel_name(g: NativeFunctionsGroup, backend_index: BackendIndex) -> str:\\n kernel = backend_index.get_kernel(g.functional)\\n if g.structured or kernel is None:\\n return cpp.name(g.functional.func)\\n return kernel.kernel\\n\\n\\ndef get_out_kernel_name(g: NativeFunctionsGroup, backend_index: BackendIndex) -> str:\\n kernel = backend_index.get_kernel(g.out)\\n if g.structured or kernel is None:\\n return cpp.name(g.out.func)\\n return kernel.kernel\\n\\n\\ndef generate_non_out_variant_call(\\n g: NativeFunctionsGroup, backend_index: BackendIndex\\n) -> str:\\n schema = g.functional.func\\n assert not schema.is_out_fn()\\n kernel_name = get_kernel_name(g, backend_index)\\n arg_names = (arg.name for arg in schema.schema_order_arguments())\\n namespace_name = \\\"cpu\\\" if g.structured else \\\"native\\\"\\n return f'at::{namespace_name}::{kernel_name}({\\\",\\\".join(arg_names)})'\\n\\n\\ndef generate_call_to_view_ops(\\n g: NativeFunctionsViewGroup, backend_index: BackendIndex\\n) -> str:\\n schema = g.view.func\\n kernel_name = cpp.name(schema)\\n kernel = backend_index.get_kernel(g.view)\\n if kernel:\\n kernel_name = kernel.kernel\\n arg_names = (arg.name for arg in schema.schema_order_arguments())\\n namespace_name = \\\"native\\\"\\n return f'at::{namespace_name}::{kernel_name}({\\\",\\\".join(arg_names)})'\\n\\n\\ndef generate_out_variant_call(\\n g: NativeFunctionsGroup, backend_index: BackendIndex\\n) -> str:\\n schema = g.out.func\\n assert schema.is_out_fn()\\n arg_names = []\\n kernel_name = get_out_kernel_name(g, backend_index)\\n if g.structured:\\n # structured op starts with the output tensor argument.\\n arg_names = [out_arg.name for out_arg in schema.arguments.out]\\n else:\\n arg_names = []\\n for arg in schema.arguments.non_out:\\n if isinstance(arg, SelfArgument):\\n arg_names.append(arg.argument.name)\\n else:\\n assert isinstance(arg, Argument)\\n arg_names.append(arg.name)\\n if not g.structured:\\n assert len(schema.arguments.out) == 1\\n arg_names.append(schema.arguments.out[0].name)\\n cpp_arg_names = \\\",\\\".join(arg_names)\\n namespace_name = \\\"cpu\\\" if g.structured else \\\"native\\\"\\n return f\\\"at::{namespace_name}::{kernel_name}({cpp_arg_names})\\\"\\n\\n\\nno_memory_resize_ops = frozenset(\\n (\\n \\\"isin.Scalar_Tensor\\\",\\n \\\"index_add\\\",\\n \\\"dot\\\",\\n \\\"vdot\\\",\\n \\\"nuclear_norm\\\",\\n \\\"histc\\\",\\n \\\"l1_loss\\\",\\n \\\"multi_margin_loss\\\",\\n \\\"multilabel_margin_loss\\\",\\n \\\"nll_loss\\\",\\n \\\"nll_loss2d\\\",\\n \\\"prod\\\",\\n )\\n)\\n\\n\\ndef should_check_resize(schema: FunctionSchema) -> bool:\\n schema_str = str(schema)\\n type_variant_op_name = schema_str[: schema_str.find(\\\"(\\\")]\\n return type_variant_op_name not in no_memory_resize_ops\\n\\n\\ndef op_name_from_group(g: NativeFunctionsGroup) -> str:\\n return g.functional.func.name.name.base\\n\\n\\nclass GenOpDispatcher:\\n def out_variant(\\n self, groups: Sequence[NativeFunctionsGroup], backend_index: BackendIndex\\n ) -> str:\\n if not groups:\\n return \\\"\\\"\\n generated_type_variants = []\\n for g in groups:\\n with native_function_manager(g):\\n assert is_supported(g)\\n assert isinstance(g, NativeFunctionsGroup)\\n generated_type_variant = self.out_variant_op_generator(g, backend_index)\\n generated_type_variants.append(generated_type_variant)\\n op_name = op_name_from_group(groups[0])\\n body = \\\"\\\\n\\\".join(generated_type_variants)\\n generated = f\\\"\\\"\\\"\\nREGISTER_OPERATOR_FUNCTOR(\\n aten::{op_name},\\n aten_{op_name},\\n [](Node* n) -> SROperator {{\\n {body}\\n LogAndDumpSchema(n);\\n return nullptr;\\n }});\\n\\\"\\\"\\\"\\n return generated\\n\\n def view(\\n self, groups: Sequence[NativeFunctionsViewGroup], backend_index: BackendIndex\\n ) -> str:\\n if not groups:\\n return \\\"\\\"\\n generated_type_variants = []\\n for g in groups:\\n with native_function_manager(g):\\n assert is_supported(g)\\n assert isinstance(g, NativeFunctionsViewGroup)\\n generated_type_variant = self.view_op_generator(g, backend_index)\\n generated_type_variants.append(generated_type_variant)\\n op_name = config.func_name_base_str(groups[0])\\n body = \\\"\\\\n\\\".join(generated_type_variants)\\n generated = f\\\"\\\"\\\"\\nREGISTER_NATIVE_OPERATOR_FUNCTOR(\\n aten::{op_name},\\n aten_{op_name},\\n [](Node* n) -> SROperator {{\\n {body}\\n LogAndDumpSchema(n);\\n return nullptr;\\n }});\\n\\\"\\\"\\\"\\n return generated\\n\\n def out_variant_op_generator(\\n self, g: NativeFunctionsGroup, backend_index: BackendIndex\\n ) -> str:\\n functional = g.functional\\n schema = str(functional.func)\\n populated_argument = generate_arg_extraction(g.functional.func)\\n functional_variant_call = generate_non_out_variant_call(g, backend_index)\\n assert len(g.out.func.arguments.out) == 1\\n out_variable_name = str(g.out.func.arguments.out[0].name)\\n out_variant_call = generate_out_variant_call(g, backend_index)\\n generated = f\\\"\\\"\\\"\\n if (n->matches(torch::schema(\\\"aten::{schema}\\\"))) {{\\n return [](ProcessedNode* p_node) {{\\n {populated_argument}\\n if (p_node->Output(0).isNone()) {{\\n p_node->Output(0) = {functional_variant_call};\\n return;\\n }}\\n auto& {out_variable_name} = p_node->Output(0).toTensor();\\n fastResizeToZero({out_variable_name});\\n {out_variant_call};\\n }};\\n }}\\\"\\\"\\\"\\n return generated\\n\\n def view_op_generator(\\n self, g: NativeFunctionsViewGroup, backend_index: BackendIndex\\n ) -> str:\\n schema = str(g.view.func)\\n populated_argument = generate_arg_extraction(g.view.func)\\n functional_variant_call = generate_call_to_view_ops(g, backend_index)\\n generated = f\\\"\\\"\\\"\\n if (n->matches(torch::schema(\\\"aten::{schema}\\\"))) {{\\n return [](ProcessedNode* p_node) {{\\n {populated_argument}\\n p_node->Output(0) = {functional_variant_call};\\n }};\\n }}\\\"\\\"\\\"\\n return generated\\n\\n\\nclass GenOpTestCase:\\n def out_variant(self, groups: Sequence[NativeFunctionsGroup]) -> str:\\n if not groups:\\n return \\\"\\\"\\n generated_type_variants = []\\n for g in groups:\\n with native_function_manager(g):\\n assert is_supported(g)\\n assert isinstance(g, NativeFunctionsGroup)\\n generated_type_variant = self.out_variant_op_test_case_generator(g)\\n generated_type_variants.append(generated_type_variant)\\n return \\\"\\\\n\\\".join(generated_type_variants)\\n\\n def view(self, groups: Sequence[NativeFunctionsViewGroup]) -> str:\\n if not groups:\\n return \\\"\\\"\\n generated_type_variants = []\\n for g in groups:\\n with native_function_manager(g):\\n assert is_supported(g)\\n assert isinstance(g, NativeFunctionsViewGroup)\\n generated_type_variant = self.view_op_test_case_generator(g)\\n generated_type_variants.append(generated_type_variant)\\n return \\\"\\\\n\\\".join(generated_type_variants)\\n\\n def out_variant_op_test_case_generator(self, g: NativeFunctionsGroup) -> str:\\n schema = g.functional.func\\n schema_str = str(schema)\\n assert schema_str.find(\\\"(\\\") > 0\\n type_variant_op_name = schema_str[: schema_str.find(\\\"(\\\")].replace(\\\".\\\", \\\"_\\\")\\n op_name = op_name_from_group(g)\\n assert type_variant_op_name.startswith(op_name)\\n\\n arg_types = generate_test_ir_arguments(schema)\\n arg_declarations = \\\", \\\".join(\\n (\\n arg_name if arg_type is None else f\\\"{arg_name}: {arg_type}\\\"\\n for arg_name, arg_type in arg_types\\n )\\n )\\n arg_names = \\\", \\\".join((arg_name for arg_name, _ in arg_types))\\n assert (\\n len(schema.returns) == 1\\n and isinstance(schema.returns[0].type, BaseType)\\n and schema.returns[0].type.name is BaseTy.Tensor\\n )\\n test_value_definitions = generate_test_value_definitions(schema, 0)\\n test_value_names = generate_test_value_names(schema, 0)\\n test_value_definitions2 = generate_test_value_definitions(schema, 1)\\n test_value_names2 = generate_test_value_names(schema, 1)\\n check_resize = \\\"true\\\" if should_check_resize(schema) else \\\"false\\\"\\n generated = f\\\"\\\"\\\"\\nTEST(StaticRuntime, autogen_{type_variant_op_name}) {{\\n const std::string script = R\\\"IR(\\n graph({arg_declarations}):\\n %bias: None = prim::Constant()\\n %ret = aten::{op_name}({arg_names})\\n %cloned = aten::clone(%ret, %bias)\\n return (%cloned)\\n )IR\\\";\\n\\n {test_value_definitions}\\n std::vector args{{{test_value_names}}};\\n testStaticRuntime(script, args, {{}}, /*use_allclose=*/false, /*use_equalnan=*/false, /*check_resize=*/{check_resize});\\n\\n {test_value_definitions2}\\n std::vector args2{{{test_value_names2}}};\\n testStaticRuntime(script, args, args2, /*use_allclose=*/false, /*use_equalnan=*/false, /*check_resize=*/{check_resize});\\n\\n}}\\n\\\"\\\"\\\"\\n return generated\\n\\n def view_op_test_case_generator(self, g: NativeFunctionsViewGroup) -> str:\\n schema = g.view.func\\n schema_str = str(schema)\\n assert schema_str.find(\\\"(\\\") > 0\\n type_variant_op_name = schema_str[: schema_str.find(\\\"(\\\")].replace(\\\".\\\", \\\"_\\\")\\n op_name = g.view.root_name\\n assert type_variant_op_name.startswith(op_name)\\n\\n arg_types = generate_test_ir_arguments(schema)\\n arg_declarations = \\\", \\\".join(\\n (\\n arg_name if arg_type is None else f\\\"{arg_name}: {arg_type}\\\"\\n for arg_name, arg_type in arg_types\\n )\\n )\\n arg_names = \\\", \\\".join((arg_name for arg_name, _ in arg_types))\\n assert (\\n len(schema.returns) == 1\\n and isinstance(schema.returns[0].type, BaseType)\\n and schema.returns[0].type.name is BaseTy.Tensor\\n )\\n test_value_definitions = generate_test_value_definitions(schema, 0)\\n test_value_names = generate_test_value_names(schema, 0)\\n generated = f\\\"\\\"\\\"\\nTEST(StaticRuntime, autogen_{type_variant_op_name}) {{\\n const std::string script = R\\\"IR(\\n graph({arg_declarations}):\\n %bias: None = prim::Constant()\\n %ret = aten::{op_name}({arg_names})\\n %cloned = aten::clone(%ret, %bias)\\n return (%cloned)\\n )IR\\\";\\n\\n {test_value_definitions}\\n std::vector args{{{test_value_names}}};\\n testStaticRuntime(script, args);\\n}}\\n\\\"\\\"\\\"\\n\\n return generated\\n\\n\\nfrom __future__ import annotations\\n\\nimport argparse\\nimport itertools\\nimport os\\nfrom typing import Sequence, TypeVar, Union\\n\\nfrom libfb.py.log import set_simple_logging # type: ignore[import]\\n\\nfrom torchgen import gen\\nfrom torchgen.context import native_function_manager\\nfrom torchgen.model import DispatchKey, NativeFunctionsGroup, NativeFunctionsViewGroup\\nfrom torchgen.static_runtime import config, generator\\n\\n\\n# Given a list of `grouped_native_functions` sorted by their op names, return a list of\\n# lists each of which groups ops that share the base name. For example, `mean` and\\n# `mean.dim` are grouped together by this function.\\n\\nNativeGroupT = TypeVar(\\n \\\"NativeGroupT\\\",\\n bound=Union[NativeFunctionsGroup, NativeFunctionsViewGroup],\\n)\\n\\n\\ndef group_functions_by_op_name(\\n grouped_native_functions: Sequence[NativeGroupT],\\n) -> Sequence[Sequence[NativeGroupT]]:\\n if not grouped_native_functions:\\n return []\\n groups = []\\n\\n def is_supported(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> bool:\\n with native_function_manager(g):\\n return generator.is_supported(g)\\n\\n eligible_ops = (g for g in grouped_native_functions if is_supported(g))\\n groups = [\\n list(group)\\n for k, group in (\\n itertools.groupby(\\n eligible_ops,\\n key=config.func_name_base_str,\\n )\\n )\\n ]\\n\\n return groups\\n\\n\\ndef clang_format(cpp_file_path: str) -> None:\\n import subprocess\\n\\n subprocess.check_call([\\\"clang-format\\\", \\\"-i\\\", cpp_file_path])\\n\\n\\ndef write_cpp(cpp_ops: Sequence[str], file_path: str) -> None:\\n code = \\\"\\\\n\\\".join(cpp_ops)\\n generated = f\\\"\\\"\\\"// @lint-ignore-every CLANGTIDY HOWTOEVEN\\n// AUTO-GENERATED FROM: torchgen/static_runtime/gen_static_runtime_ops.py\\n#include \\n\\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n#include \\n\\nnamespace torch {{\\nnamespace jit {{\\n\\n{code}\\n\\n}} // namespace jit\\n}} // namespace torch\\n\\\"\\\"\\\"\\n with open(file_path, \\\"w\\\") as f:\\n f.write(generated)\\n clang_format(file_path)\\n\\n\\ndef write_test_cpp(cpp_ops: Sequence[str], file_path: str) -> None:\\n code = \\\"\\\\n\\\".join(cpp_ops)\\n generated = f\\\"\\\"\\\"// @lint-ignore-every CLANGTIDY HOWTOEVEN\\n// AUTO-GENERATED FROM: torchgen/static_runtime/gen_static_runtime_ops.py\\n#include \\n#include \\n#include \\n\\n#include \\\"test_utils.h\\\"\\n\\nusing namespace caffe2;\\nusing namespace torch;\\nusing namespace torch::jit;\\nusing namespace torch::jit::test;\\nusing c10::IValue;\\n\\n{code}\\n\\n\\\"\\\"\\\"\\n with open(file_path, \\\"w\\\") as f:\\n f.write(generated)\\n clang_format(file_path)\\n\\n\\ndef main() -> None:\\n parser = argparse.ArgumentParser(description=\\\"Generate ATen source files\\\")\\n parser.add_argument(\\n \\\"-s\\\",\\n \\\"--source-path\\\",\\n help=\\\"path to source directory for ATen\\\",\\n default=\\\"caffe2/aten/src/ATen\\\",\\n )\\n parser.add_argument(\\n \\\"-p\\\",\\n \\\"--generated-ops-cpp-path\\\",\\n help=\\\"path to directory to generate op dispatcher .cpp file\\\",\\n default=\\\"caffe2/torch/csrc/jit/runtime/static/generated_ops.cpp\\\",\\n )\\n parser.add_argument(\\n \\\"-t\\\",\\n \\\"--generated-ops-test-cpp-path\\\",\\n help=\\\"path to directory to generate op dispatcher .cpp file\\\",\\n default=\\\"caffe2/benchmarks/static_runtime/test_generated_ops.cc\\\",\\n )\\n options = parser.parse_args()\\n native_yaml_path = os.path.join(options.source_path, \\\"native/native_functions.yaml\\\")\\n tags_yaml_path = os.path.join(options.source_path, \\\"native/tags.yaml\\\")\\n parsed_yaml = gen.parse_native_yaml(native_yaml_path, tags_yaml_path)\\n native_functions, backend_indices = (\\n parsed_yaml.native_functions,\\n parsed_yaml.backend_indices,\\n )\\n\\n op_generator = generator.GenOpDispatcher()\\n test_case_generator = generator.GenOpTestCase()\\n\\n native_functions_groups = [\\n g\\n for g in gen.get_grouped_native_functions(native_functions)\\n if isinstance(g, NativeFunctionsGroup)\\n ]\\n\\n supported_functions_groups = group_functions_by_op_name(native_functions_groups)\\n\\n out_variant_op_result = [\\n op_generator.out_variant(groups, backend_indices[DispatchKey.CPU])\\n for groups in supported_functions_groups\\n ]\\n out_variant_test_result = [\\n test_case_generator.out_variant(groups) for groups in supported_functions_groups\\n ]\\n\\n native_functions_view_groups = [\\n g\\n for g in gen.get_grouped_by_view_native_functions(native_functions)\\n if isinstance(g, NativeFunctionsViewGroup)\\n ]\\n\\n supported_functions_view_groups = group_functions_by_op_name(\\n native_functions_view_groups\\n )\\n\\n view_op_result = [\\n op_generator.view(groups, backend_indices[DispatchKey.CPU])\\n for groups in supported_functions_view_groups\\n ]\\n view_test_result = [\\n test_case_generator.view(groups) for groups in supported_functions_view_groups\\n ]\\n\\n op_result = out_variant_op_result + [\\\"\\\\n\\\\n\\\"] + view_op_result\\n test_result = out_variant_test_result + [\\\"\\\\n\\\\n\\\"] + view_test_result\\n\\n write_cpp(op_result, options.generated_ops_cpp_path)\\n write_test_cpp(test_result, options.generated_ops_test_cpp_path)\\n\\n print(\\n \\\"\\\\ntotal grouped native ops: %d\\\"\\n % len(gen.get_grouped_native_functions(native_functions))\\n )\\n\\n print(\\\"grouped native ops with out variant: %d\\\" % len(native_functions_groups))\\n supported_functions_num = sum(len(groups) for groups in supported_functions_groups)\\n print(\\\"generated functions groups with out variant: %d\\\" % supported_functions_num)\\n\\n print(\\\"\\\\nview grouped native ops: %d\\\" % len(native_functions_view_groups))\\n supported_view_functions_num = sum(\\n len(groups) for groups in supported_functions_view_groups\\n )\\n print(\\\"generated functions view groups: %d\\\" % supported_view_functions_num)\\n\\n print(\\n \\\"\\\\noverall generated : %d\\\"\\n % (supported_functions_num + supported_view_functions_num)\\n )\\n\\n\\nif __name__ == \\\"__main__\\\":\\n set_simple_logging(escape_newlines=False)\\n main()\\n\\n\\nfrom __future__ import annotations\\n\\nfrom torchgen.model import NativeFunctionsGroup, NativeFunctionsViewGroup\\n\\n\\ndef func_name_base_str(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> str:\\n if isinstance(g, NativeFunctionsGroup):\\n return str(g.functional.func.name.name.base)\\n else:\\n return str(g.view.root_name)\\n\\n\\nis_hand_written_ops_ = frozenset(\\n (\\n \\\"abs\\\",\\n \\\"add\\\",\\n \\\"addmm\\\",\\n \\\"all\\\",\\n \\\"any\\\",\\n \\\"argmin\\\",\\n \\\"bmm\\\",\\n \\\"clamp\\\",\\n \\\"clamp_min\\\",\\n \\\"cumsum\\\",\\n \\\"div\\\",\\n \\\"fmod\\\",\\n \\\"index_select\\\",\\n \\\"leaky_relu\\\",\\n \\\"linear\\\",\\n \\\"log\\\",\\n \\\"matmul\\\",\\n \\\"mul\\\",\\n \\\"narrow_copy\\\",\\n \\\"nonzero\\\",\\n \\\"pow\\\",\\n \\\"remainder\\\",\\n \\\"sigmoid\\\",\\n \\\"sign\\\",\\n \\\"sub\\\",\\n \\\"tanh\\\",\\n \\\"detach\\\",\\n \\\"expand_as\\\",\\n \\\"flatten\\\",\\n \\\"narrow\\\",\\n \\\"reshape_as\\\",\\n \\\"select\\\",\\n \\\"slice\\\",\\n \\\"softmax\\\",\\n \\\"split\\\",\\n \\\"squeeze\\\",\\n \\\"transpose\\\",\\n \\\"view\\\",\\n \\\"where\\\",\\n )\\n)\\n\\n\\ndef is_hand_written(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> bool:\\n name_base = func_name_base_str(g)\\n return name_base in is_hand_written_ops_\\n\\n\\ndef override_test_values(arg_map: dict[str, str], op_name: str, index: int) -> None:\\n assert index == 0 or index == 1\\n if op_name == \\\"addr\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6})\\\"\\n arg_map[\\\"vec1\\\"] = \\\"at::rand({6})\\\"\\n arg_map[\\\"vec2\\\"] = \\\"at::rand({6})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22})\\\"\\n arg_map[\\\"vec1\\\"] = \\\"at::rand({22})\\\"\\n arg_map[\\\"vec2\\\"] = \\\"at::rand({22})\\\"\\n return\\n if op_name == \\\"mv\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6})\\\"\\n arg_map[\\\"vec\\\"] = \\\"at::rand({6})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22})\\\"\\n arg_map[\\\"vec\\\"] = \\\"at::rand({22})\\\"\\n return\\n if op_name == \\\"addbmm\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22})\\\"\\n return\\n if op_name == \\\"cross\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({3, 3, 3})\\\"\\n arg_map[\\\"other\\\"] = \\\"at::rand({3, 3, 3})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 3, 22})\\\"\\n arg_map[\\\"other\\\"] = \\\"at::rand({22, 3, 22})\\\"\\n return\\n if op_name == \\\"take\\\":\\n if index == 0:\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 216, {20}, torch::kInt64)\\\"\\n else:\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 1000, {100}, torch::kInt64)\\\"\\n return\\n if op_name == \\\"take_along_dim\\\":\\n if index == 0:\\n arg_map[\\\"indices\\\"] = \\\"at::argsort(self0, 1, true)\\\"\\n else:\\n arg_map[\\\"indices\\\"] = \\\"at::argsort(self1, 1, true)\\\"\\n return\\n if op_name == \\\"masked_select\\\":\\n if index == 0:\\n arg_map[\\\"mask\\\"] = \\\"at::randn({6, 6, 6}) > 0.5\\\"\\n else:\\n arg_map[\\\"mask\\\"] = \\\"at::rand({22, 22, 22}) > 0.5\\\"\\n return\\n if op_name == \\\"orgqr\\\":\\n if index == 0:\\n arg_map[\\\"input2\\\"] = \\\"at::rand({6, 6})\\\"\\n else:\\n arg_map[\\\"input2\\\"] = \\\"at::rand({22, 22})\\\"\\n return\\n if op_name == \\\"ormqr\\\":\\n if index == 0:\\n arg_map[\\\"input2\\\"] = \\\"at::rand({6, 6})\\\"\\n else:\\n arg_map[\\\"input2\\\"] = \\\"at::rand({22, 22})\\\"\\n return\\n if op_name == \\\"quantile\\\":\\n if index == 0:\\n arg_map[\\\"q\\\"] = \\\"at::rand({6})\\\"\\n arg_map[\\\"interpolation\\\"] = '\\\"linear\\\"'\\n else:\\n arg_map[\\\"q\\\"] = \\\"at::rand({22})\\\"\\n arg_map[\\\"interpolation\\\"] = '\\\"linear\\\"'\\n return\\n if op_name == \\\"nanquantile\\\":\\n if index == 0:\\n arg_map[\\\"q\\\"] = \\\"at::rand({6})\\\"\\n arg_map[\\\"interpolation\\\"] = '\\\"linear\\\"'\\n else:\\n arg_map[\\\"q\\\"] = \\\"at::rand({22})\\\"\\n arg_map[\\\"interpolation\\\"] = '\\\"linear\\\"'\\n return\\n if op_name == \\\"multi_margin_loss\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(6, {6}, torch::kInt64)\\\"\\n arg_map[\\\"weight\\\"] = \\\"at::rand({6})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(22, {22}, torch::kInt64)\\\"\\n arg_map[\\\"weight\\\"] = \\\"at::rand({22})\\\"\\n return\\n if op_name == \\\"multilabel_margin_loss\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(6, {6, 6}, torch::kInt64)\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(22, {22, 22}, torch::kInt64)\\\"\\n return\\n if op_name == \\\"nll_loss\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(6, {6}, torch::kInt64)\\\"\\n arg_map[\\\"weight\\\"] = \\\"at::rand({6})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(22, {22}, torch::kInt64)\\\"\\n arg_map[\\\"weight\\\"] = \\\"at::rand({22})\\\"\\n return\\n if op_name == \\\"nll_loss2d\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6, 6, 6})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(6, {6, 6, 6}, torch::kInt64)\\\"\\n arg_map[\\\"weight\\\"] = \\\"at::rand({6})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22, 22, 22})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(22, {22, 22, 22}, torch::kInt64)\\\"\\n arg_map[\\\"weight\\\"] = \\\"at::rand({22})\\\"\\n return\\n if op_name in (\\n \\\"fft_fft\\\",\\n \\\"fft_ifft\\\",\\n \\\"fft_rfft\\\",\\n \\\"fft_irfft\\\",\\n \\\"fft_hfft\\\",\\n \\\"fft_ihfft\\\",\\n ):\\n arg_map[\\\"norm\\\"] = '\\\"forward\\\"'\\n return\\n if op_name == \\\"linalg_tensorinv\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6, 6, 6})\\\"\\n arg_map[\\\"ind\\\"] = \\\"2\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22, 22, 22})\\\"\\n arg_map[\\\"ind\\\"] = \\\"2\\\"\\n return\\n if op_name == \\\"addmv\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({2})\\\"\\n arg_map[\\\"mat\\\"] = \\\"at::rand({2, 2})\\\"\\n arg_map[\\\"vec\\\"] = \\\"at::rand({2})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({35})\\\"\\n arg_map[\\\"mat\\\"] = \\\"at::rand({35, 35})\\\"\\n arg_map[\\\"vec\\\"] = \\\"at::rand({35})\\\"\\n return\\n if op_name == \\\"acosh\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({2, 2, 2}) + at::ones({2, 2, 2})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({5, 5, 5}) + at::ones({5, 5, 5})\\\"\\n return\\n if op_name == \\\"adaptive_max_pool2d_backward\\\":\\n if index == 0:\\n arg_map[\\\"grad_output\\\"] = \\\"at::rand({2, 2, 2}, at::kFloat)\\\"\\n arg_map[\\\"self\\\"] = \\\"at::rand({2, 2, 2}, at::kFloat)\\\"\\n arg_map[\\\"indices\\\"] = \\\"at::randint(0, 1, {2, 2, 2}, at::kLong)\\\"\\n else:\\n arg_map[\\\"grad_output\\\"] = \\\"at::rand({3, 3, 3}, at::kFloat)\\\"\\n arg_map[\\\"self\\\"] = \\\"at::rand({3, 3, 3}, at::kFloat)\\\"\\n arg_map[\\\"indices\\\"] = \\\"at::randint(0, 1, {3, 3, 3}, at::kLong)\\\"\\n return\\n if op_name == \\\"adaptive_max_pool3d_backward\\\":\\n if index == 0:\\n arg_map[\\\"grad_output\\\"] = \\\"at::rand({2, 2, 2, 2}, at::kFloat)\\\"\\n arg_map[\\\"self\\\"] = \\\"at::rand({2, 2, 2, 2}, at::kFloat)\\\"\\n arg_map[\\\"indices\\\"] = \\\"at::randint(0, 1, {2, 2, 2, 2}, at::kLong)\\\"\\n else:\\n arg_map[\\\"grad_output\\\"] = \\\"at::rand({3, 3, 3, 3}, at::kFloat)\\\"\\n arg_map[\\\"self\\\"] = \\\"at::rand({3, 3, 3, 3}, at::kFloat)\\\"\\n arg_map[\\\"indices\\\"] = \\\"at::randint(0, 1, {3, 3, 3, 3}, at::kLong)\\\"\\n return\\n if op_name == \\\"bitwise_left_shift\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::randint(1, 1 << 4, {6, 6, 6}, at::kInt)\\\"\\n arg_map[\\\"other\\\"] = \\\"at::randint(1, 26, {6, 6, 6}, at::kInt)\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::randint(1, 1 << 4, {22, 22, 22}, at::kInt)\\\"\\n arg_map[\\\"other\\\"] = \\\"at::randint(1, 26, {22, 22, 22}, at::kInt)\\\"\\n return\\n if op_name == \\\"bitwise_right_shift\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::randint(1 << 21, 1 << 30, {6, 6, 6}, at::kInt)\\\"\\n arg_map[\\\"other\\\"] = \\\"at::randint(1, 22, {6, 6, 6}, at::kInt)\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::randint(1 << 21, 1 << 30, {22, 22, 22}, at::kInt)\\\"\\n arg_map[\\\"other\\\"] = \\\"at::randint(1, 22, {22, 22, 22}, at::kInt)\\\"\\n return\\n if op_name == \\\"gather\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::randint(1, 100, {2,2,2}, at::kInt)\\\"\\n arg_map[\\\"dim\\\"] = \\\"1\\\"\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 1, {2,2,2}, torch::kInt64)\\\"\\n arg_map[\\\"sparse_grad\\\"] = \\\"false\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::randint(1, 100, {5,5,5}, at::kInt)\\\"\\n arg_map[\\\"dim\\\"] = \\\"1\\\"\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 4, {5,5,5}, torch::kInt64)\\\"\\n arg_map[\\\"sparse_grad\\\"] = \\\"false\\\"\\n return\\n if op_name == \\\"gelu\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6, 6})\\\"\\n arg_map[\\\"approximate\\\"] = '\\\"tanh\\\"'\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22, 22})\\\"\\n arg_map[\\\"approximate\\\"] = '\\\"tanh\\\"'\\n return\\n if op_name == \\\"gelu_backward\\\":\\n if index == 0:\\n arg_map[\\\"grad_output\\\"] = \\\"at::rand({6, 6, 6})\\\"\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 6, 6})\\\"\\n arg_map[\\\"approximate\\\"] = '\\\"tanh\\\"'\\n else:\\n arg_map[\\\"grad_output\\\"] = \\\"at::rand({22, 22, 22})\\\"\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 22, 22})\\\"\\n arg_map[\\\"approximate\\\"] = '\\\"tanh\\\"'\\n return\\n if op_name == \\\"index_add\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({2})\\\"\\n arg_map[\\\"dim\\\"] = \\\"0\\\"\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 1, {2}, at::kInt)\\\"\\n arg_map[\\\"source\\\"] = \\\"at::rand({2})\\\"\\n arg_map[\\\"alpha\\\"] = \\\"2\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({16})\\\"\\n arg_map[\\\"dim\\\"] = \\\"0\\\"\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 10, {16}, at::kInt)\\\"\\n arg_map[\\\"source\\\"] = \\\"at::rand({16})\\\"\\n arg_map[\\\"alpha\\\"] = \\\"2\\\"\\n return\\n if op_name == \\\"index_copy\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({2})\\\"\\n arg_map[\\\"dim\\\"] = \\\"0\\\"\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 1, {2}, at::kLong)\\\"\\n arg_map[\\\"source\\\"] = \\\"at::rand({2})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({32})\\\"\\n arg_map[\\\"dim\\\"] = \\\"0\\\"\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 10, {32}, at::kLong)\\\"\\n arg_map[\\\"source\\\"] = \\\"at::rand({32})\\\"\\n return\\n if op_name == \\\"linalg_cross\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({6, 3, 6})\\\"\\n arg_map[\\\"other\\\"] = \\\"at::rand({6, 3, 6})\\\"\\n arg_map[\\\"dim\\\"] = \\\"1\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({22, 3, 22})\\\"\\n arg_map[\\\"other\\\"] = \\\"at::rand({22, 3, 22})\\\"\\n arg_map[\\\"dim\\\"] = \\\"1\\\"\\n return\\n if op_name == \\\"nll_loss_backward\\\":\\n if index == 0:\\n arg_map[\\\"grad_output\\\"] = \\\"at::rand({})\\\"\\n arg_map[\\\"self\\\"] = \\\"at::rand({6})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(0, 5, {6}, torch::kInt64)\\\"\\n arg_map[\\\"weight\\\"] = \\\"at::rand({6})\\\"\\n arg_map[\\\"reduction\\\"] = \\\"1\\\"\\n arg_map[\\\"ignore_index\\\"] = \\\"1\\\"\\n arg_map[\\\"total_weight\\\"] = \\\"at::rand({})\\\"\\n else:\\n arg_map[\\\"grad_output\\\"] = \\\"at::rand({})\\\"\\n arg_map[\\\"self\\\"] = \\\"at::rand({36})\\\"\\n arg_map[\\\"target\\\"] = \\\"at::randint(0, 11, {36}, torch::kInt64)\\\"\\n arg_map[\\\"weight\\\"] = \\\"at::rand({36})\\\"\\n arg_map[\\\"reduction\\\"] = \\\"1\\\"\\n arg_map[\\\"ignore_index\\\"] = \\\"1\\\"\\n arg_map[\\\"total_weight\\\"] = \\\"at::rand({})\\\"\\n return\\n if op_name in [\\\"scatter\\\", \\\"scatter_add\\\", \\\"_scatter_reduce\\\"]:\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::randint(1, 100, {2,2,2}, torch::kInt64)\\\"\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 1, {2,2,2}, torch::kInt64)\\\"\\n arg_map[\\\"src\\\"] = \\\"at::randint(1, 100, {2,2,2}, torch::kInt64)\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::randint(1, 100, {5,5,5}, torch::kInt64)\\\"\\n arg_map[\\\"index\\\"] = \\\"at::randint(0, 1, {5,5,5}, torch::kInt64)\\\"\\n arg_map[\\\"src\\\"] = \\\"at::randint(1, 100, {5,5,5}, torch::kInt64)\\\"\\n if \\\"reduce\\\" in arg_map:\\n arg_map[\\\"reduce\\\"] = '\\\"sum\\\"' if op_name == \\\"_scatter_reduce\\\" else '\\\"add\\\"'\\n return\\n if op_name == \\\"scatter_reduce\\\":\\n arg_map[\\\"reduce\\\"] = '\\\"mean\\\"'\\n if index == 0:\\n arg_map[\\\"index\\\"] = \\\"at::randint(6, {6, 6, 6}, torch::kInt64)\\\"\\n else:\\n arg_map[\\\"index\\\"] = \\\"at::randint(22, {22, 22, 22}, torch::kInt64)\\\"\\n return\\n if op_name == \\\"special_zeta\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::rand({2,2,2}, at::kDouble) + at::ones({2,2,2})\\\"\\n arg_map[\\\"other\\\"] = \\\"at::rand({2,2,2}, at::kDouble) + at::ones({2,2,2})\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::rand({5,5,5}, at::kDouble) + at::ones({5,5,5})\\\"\\n arg_map[\\\"other\\\"] = \\\"at::rand({5,5,5}, at::kDouble) + at::ones({5,5,5})\\\"\\n return\\n if op_name == \\\"_convert_indices_from_csr_to_coo\\\":\\n if index == 0:\\n arg_map[\\\"crow_indices\\\"] = \\\"torch::tensor({1}, torch::kInt32)\\\"\\n arg_map[\\\"col_indices\\\"] = \\\"torch::tensor({0, 1, 0}, torch::kInt32)\\\"\\n arg_map[\\\"out_int32\\\"] = \\\"false\\\"\\n else:\\n arg_map[\\\"crow_indices\\\"] = \\\"torch::tensor({0}, torch::kInt32)\\\"\\n arg_map[\\n \\\"col_indices\\\"\\n ] = \\\"torch::tensor({0, 1, 0, 2, 1, 2, 0, 1, 0, 2, 1, 2}, torch::kInt32)\\\"\\n arg_map[\\\"out_int32\\\"] = \\\"false\\\"\\n return\\n if op_name == \\\"_convert_indices_from_coo_to_csr\\\":\\n if index == 0:\\n arg_map[\\\"self\\\"] = \\\"at::randint(0, 3, {2}, at::kInt)\\\"\\n arg_map[\\\"size\\\"] = \\\"10\\\"\\n arg_map[\\\"out_int32\\\"] = \\\"false\\\"\\n else:\\n arg_map[\\\"self\\\"] = \\\"at::randint(0, 3, {12}, at::kInt)\\\"\\n arg_map[\\\"size\\\"] = \\\"24\\\"\\n arg_map[\\\"out_int32\\\"] = \\\"false\\\"\\n return\\n if op_name in (\\\"diagonal\\\", \\\"linalg_diagonal\\\"):\\n arg_map[\\\"offset\\\"] = \\\"0\\\"\\n arg_map[\\\"dim1\\\"] = \\\"2\\\"\\n arg_map[\\\"dim2\\\"] = \\\"1\\\"\\n return\",\"difficulty\":\"easy\",\"domain\":\"Code Repository Understanding\",\"length\":\"medium\",\"question\":\"In the FileManager class, which of the following wrongly describes the purpose of the write_with_template method, and how it handles file writing while ensuring template substitution?\",\"sub_domain\":\"Code repo QA\"}","display_format":"text","language":"","answer_status":"published","assets":[],"source_url":"https://huggingface.co/datasets/zai-org/LongBench-v2","history":"initial import","indexing_mode":"noindex","subproblems":[],"grids":[]}