diff --git a/torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/data_sparsity.py b/torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/data_sparsity.py index 77ca61d599cb..554ad27dd357 100644 --- a/torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/data_sparsity.py +++ b/torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/data_sparsity.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from collections import defaultdict from copy import deepcopy from typing import Any, Optional, Dict, TYPE_CHECKING diff --git a/torch/ao/pruning/_experimental/data_sparsifier/lightning/tests/test_callbacks.py b/torch/ao/pruning/_experimental/data_sparsifier/lightning/tests/test_callbacks.py index 252405de4968..957254284215 100644 --- a/torch/ao/pruning/_experimental/data_sparsifier/lightning/tests/test_callbacks.py +++ b/torch/ao/pruning/_experimental/data_sparsifier/lightning/tests/test_callbacks.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from torch.ao.pruning._experimental.data_sparsifier.data_norm_sparsifier import DataNormSparsifier from torch.ao.pruning._experimental.data_scheduler.base_data_scheduler import BaseDataScheduler import torch diff --git a/torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py b/torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py index 1e76cfc345ac..0e907f42d3bf 100644 --- a/torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py +++ b/torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch import torch.nn as nn from torch.ao.pruning.sparsifier.utils import module_to_fqn, fqn_to_module diff --git a/torch/ao/pruning/_experimental/pruner/FPGM_pruner.py b/torch/ao/pruning/_experimental/pruner/FPGM_pruner.py index d8c3d20052ba..fe874c6effc7 100644 --- a/torch/ao/pruning/_experimental/pruner/FPGM_pruner.py +++ b/torch/ao/pruning/_experimental/pruner/FPGM_pruner.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Callable, Optional, Union import torch diff --git a/torch/ao/pruning/_experimental/pruner/base_structured_sparsifier.py b/torch/ao/pruning/_experimental/pruner/base_structured_sparsifier.py index 357421fb5529..b380ae00adce 100644 --- a/torch/ao/pruning/_experimental/pruner/base_structured_sparsifier.py +++ b/torch/ao/pruning/_experimental/pruner/base_structured_sparsifier.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from itertools import chain from operator import getitem import torch diff --git a/torch/ao/pruning/_experimental/pruner/lstm_saliency_pruner.py b/torch/ao/pruning/_experimental/pruner/lstm_saliency_pruner.py index 9e569c14a6c8..3b65ce59fecc 100644 --- a/torch/ao/pruning/_experimental/pruner/lstm_saliency_pruner.py +++ b/torch/ao/pruning/_experimental/pruner/lstm_saliency_pruner.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import cast import torch diff --git a/torch/ao/pruning/_experimental/pruner/parametrization.py b/torch/ao/pruning/_experimental/pruner/parametrization.py index df94f7093b53..c5aa74e3bc52 100644 --- a/torch/ao/pruning/_experimental/pruner/parametrization.py +++ b/torch/ao/pruning/_experimental/pruner/parametrization.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch import nn from torch.nn.utils.parametrize import is_parametrized diff --git a/torch/ao/pruning/_experimental/pruner/prune_functions.py b/torch/ao/pruning/_experimental/pruner/prune_functions.py index 2b16d4b327a0..f7dcf120f9c3 100644 --- a/torch/ao/pruning/_experimental/pruner/prune_functions.py +++ b/torch/ao/pruning/_experimental/pruner/prune_functions.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs """ Collection of conversion functions for linear / conv2d structured pruning Also contains utilities for bias propagation diff --git a/torch/ao/pruning/_experimental/pruner/saliency_pruner.py b/torch/ao/pruning/_experimental/pruner/saliency_pruner.py index 7f96f0865d30..cf932c272005 100644 --- a/torch/ao/pruning/_experimental/pruner/saliency_pruner.py +++ b/torch/ao/pruning/_experimental/pruner/saliency_pruner.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from .base_structured_sparsifier import BaseStructuredSparsifier diff --git a/torch/ao/pruning/_mappings.py b/torch/ao/pruning/_mappings.py index 726cbc6b0fc8..70a0c785190f 100644 --- a/torch/ao/pruning/_mappings.py +++ b/torch/ao/pruning/_mappings.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs __all__ = [ "get_static_sparse_quantized_mapping", "get_dynamic_sparse_quantized_mapping", diff --git a/torch/ao/pruning/scheduler/base_scheduler.py b/torch/ao/pruning/scheduler/base_scheduler.py index 3391d3e73cd6..82f02399b7ec 100644 --- a/torch/ao/pruning/scheduler/base_scheduler.py +++ b/torch/ao/pruning/scheduler/base_scheduler.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from torch.ao.pruning import BaseSparsifier diff --git a/torch/ao/pruning/scheduler/cubic_scheduler.py b/torch/ao/pruning/scheduler/cubic_scheduler.py index 76fc61daa288..1a883059f569 100644 --- a/torch/ao/pruning/scheduler/cubic_scheduler.py +++ b/torch/ao/pruning/scheduler/cubic_scheduler.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import warnings from .base_scheduler import BaseScheduler diff --git a/torch/ao/pruning/scheduler/lambda_scheduler.py b/torch/ao/pruning/scheduler/lambda_scheduler.py index a88d99a1f83b..5236ebc33a26 100644 --- a/torch/ao/pruning/scheduler/lambda_scheduler.py +++ b/torch/ao/pruning/scheduler/lambda_scheduler.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import warnings from .base_scheduler import BaseScheduler diff --git a/torch/ao/pruning/sparsifier/base_sparsifier.py b/torch/ao/pruning/sparsifier/base_sparsifier.py index 1c210ace344d..8afed4d68945 100644 --- a/torch/ao/pruning/sparsifier/base_sparsifier.py +++ b/torch/ao/pruning/sparsifier/base_sparsifier.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import abc import copy from collections import defaultdict diff --git a/torch/ao/pruning/sparsifier/nearly_diagonal_sparsifier.py b/torch/ao/pruning/sparsifier/nearly_diagonal_sparsifier.py index 4f44e81485df..419323e68f93 100644 --- a/torch/ao/pruning/sparsifier/nearly_diagonal_sparsifier.py +++ b/torch/ao/pruning/sparsifier/nearly_diagonal_sparsifier.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from . import base_sparsifier diff --git a/torch/ao/pruning/sparsifier/utils.py b/torch/ao/pruning/sparsifier/utils.py index 98f489904cc4..7fd93e4d9da7 100644 --- a/torch/ao/pruning/sparsifier/utils.py +++ b/torch/ao/pruning/sparsifier/utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Any, Dict, Optional, Type from torch.nn.utils.parametrize import type_before_parametrizations, is_parametrized from itertools import chain diff --git a/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py b/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py index 2b24ca3d82e3..2f50d51f2a38 100644 --- a/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py +++ b/torch/ao/pruning/sparsifier/weight_norm_sparsifier.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from functools import reduce from typing import Callable, Optional, Tuple, Union diff --git a/torch/ao/quantization/__init__.py b/torch/ao/quantization/__init__.py index e2b8ee5c810a..f77969b32149 100644 --- a/torch/ao/quantization/__init__.py +++ b/torch/ao/quantization/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs # flake8: noqa: F403 from .fake_quantize import * # noqa: F403 diff --git a/torch/ao/quantization/_correct_bias.py b/torch/ao/quantization/_correct_bias.py index 83cc81bb6b00..bf6b42a4a0dc 100644 --- a/torch/ao/quantization/_correct_bias.py +++ b/torch/ao/quantization/_correct_bias.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch import torch.nn as nn import torch.ao.nn.quantized as nnq diff --git a/torch/ao/quantization/_equalize.py b/torch/ao/quantization/_equalize.py index 7d39dbcf1ca8..4fed532c56f0 100644 --- a/torch/ao/quantization/_equalize.py +++ b/torch/ao/quantization/_equalize.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch import copy from typing import Dict, Any diff --git a/torch/ao/quantization/_learnable_fake_quantize.py b/torch/ao/quantization/_learnable_fake_quantize.py index cdf44c5ea7b2..ce23e80de150 100644 --- a/torch/ao/quantization/_learnable_fake_quantize.py +++ b/torch/ao/quantization/_learnable_fake_quantize.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch.nn.parameter import Parameter from typing import List diff --git a/torch/ao/quantization/backend_config/_common_operator_config_utils.py b/torch/ao/quantization/backend_config/_common_operator_config_utils.py index 4e946a25ffbb..d76bdfddddaf 100644 --- a/torch/ao/quantization/backend_config/_common_operator_config_utils.py +++ b/torch/ao/quantization/backend_config/_common_operator_config_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import copy import operator import torch diff --git a/torch/ao/quantization/backend_config/_qnnpack_pt2e.py b/torch/ao/quantization/backend_config/_qnnpack_pt2e.py index 01e112b688c0..871d26dd9ff7 100644 --- a/torch/ao/quantization/backend_config/_qnnpack_pt2e.py +++ b/torch/ao/quantization/backend_config/_qnnpack_pt2e.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import operator import torch from torch.ao.quantization.backend_config import ( diff --git a/torch/ao/quantization/backend_config/backend_config.py b/torch/ao/quantization/backend_config/backend_config.py index 2288aced0995..96fb66662d6f 100644 --- a/torch/ao/quantization/backend_config/backend_config.py +++ b/torch/ao/quantization/backend_config/backend_config.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from __future__ import annotations from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Type, Union, TYPE_CHECKING diff --git a/torch/ao/quantization/backend_config/native.py b/torch/ao/quantization/backend_config/native.py index 81cfc928adb5..84e0fbc45c62 100644 --- a/torch/ao/quantization/backend_config/native.py +++ b/torch/ao/quantization/backend_config/native.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from ._common_operator_config_utils import ( _get_binary_op_configs, diff --git a/torch/ao/quantization/backend_config/onednn.py b/torch/ao/quantization/backend_config/onednn.py index 6eab945f7d74..88dffedfd81b 100644 --- a/torch/ao/quantization/backend_config/onednn.py +++ b/torch/ao/quantization/backend_config/onednn.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch import torch.nn as nn import torch.ao.nn.intrinsic as nni diff --git a/torch/ao/quantization/backend_config/tensorrt.py b/torch/ao/quantization/backend_config/tensorrt.py index 1c5f761508bb..7a80d1883cfd 100644 --- a/torch/ao/quantization/backend_config/tensorrt.py +++ b/torch/ao/quantization/backend_config/tensorrt.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from .backend_config import ( BackendConfig, diff --git a/torch/ao/quantization/backend_config/utils.py b/torch/ao/quantization/backend_config/utils.py index 2e7382274079..13bf632e251a 100644 --- a/torch/ao/quantization/backend_config/utils.py +++ b/torch/ao/quantization/backend_config/utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Dict, Any, List, Callable, Union, Tuple, Type import torch diff --git a/torch/ao/quantization/experimental/APoT_tensor.py b/torch/ao/quantization/experimental/APoT_tensor.py index debda7aea8c0..6caa2334be07 100644 --- a/torch/ao/quantization/experimental/APoT_tensor.py +++ b/torch/ao/quantization/experimental/APoT_tensor.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch.ao.quantization.experimental.quantizer import APoTQuantizer diff --git a/torch/ao/quantization/experimental/adaround_fake_quantize.py b/torch/ao/quantization/experimental/adaround_fake_quantize.py index 4d988bbb25bb..d035a02b047a 100644 --- a/torch/ao/quantization/experimental/adaround_fake_quantize.py +++ b/torch/ao/quantization/experimental/adaround_fake_quantize.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Tuple import torch diff --git a/torch/ao/quantization/experimental/adaround_optimization.py b/torch/ao/quantization/experimental/adaround_optimization.py index 808b7abe2c78..f7eedd9fef12 100644 --- a/torch/ao/quantization/experimental/adaround_optimization.py +++ b/torch/ao/quantization/experimental/adaround_optimization.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import copy from typing import Any, Callable, List, Optional, Tuple, Type, Union diff --git a/torch/ao/quantization/experimental/apot_utils.py b/torch/ao/quantization/experimental/apot_utils.py index ad7a7bed1fbe..c2f2f0746ca5 100644 --- a/torch/ao/quantization/experimental/apot_utils.py +++ b/torch/ao/quantization/experimental/apot_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs r""" This file contains utility functions to convert values using APoT nonuniform quantization methods. diff --git a/torch/ao/quantization/experimental/fake_quantize.py b/torch/ao/quantization/experimental/fake_quantize.py index 7541106a61c8..6b4da74541f2 100644 --- a/torch/ao/quantization/experimental/fake_quantize.py +++ b/torch/ao/quantization/experimental/fake_quantize.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch import Tensor from torch.ao.quantization.experimental.observer import APoTObserver diff --git a/torch/ao/quantization/experimental/fake_quantize_function.py b/torch/ao/quantization/experimental/fake_quantize_function.py index cac01fd8c002..924c81fc08df 100644 --- a/torch/ao/quantization/experimental/fake_quantize_function.py +++ b/torch/ao/quantization/experimental/fake_quantize_function.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch import Tensor from torch.ao.quantization.experimental.quantizer import quantize_APoT, dequantize_APoT diff --git a/torch/ao/quantization/experimental/linear.py b/torch/ao/quantization/experimental/linear.py index 154023b16183..cb46c99b01af 100644 --- a/torch/ao/quantization/experimental/linear.py +++ b/torch/ao/quantization/experimental/linear.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch import numpy as np diff --git a/torch/ao/quantization/experimental/observer.py b/torch/ao/quantization/experimental/observer.py index 76a63815bdc6..8474f69c26a2 100644 --- a/torch/ao/quantization/experimental/observer.py +++ b/torch/ao/quantization/experimental/observer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs """ This module implements nonuniform observers used to collect statistics about the values observed during calibration (PTQ) or training (QAT). diff --git a/torch/ao/quantization/experimental/quantizer.py b/torch/ao/quantization/experimental/quantizer.py index df9c0f27847e..b386ce20bbd3 100644 --- a/torch/ao/quantization/experimental/quantizer.py +++ b/torch/ao/quantization/experimental/quantizer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch import Tensor import numpy as np diff --git a/torch/ao/quantization/fake_quantize.py b/torch/ao/quantization/fake_quantize.py index 9f0503cf06a5..b921df39217a 100644 --- a/torch/ao/quantization/fake_quantize.py +++ b/torch/ao/quantization/fake_quantize.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs """Implements modules used to perform fake quantization.""" import torch diff --git a/torch/ao/quantization/fuse_modules.py b/torch/ao/quantization/fuse_modules.py index 2caa0a2b7f2d..b9447ff37e39 100644 --- a/torch/ao/quantization/fuse_modules.py +++ b/torch/ao/quantization/fuse_modules.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import copy import torch.nn as nn diff --git a/torch/ao/quantization/fuser_method_mappings.py b/torch/ao/quantization/fuser_method_mappings.py index 16c0c3a85b8f..a989ae298825 100644 --- a/torch/ao/quantization/fuser_method_mappings.py +++ b/torch/ao/quantization/fuser_method_mappings.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch.nn as nn import torch.ao.nn.intrinsic as nni diff --git a/torch/ao/quantization/fx/_decomposed.py b/torch/ao/quantization/fx/_decomposed.py index f2e774590be3..72ce4b2471f5 100644 --- a/torch/ao/quantization/fx/_decomposed.py +++ b/torch/ao/quantization/fx/_decomposed.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import math from typing import Optional, Tuple diff --git a/torch/ao/quantization/fx/_equalize.py b/torch/ao/quantization/fx/_equalize.py index b0965b9a7051..40a7e7bbff3b 100644 --- a/torch/ao/quantization/fx/_equalize.py +++ b/torch/ao/quantization/fx/_equalize.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import warnings from collections import namedtuple diff --git a/torch/ao/quantization/fx/_lower_to_native_backend.py b/torch/ao/quantization/fx/_lower_to_native_backend.py index 049f4e3135d9..92620a169383 100644 --- a/torch/ao/quantization/fx/_lower_to_native_backend.py +++ b/torch/ao/quantization/fx/_lower_to_native_backend.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch.fx import map_arg, Node from torch.fx.graph import Graph diff --git a/torch/ao/quantization/fx/_model_report/detector.py b/torch/ao/quantization/fx/_model_report/detector.py index b5c7f9fd2976..8e59df51c6ff 100644 --- a/torch/ao/quantization/fx/_model_report/detector.py +++ b/torch/ao/quantization/fx/_model_report/detector.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Any, Dict, Set, Tuple, Callable, List import torch diff --git a/torch/ao/quantization/fx/_model_report/model_report.py b/torch/ao/quantization/fx/_model_report/model_report.py index 724e76ad576f..3370d8c9baf6 100644 --- a/torch/ao/quantization/fx/_model_report/model_report.py +++ b/torch/ao/quantization/fx/_model_report/model_report.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Any, Dict, Set, Tuple, Callable from collections import OrderedDict import torch diff --git a/torch/ao/quantization/fx/_model_report/model_report_observer.py b/torch/ao/quantization/fx/_model_report/model_report_observer.py index eaa45264be7e..f04d6da8a054 100644 --- a/torch/ao/quantization/fx/_model_report/model_report_observer.py +++ b/torch/ao/quantization/fx/_model_report/model_report_observer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch.ao.quantization.observer import ObserverBase diff --git a/torch/ao/quantization/fx/_model_report/model_report_visualizer.py b/torch/ao/quantization/fx/_model_report/model_report_visualizer.py index 5463862aa1cd..e6288c6f71d9 100644 --- a/torch/ao/quantization/fx/_model_report/model_report_visualizer.py +++ b/torch/ao/quantization/fx/_model_report/model_report_visualizer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from typing import Any, Set, Dict, List, Tuple, OrderedDict from collections import OrderedDict as OrdDict diff --git a/torch/ao/quantization/fx/custom_config.py b/torch/ao/quantization/fx/custom_config.py index 4fb2c3a28cb0..72f28ddbc777 100644 --- a/torch/ao/quantization/fx/custom_config.py +++ b/torch/ao/quantization/fx/custom_config.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from __future__ import annotations from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Type diff --git a/torch/ao/quantization/fx/fuse.py b/torch/ao/quantization/fx/fuse.py index 6b2b614728f8..b555789f673a 100644 --- a/torch/ao/quantization/fx/fuse.py +++ b/torch/ao/quantization/fx/fuse.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from torch.fx import ( GraphModule, Node, diff --git a/torch/ao/quantization/fx/fuse_handler.py b/torch/ao/quantization/fx/fuse_handler.py index 718cc561bfa0..2766211e8e1b 100644 --- a/torch/ao/quantization/fx/fuse_handler.py +++ b/torch/ao/quantization/fx/fuse_handler.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch.ao.quantization.backend_config import BackendConfig from torch.fx.graph import Node, Graph diff --git a/torch/ao/quantization/fx/graph_module.py b/torch/ao/quantization/fx/graph_module.py index cc9187285ae6..224f71745157 100644 --- a/torch/ao/quantization/fx/graph_module.py +++ b/torch/ao/quantization/fx/graph_module.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch import copy from torch.fx import GraphModule diff --git a/torch/ao/quantization/fx/match_utils.py b/torch/ao/quantization/fx/match_utils.py index cf287db8c524..b5a6657103fc 100644 --- a/torch/ao/quantization/fx/match_utils.py +++ b/torch/ao/quantization/fx/match_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import sys import torch from torch.fx.graph import ( diff --git a/torch/ao/quantization/fx/pattern_utils.py b/torch/ao/quantization/fx/pattern_utils.py index d8648a0aed5e..3665f75f7567 100644 --- a/torch/ao/quantization/fx/pattern_utils.py +++ b/torch/ao/quantization/fx/pattern_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from collections import OrderedDict from typing import Dict, Any from torch.ao.quantization.utils import Pattern diff --git a/torch/ao/quantization/fx/prepare.py b/torch/ao/quantization/fx/prepare.py index d8e25f1260f5..80f50581cc72 100644 --- a/torch/ao/quantization/fx/prepare.py +++ b/torch/ao/quantization/fx/prepare.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import copy import torch import warnings diff --git a/torch/ao/quantization/fx/qconfig_mapping_utils.py b/torch/ao/quantization/fx/qconfig_mapping_utils.py index 0b906a1777de..378c51b6805d 100644 --- a/torch/ao/quantization/fx/qconfig_mapping_utils.py +++ b/torch/ao/quantization/fx/qconfig_mapping_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch import re from collections import defaultdict, OrderedDict diff --git a/torch/ao/quantization/fx/quantize_handler.py b/torch/ao/quantization/fx/quantize_handler.py index e70040f7e649..83fee8efcd99 100644 --- a/torch/ao/quantization/fx/quantize_handler.py +++ b/torch/ao/quantization/fx/quantize_handler.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from abc import ABC from typing import Callable, Dict, List, Optional, Type diff --git a/torch/ao/quantization/fx/utils.py b/torch/ao/quantization/fx/utils.py index 5cfedde4bc24..5029db47961f 100644 --- a/torch/ao/quantization/fx/utils.py +++ b/torch/ao/quantization/fx/utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import copy import torch import torch.nn as nn diff --git a/torch/ao/quantization/observer.py b/torch/ao/quantization/observer.py index 5f075df1cd83..656372d37555 100644 --- a/torch/ao/quantization/observer.py +++ b/torch/ao/quantization/observer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs """ This module implements observers which are used to collect statistics about the values observed during calibration (PTQ) or training (QAT). diff --git a/torch/ao/quantization/pt2e/duplicate_dq_pass.py b/torch/ao/quantization/pt2e/duplicate_dq_pass.py index 48c7d7247b99..a6cfbce611fa 100644 --- a/torch/ao/quantization/pt2e/duplicate_dq_pass.py +++ b/torch/ao/quantization/pt2e/duplicate_dq_pass.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import logging import operator diff --git a/torch/ao/quantization/pt2e/export_utils.py b/torch/ao/quantization/pt2e/export_utils.py index 139042c326b8..78c69b718d7d 100644 --- a/torch/ao/quantization/pt2e/export_utils.py +++ b/torch/ao/quantization/pt2e/export_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import types import torch diff --git a/torch/ao/quantization/pt2e/graph_utils.py b/torch/ao/quantization/pt2e/graph_utils.py index bacb4d8a28f1..6ae93ba1d260 100644 --- a/torch/ao/quantization/pt2e/graph_utils.py +++ b/torch/ao/quantization/pt2e/graph_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import itertools from typing import Any, List, OrderedDict, Set, Optional, Callable import operator diff --git a/torch/ao/quantization/pt2e/port_metadata_pass.py b/torch/ao/quantization/pt2e/port_metadata_pass.py index 5ea1f939a3b6..313b420e7a22 100644 --- a/torch/ao/quantization/pt2e/port_metadata_pass.py +++ b/torch/ao/quantization/pt2e/port_metadata_pass.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import logging from typing import Optional diff --git a/torch/ao/quantization/pt2e/prepare.py b/torch/ao/quantization/pt2e/prepare.py index 169a982f62ce..162ee45623ee 100644 --- a/torch/ao/quantization/pt2e/prepare.py +++ b/torch/ao/quantization/pt2e/prepare.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch._subclasses import FakeTensor from torch.ao.quantization.fx.prepare import ( diff --git a/torch/ao/quantization/pt2e/qat_utils.py b/torch/ao/quantization/pt2e/qat_utils.py index 45f5c265d2cb..c4c1f804d41c 100644 --- a/torch/ao/quantization/pt2e/qat_utils.py +++ b/torch/ao/quantization/pt2e/qat_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import dataclasses import itertools import operator diff --git a/torch/ao/quantization/pt2e/representation/rewrite.py b/torch/ao/quantization/pt2e/representation/rewrite.py index 7f5cb2eeb13b..40801344740b 100644 --- a/torch/ao/quantization/pt2e/representation/rewrite.py +++ b/torch/ao/quantization/pt2e/representation/rewrite.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch.fx import GraphModule from ..export_utils import _WrapperModule diff --git a/torch/ao/quantization/pt2e/utils.py b/torch/ao/quantization/pt2e/utils.py index 25f82f04e4e3..cde22426ae5b 100644 --- a/torch/ao/quantization/pt2e/utils.py +++ b/torch/ao/quantization/pt2e/utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import operator import types diff --git a/torch/ao/quantization/qconfig.py b/torch/ao/quantization/qconfig.py index 88e7b47aff2b..dc93d7938f0c 100644 --- a/torch/ao/quantization/qconfig.py +++ b/torch/ao/quantization/qconfig.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from collections import namedtuple from typing import Optional, Any, Union, Type from typing_extensions import deprecated diff --git a/torch/ao/quantization/qconfig_mapping.py b/torch/ao/quantization/qconfig_mapping.py index 6bf4b41c724a..37f71465afea 100644 --- a/torch/ao/quantization/qconfig_mapping.py +++ b/torch/ao/quantization/qconfig_mapping.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from __future__ import annotations from collections import OrderedDict from typing import Any, Callable, Dict, Tuple, Union, List diff --git a/torch/ao/quantization/quantize.py b/torch/ao/quantization/quantize.py index 534def354573..be00be0e295b 100644 --- a/torch/ao/quantization/quantize.py +++ b/torch/ao/quantization/quantize.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import copy import itertools import warnings diff --git a/torch/ao/quantization/quantize_jit.py b/torch/ao/quantization/quantize_jit.py index 632fc1db2327..3001deb6ab9c 100644 --- a/torch/ao/quantization/quantize_jit.py +++ b/torch/ao/quantization/quantize_jit.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch.ao.quantization.qconfig import QConfig diff --git a/torch/ao/quantization/quantizer/embedding_quantizer.py b/torch/ao/quantization/quantizer/embedding_quantizer.py index 81306943264b..bd3d2773e628 100644 --- a/torch/ao/quantization/quantizer/embedding_quantizer.py +++ b/torch/ao/quantization/quantizer/embedding_quantizer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from __future__ import annotations import copy diff --git a/torch/ao/quantization/quantizer/quantizer.py b/torch/ao/quantization/quantizer/quantizer.py index a521ff56c34c..4cecfee28f2b 100644 --- a/torch/ao/quantization/quantizer/quantizer.py +++ b/torch/ao/quantization/quantizer/quantizer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import Callable, Dict, List, Optional, Tuple, Union diff --git a/torch/ao/quantization/quantizer/utils.py b/torch/ao/quantization/quantizer/utils.py index f25d0916018b..f948dbb112dc 100644 --- a/torch/ao/quantization/quantizer/utils.py +++ b/torch/ao/quantization/quantizer/utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import List from torch.ao.quantization.pt2e.utils import _is_sym_size_node diff --git a/torch/ao/quantization/quantizer/x86_inductor_quantizer.py b/torch/ao/quantization/quantizer/x86_inductor_quantizer.py index 4cc05e46c6a7..89e4966bf4eb 100644 --- a/torch/ao/quantization/quantizer/x86_inductor_quantizer.py +++ b/torch/ao/quantization/quantizer/x86_inductor_quantizer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import copy import functools import itertools diff --git a/torch/ao/quantization/quantizer/xnnpack_quantizer.py b/torch/ao/quantization/quantizer/xnnpack_quantizer.py index f3d1b6ca8b39..ae9ae60b8a3b 100644 --- a/torch/ao/quantization/quantizer/xnnpack_quantizer.py +++ b/torch/ao/quantization/quantizer/xnnpack_quantizer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from __future__ import annotations import copy diff --git a/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py b/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py index 9f1732e57370..928ee0d3ac45 100644 --- a/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py +++ b/torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import itertools import operator from dataclasses import dataclass diff --git a/torch/ao/quantization/stubs.py b/torch/ao/quantization/stubs.py index 10a63fb8f0ee..f62a227f1d77 100644 --- a/torch/ao/quantization/stubs.py +++ b/torch/ao/quantization/stubs.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from torch import nn diff --git a/torch/ao/quantization/utils.py b/torch/ao/quantization/utils.py index 5ce1d1109e72..fadbf33a70b6 100644 --- a/torch/ao/quantization/utils.py +++ b/torch/ao/quantization/utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs """ Utils shared by different modes of quantization (eager/graph) """ diff --git a/torch/autograd/__init__.py b/torch/autograd/__init__.py index adf47ad1727d..aca9abb24070 100644 --- a/torch/autograd/__init__.py +++ b/torch/autograd/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs """ ``torch.autograd`` provides classes and functions implementing automatic differentiation of arbitrary scalar valued functions. It requires minimal diff --git a/torch/autograd/_functions/tensor.py b/torch/autograd/_functions/tensor.py index d2b3149bfc81..9c982b074b65 100644 --- a/torch/autograd/_functions/tensor.py +++ b/torch/autograd/_functions/tensor.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import operator from functools import reduce from typing_extensions import deprecated diff --git a/torch/autograd/_functions/utils.py b/torch/autograd/_functions/utils.py index 7111d893400f..56baae4aae3b 100644 --- a/torch/autograd/_functions/utils.py +++ b/torch/autograd/_functions/utils.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import operator from functools import reduce diff --git a/torch/autograd/anomaly_mode.py b/torch/autograd/anomaly_mode.py index 80a2526a81de..7e73ad4ef2c3 100644 --- a/torch/autograd/anomaly_mode.py +++ b/torch/autograd/anomaly_mode.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import warnings import torch diff --git a/torch/autograd/forward_ad.py b/torch/autograd/forward_ad.py index 747b18f0f369..4187e220ceab 100644 --- a/torch/autograd/forward_ad.py +++ b/torch/autograd/forward_ad.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import os from collections import namedtuple diff --git a/torch/autograd/function.py b/torch/autograd/function.py index 9aca2b2a1b32..62ec1183a365 100644 --- a/torch/autograd/function.py +++ b/torch/autograd/function.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import functools import inspect import itertools diff --git a/torch/autograd/functional.py b/torch/autograd/functional.py index 6701efbedac1..8cf3955a6927 100644 --- a/torch/autograd/functional.py +++ b/torch/autograd/functional.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import List, Tuple import torch diff --git a/torch/autograd/grad_mode.py b/torch/autograd/grad_mode.py index be173c9b9de0..1c97ab58298b 100644 --- a/torch/autograd/grad_mode.py +++ b/torch/autograd/grad_mode.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Any import torch diff --git a/torch/autograd/gradcheck.py b/torch/autograd/gradcheck.py index a0d874038761..5bf74afacb66 100644 --- a/torch/autograd/gradcheck.py +++ b/torch/autograd/gradcheck.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import collections import functools import warnings diff --git a/torch/autograd/graph.py b/torch/autograd/graph.py index 19938c183557..cde56a6f26c7 100644 --- a/torch/autograd/graph.py +++ b/torch/autograd/graph.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import abc import collections import contextlib diff --git a/torch/autograd/profiler.py b/torch/autograd/profiler.py index 162dfe1eeaef..0392a8769846 100644 --- a/torch/autograd/profiler.py +++ b/torch/autograd/profiler.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from collections import defaultdict from dataclasses import dataclass from time import perf_counter_ns diff --git a/torch/autograd/profiler_legacy.py b/torch/autograd/profiler_legacy.py index cb573faf4410..40baafd441ae 100644 --- a/torch/autograd/profiler_legacy.py +++ b/torch/autograd/profiler_legacy.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import itertools import warnings from typing_extensions import deprecated diff --git a/torch/autograd/profiler_util.py b/torch/autograd/profiler_util.py index 23243733aaa8..a5cff1ea12a8 100644 --- a/torch/autograd/profiler_util.py +++ b/torch/autograd/profiler_util.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import bisect import itertools import math diff --git a/torch/autograd/variable.py b/torch/autograd/variable.py index ed841d4da7d4..84b504a9c82c 100644 --- a/torch/autograd/variable.py +++ b/torch/autograd/variable.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from torch._C import _ImperativeEngine as ImperativeEngine diff --git a/torch/backends/__init__.py b/torch/backends/__init__.py index 2236230e8c6d..086147b87a81 100644 --- a/torch/backends/__init__.py +++ b/torch/backends/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import types from contextlib import contextmanager diff --git a/torch/backends/_coreml/preprocess.py b/torch/backends/_coreml/preprocess.py index f393929bb7c2..18cb8229db9a 100644 --- a/torch/backends/_coreml/preprocess.py +++ b/torch/backends/_coreml/preprocess.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import hashlib import json from typing import Dict, Tuple diff --git a/torch/backends/_nnapi/prepare.py b/torch/backends/_nnapi/prepare.py index 8b07c3d6e0c6..6ba389902c9f 100644 --- a/torch/backends/_nnapi/prepare.py +++ b/torch/backends/_nnapi/prepare.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import List, Optional import torch diff --git a/torch/backends/_nnapi/serializer.py b/torch/backends/_nnapi/serializer.py index 551fa821df68..34bcc42f8927 100644 --- a/torch/backends/_nnapi/serializer.py +++ b/torch/backends/_nnapi/serializer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import array import enum import functools diff --git a/torch/backends/cuda/__init__.py b/torch/backends/cuda/__init__.py index c35a962ba693..cb5f511bc5db 100644 --- a/torch/backends/cuda/__init__.py +++ b/torch/backends/cuda/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import contextlib from typing import Union diff --git a/torch/backends/cudnn/__init__.py b/torch/backends/cudnn/__init__.py index e00d92f44b28..e528ac68552d 100644 --- a/torch/backends/cudnn/__init__.py +++ b/torch/backends/cudnn/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import os import sys import warnings diff --git a/torch/backends/cudnn/rnn.py b/torch/backends/cudnn/rnn.py index aaf0bd02e8af..f2e9d4321a02 100644 --- a/torch/backends/cudnn/rnn.py +++ b/torch/backends/cudnn/rnn.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch.cuda try: diff --git a/torch/backends/mkl/__init__.py b/torch/backends/mkl/__init__.py index 261ee764485b..9f96d692ae02 100644 --- a/torch/backends/mkl/__init__.py +++ b/torch/backends/mkl/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch diff --git a/torch/backends/mkldnn/__init__.py b/torch/backends/mkldnn/__init__.py index 9cdee1cbd565..669ed59a1132 100644 --- a/torch/backends/mkldnn/__init__.py +++ b/torch/backends/mkldnn/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import sys from contextlib import contextmanager diff --git a/torch/backends/mps/__init__.py b/torch/backends/mps/__init__.py index 8d5e70f06a0a..06eda58e82f9 100644 --- a/torch/backends/mps/__init__.py +++ b/torch/backends/mps/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from functools import lru_cache as _lru_cache from typing import Optional diff --git a/torch/backends/nnpack/__init__.py b/torch/backends/nnpack/__init__.py index 892dfa022cfc..1a30e977cab3 100644 --- a/torch/backends/nnpack/__init__.py +++ b/torch/backends/nnpack/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from contextlib import contextmanager import torch diff --git a/torch/backends/openmp/__init__.py b/torch/backends/openmp/__init__.py index 4a7fcca12d0c..aff8d46cd4ac 100644 --- a/torch/backends/openmp/__init__.py +++ b/torch/backends/openmp/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch diff --git a/torch/backends/opt_einsum/__init__.py b/torch/backends/opt_einsum/__init__.py index 2e66cd37542d..993a219fa9aa 100644 --- a/torch/backends/opt_einsum/__init__.py +++ b/torch/backends/opt_einsum/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import sys import warnings from contextlib import contextmanager diff --git a/torch/backends/quantized/__init__.py b/torch/backends/quantized/__init__.py index 85009753e0ae..3cb795dd39fc 100644 --- a/torch/backends/quantized/__init__.py +++ b/torch/backends/quantized/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import sys import types from typing import List diff --git a/torch/backends/xeon/run_cpu.py b/torch/backends/xeon/run_cpu.py index 0344631ee6b4..bdf07e286174 100644 --- a/torch/backends/xeon/run_cpu.py +++ b/torch/backends/xeon/run_cpu.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs """ This is a script for launching PyTorch inference on Intel(R) Xeon(R) Scalable Processors with optimal configurations. diff --git a/torch/backends/xnnpack/__init__.py b/torch/backends/xnnpack/__init__.py index c26dc11deb47..31e69876927d 100644 --- a/torch/backends/xnnpack/__init__.py +++ b/torch/backends/xnnpack/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import sys import types diff --git a/torch/compiler/__init__.py b/torch/compiler/__init__.py index a27238c3d833..812bbaa4c660 100644 --- a/torch/compiler/__init__.py +++ b/torch/compiler/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import torch from typing import List diff --git a/torch/contrib/_tensorboard_vis.py b/torch/contrib/_tensorboard_vis.py index 87c325948a8b..ed1445dd7bce 100644 --- a/torch/contrib/_tensorboard_vis.py +++ b/torch/contrib/_tensorboard_vis.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import time from collections import defaultdict from functools import partial diff --git a/torch/cpu/__init__.py b/torch/cpu/__init__.py index a36594a3cb15..d2b8069048cc 100644 --- a/torch/cpu/__init__.py +++ b/torch/cpu/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs r""" This package implements abstractions found in ``torch.cuda`` to facilitate writing device-agnostic code. diff --git a/torch/cpu/amp/autocast_mode.py b/torch/cpu/amp/autocast_mode.py index b545e91dd6f4..b61e9b542dba 100644 --- a/torch/cpu/amp/autocast_mode.py +++ b/torch/cpu/amp/autocast_mode.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Any from typing_extensions import deprecated diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py index 2b2fe32154b2..6722114e295b 100644 --- a/torch/cuda/__init__.py +++ b/torch/cuda/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs r""" This package adds support for CUDA tensor types. diff --git a/torch/cuda/_memory_viz.py b/torch/cuda/_memory_viz.py index 7d211fd3b8cb..2047ec4efb28 100644 --- a/torch/cuda/_memory_viz.py +++ b/torch/cuda/_memory_viz.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import pickle import sys import os diff --git a/torch/cuda/_sanitizer.py b/torch/cuda/_sanitizer.py index 89766ba8c1a4..bf72f277dd8a 100644 --- a/torch/cuda/_sanitizer.py +++ b/torch/cuda/_sanitizer.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs r""" This module introduces CUDA Sanitizer, a tool for detecting synchronization errors between kernels ran on different streams. diff --git a/torch/cuda/amp/autocast_mode.py b/torch/cuda/amp/autocast_mode.py index eb17d7a75e69..049ff41c590f 100644 --- a/torch/cuda/amp/autocast_mode.py +++ b/torch/cuda/amp/autocast_mode.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import functools from typing import Any from typing_extensions import deprecated diff --git a/torch/cuda/amp/common.py b/torch/cuda/amp/common.py index c4e8c1cc99b0..30ccaeede8d9 100644 --- a/torch/cuda/amp/common.py +++ b/torch/cuda/amp/common.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from importlib.util import find_spec import torch diff --git a/torch/cuda/graphs.py b/torch/cuda/graphs.py index 9d9df283ced6..78c572a1822d 100644 --- a/torch/cuda/graphs.py +++ b/torch/cuda/graphs.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import gc import typing diff --git a/torch/cuda/jiterator.py b/torch/cuda/jiterator.py index 1be552555945..294670f8819e 100644 --- a/torch/cuda/jiterator.py +++ b/torch/cuda/jiterator.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import re from typing import Callable, List diff --git a/torch/cuda/memory.py b/torch/cuda/memory.py index 0f12395ac778..9634d1c0d80b 100644 --- a/torch/cuda/memory.py +++ b/torch/cuda/memory.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs r"""This package adds support for device memory management implemented in CUDA.""" import collections diff --git a/torch/cuda/nccl.py b/torch/cuda/nccl.py index 67d528771215..4c28443c9e29 100644 --- a/torch/cuda/nccl.py +++ b/torch/cuda/nccl.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import collections import warnings from typing import Optional, Sequence, Union diff --git a/torch/cuda/nvtx.py b/torch/cuda/nvtx.py index 4b902c0c6d4d..195509687905 100644 --- a/torch/cuda/nvtx.py +++ b/torch/cuda/nvtx.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs r"""This package adds support for NVIDIA Tools Extension (NVTX) used in profiling.""" from contextlib import contextmanager diff --git a/torch/cuda/profiler.py b/torch/cuda/profiler.py index 51c8aa46f714..7e5dc9bab8de 100644 --- a/torch/cuda/profiler.py +++ b/torch/cuda/profiler.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import contextlib import tempfile diff --git a/torch/cuda/random.py b/torch/cuda/random.py index 1cf33114d17b..b736c9d959d8 100644 --- a/torch/cuda/random.py +++ b/torch/cuda/random.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from typing import Iterable, List, Union import torch diff --git a/torch/cuda/streams.py b/torch/cuda/streams.py index d36121381586..89271b588711 100644 --- a/torch/cuda/streams.py +++ b/torch/cuda/streams.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import ctypes import torch diff --git a/torch/cuda/tunable.py b/torch/cuda/tunable.py index 0f7e0a1f3725..8b387102b43d 100644 --- a/torch/cuda/tunable.py +++ b/torch/cuda/tunable.py @@ -138,12 +138,12 @@ __all__ = [ def enable(val: bool = True) -> None: r"""This is the big on/off switch for all TunableOp implementations.""" - torch._C._cuda_tunableop_enable(val) + torch._C._cuda_tunableop_enable(val) # type: ignore[attr-defined] def is_enabled() -> bool: r"""Returns whether the TunableOp feature is enabled.""" - return torch._C._cuda_tunableop_is_enabled() + return torch._C._cuda_tunableop_is_enabled() # type: ignore[attr-defined] def tuning_enable(val: bool = True) -> None: @@ -152,12 +152,12 @@ def tuning_enable(val: bool = True) -> None: When enabled, if a tuned entry isn't found, run the tuning step and record the entry. """ - torch._C._cuda_tunableop_tuning_enable(val) + torch._C._cuda_tunableop_tuning_enable(val) # type: ignore[attr-defined] def tuning_is_enabled() -> bool: r"""Returns whether TunableOp implementations can be tuned.""" - return torch._C._cuda_tunableop_tuning_is_enabled() + return torch._C._cuda_tunableop_tuning_is_enabled() # type: ignore[attr-defined] def set_max_tuning_duration(duration: int) -> None: @@ -166,12 +166,12 @@ def set_max_tuning_duration(duration: int) -> None: If both max tuning duration and iterations are set, the smaller of the two will be honored. At minimum 1 tuning iteration will always be run. """ - torch._C._cuda_tunableop_set_max_tuning_duration(duration) + torch._C._cuda_tunableop_set_max_tuning_duration(duration) # type: ignore[attr-defined] def get_max_tuning_duration() -> int: r"""Get max time to spend tuning a given solution.""" - return torch._C._cuda_tunableop_get_max_tuning_duration() + return torch._C._cuda_tunableop_get_max_tuning_duration() # type: ignore[attr-defined] def set_max_tuning_iterations(iterations: int) -> None: @@ -180,12 +180,12 @@ def set_max_tuning_iterations(iterations: int) -> None: If both max tuning duration and iterations are set, the smaller of the two will be honored. At minimum 1 tuning iteration will always be run. """ - torch._C._cuda_tunableop_set_max_tuning_iterations(iterations) + torch._C._cuda_tunableop_set_max_tuning_iterations(iterations) # type: ignore[attr-defined] def get_max_tuning_iterations() -> int: r"""Get max iterations to spend tuning a given solution.""" - return torch._C._cuda_tunableop_get_max_tuning_iterations() + return torch._C._cuda_tunableop_get_max_tuning_iterations() # type: ignore[attr-defined] def set_filename(filename: str, insert_device_ordinal: bool = False) -> None: @@ -195,22 +195,22 @@ def set_filename(filename: str, insert_device_ordinal: bool = False) -> None: will be added to the given filename automatically. This can be used in a 1-process-per-gpu cenario to ensure all processes write to a separate file. """ - torch._C._cuda_tunableop_set_filename(filename, insert_device_ordinal) + torch._C._cuda_tunableop_set_filename(filename, insert_device_ordinal) # type: ignore[attr-defined] def get_filename() -> str: r"""Get the results filename.""" - return torch._C._cuda_tunableop_get_filename() + return torch._C._cuda_tunableop_get_filename() # type: ignore[attr-defined] def get_results() -> Tuple[str, str, str, float]: r"""Return all TunableOp results.""" - return torch._C._cuda_tunableop_get_results() + return torch._C._cuda_tunableop_get_results() # type: ignore[attr-defined] def get_validators() -> Tuple[str, str]: r"""Return the TunableOp validators.""" - return torch._C._cuda_tunableop_get_validators() + return torch._C._cuda_tunableop_get_validators() # type: ignore[attr-defined] def write_file_on_exit(val: bool) -> None: @@ -219,7 +219,7 @@ def write_file_on_exit(val: bool) -> None: This is useful as a final flush of your results to disk if your application terminates as result of normal operation or an error. Manual flushing of your results can be achieved by manually calling ``write_file()``.""" - torch._C._cuda_tunableop_write_file_on_exit(val) + torch._C._cuda_tunableop_write_file_on_exit(val) # type: ignore[attr-defined] def write_file(filename: Optional[str] = None) -> bool: @@ -229,7 +229,7 @@ def write_file(filename: Optional[str] = None) -> bool: """ if filename is None: filename = get_filename() - return torch._C._cuda_tunableop_write_file(filename) + return torch._C._cuda_tunableop_write_file(filename) # type: ignore[attr-defined] def read_file(filename: Optional[str] = None) -> bool: @@ -239,4 +239,4 @@ def read_file(filename: Optional[str] = None) -> bool: """ if filename is None: filename = get_filename() - return torch._C._cuda_tunableop_read_file(filename) + return torch._C._cuda_tunableop_read_file(filename) # type: ignore[attr-defined] diff --git a/torch/distributed/__init__.py b/torch/distributed/__init__.py index 3e7dce97b54c..b8e911c8738c 100644 --- a/torch/distributed/__init__.py +++ b/torch/distributed/__init__.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs import os import sys from enum import Enum diff --git a/torch/distributed/_composable/checkpoint_activation.py b/torch/distributed/_composable/checkpoint_activation.py index 8accef6afc34..6716f43a74a0 100644 --- a/torch/distributed/_composable/checkpoint_activation.py +++ b/torch/distributed/_composable/checkpoint_activation.py @@ -1,3 +1,4 @@ +# mypy: allow-untyped-defs from contextlib import contextmanager, nullcontext from typing import Any, Tuple