[BE] mypy: disallow untyped decorators (#131428)

Untyped decorators strip the types from their decorated function so even if the underlying function is fully typed then callers to it don't get any benefit from type annotations.

Step 1 - Enable the error and override in all the offending files.

#131429

Pull Request resolved: https://github.com/pytorch/pytorch/pull/131428
Approved by: https://github.com/justinchuby, https://github.com/oulgen

mypy.ini

@@ -12,6 +12,7 @@ show_error_codes = True
 show_column_numbers = True
 check_untyped_defs = True
 disallow_untyped_defs = True
+disallow_untyped_decorators = True
 follow_imports = normal
 local_partial_types = True
 enable_error_code = possibly-undefined

test/test_torch.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Owner(s): ["module: tests"]
 

torch/_decomp/decompositions.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import itertools

torch/_decomp/decompositions_for_jvp.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 from typing import Callable, Dict, List, Optional, Tuple

torch/_decomp/decompositions_for_rng.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from collections import defaultdict

torch/_dynamo/convert_frame.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from __future__ import annotations
 
 import collections

torch/_dynamo/output_graph.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import collections
 import contextlib

torch/_dynamo/variables/script_object.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from typing import Dict

torch/_export/__init__.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 import dataclasses

torch/_functorch/apis.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # NOTE: We allow Dynamo to see this file (via torch/_dynamo/trace_rules.py) so that it can
 #       trace through functorch transforms.

torch/_functorch/batch_norm_replacement.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch.nn as nn
 from torch._functorch.utils import exposed_in

torch/_functorch/functional_call.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from collections import Counter
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

torch/_higher_order_ops/auto_functionalize.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union

torch/_higher_order_ops/effects.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from enum import Enum
 from typing import Any, Dict, Optional, Tuple, Union

torch/_higher_order_ops/flex_attention.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Callable, Tuple, Union
 

torch/_higher_order_ops/out_dtype.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 
 import torch

torch/_inductor/codecache.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from __future__ import annotations
 

torch/_inductor/compile_fx.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import contextlib
 import functools

torch/_inductor/decomposition.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import logging

torch/_inductor/freezing.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from __future__ import annotations
 

torch/_inductor/fx_passes/post_grad.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import itertools

torch/_inductor/fx_passes/quantization.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 import functools

torch/_inductor/jagged_lowerings.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Optional, Tuple, Union
 

torch/_inductor/kernel/conv.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from __future__ import annotations
 

torch/_inductor/lowering.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import itertools

torch/_inductor/mkldnn_lowerings.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from typing import List, Optional

torch/_inductor/pattern_matcher.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 """
 # Inductor Pattern Matcher
 

torch/_inductor/runtime/triton_helpers.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import triton
 import triton.language as tl

torch/_inductor/scheduler.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: disallow-untyped-defs
 from __future__ import annotations
 

torch/_library/custom_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import logging

torch/_library/infer_schema.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import typing

torch/_meta_registrations.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 from enum import Enum

torch/_refs/__init__.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import builtins
 import collections

torch/_refs/_conversions.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import torch._prims_common as utils

torch/_refs/fft.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 import math
 
 from typing import Iterable, List, Literal, NamedTuple, Optional, Sequence, Tuple, Union

torch/_refs/linalg/__init__.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from functools import partial
 

torch/_refs/nn/functional/__init__.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 from functools import wraps

torch/_refs/special/__init__.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 from typing import Optional, Union

torch/_subclasses/fake_tensor.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from __future__ import annotations
 
 import atexit

torch/ao/nn/quantized/dynamic/modules/rnn.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import numbers
 import warnings

torch/ao/nn/quantized/modules/embedding_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import torch.nn as nn

torch/ao/nn/quantized/modules/linear.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from collections.abc import Iterable
 import torch

torch/ao/nn/sparse/quantized/linear.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Optional
 

torch/ao/ns/fx/utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import enum
 import operator

torch/ao/quantization/experimental/adaround_fake_quantize.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Tuple
 

torch/ao/quantization/fake_quantize.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """Implements modules  used to perform fake quantization."""
 

torch/ao/quantization/fx/_decomposed.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import math
 from typing import Optional, Tuple

torch/ao/quantization/observer.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 """
 This module implements observers which are used to collect statistics about

torch/ao/quantization/quantizer/xnnpack_quantizer_utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import itertools
 import operator

torch/backends/_nnapi/prepare.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Optional
 

torch/distributed/_composable/checkpoint_activation.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from contextlib import contextmanager, nullcontext
 from typing import Any, ContextManager, Dict, Optional, Tuple

torch/distributed/_composable/fsdp/_fsdp_collectives.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from typing import List, NamedTuple, Optional, Tuple, Union
 
 import torch

torch/distributed/_composable/fsdp/_fsdp_state.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 import logging

torch/distributed/_composable/fsdp/fully_shard.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import functools
 from typing import Any, cast, Iterable, List, NoReturn, Optional, Union

torch/distributed/_composable/fully_shard.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from typing import Callable, Iterable, Optional, Union
 from typing_extensions import deprecated
 

torch/distributed/_composable/replicate.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import weakref
 from typing import Any, cast, Dict, Iterable, List, NoReturn, Optional, Set, Tuple

torch/distributed/_spmd/experimental_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 from typing import cast, List, Optional, Sequence, Tuple

torch/distributed/_symmetric_memory/__init__.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 import socket
 import uuid
 

torch/distributed/_tensor/api.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import inspect

torch/distributed/_tensor/ops/conv_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 from typing import List

torch/distributed/_tensor/ops/embedding_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor

torch/distributed/_tensor/ops/experimental_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 

torch/distributed/_tensor/ops/math_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import math

torch/distributed/_tensor/ops/matrix_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 # implement matrix related ops for distributed tensor
 

torch/distributed/_tensor/ops/random_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Meta Platforms, Inc. and affiliates
 import torch
 from torch.distributed._tensor._op_schema import (

torch/distributed/_tensor/ops/tensor_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 from typing import cast, List, Optional, Sequence, Tuple

torch/distributed/_tensor/ops/view_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 # Copyright (c) Meta Platforms, Inc. and affiliates
 from dataclasses import dataclass

torch/distributed/checkpoint/state_dict_loader.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import os
 import warnings

torch/distributed/checkpoint/state_dict_saver.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import os

torch/distributed/optim/optimizer.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import logging
 from collections import defaultdict

torch/distributed/rpc/api.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 
 import collections

torch/export/_trace.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import dataclasses
 import functools

torch/export/exported_program.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import contextlib
 import copy

torch/fx/_lazy_graph_module.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from contextlib import contextmanager
 

torch/fx/_symbolic_trace.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import builtins
 import copy

torch/fx/experimental/graph_gradual_typechecker.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from functools import reduce
 import torch

torch/fx/experimental/migrate_gradual_types/constraint_generator.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import operator

torch/fx/experimental/proxy_tensor.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Copyright (c) Facebook, Inc. and its affiliates.
 # All rights reserved.
 #

torch/fx/experimental/rewriter.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import ast
 import inspect

torch/fx/graph.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from collections import defaultdict
 from .node import Node, Argument, Target, map_arg, _type_repr, _get_qualified_name

torch/fx/graph_module.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import contextlib
 import copy

torch/fx/interpreter.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from .graph_module import GraphModule
 from ._lazy_graph_module import _make_graph_module

torch/fx/node.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # Nodes represent a definition of a value in our graph of operators.
 from typing import TYPE_CHECKING, Union, Callable, Any, Tuple, List, Optional, Dict, Set
 from ._compatibility import compatibility

torch/fx/operator_schemas.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 import inspect

torch/fx/passes/graph_manipulation.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Any, Dict, List, NamedTuple, Optional
 

torch/fx/passes/infra/pass_manager.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 import logging

torch/fx/passes/operator_support.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import abc
 import typing as t

torch/fx/passes/param_fetch.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from torch.fx.graph_module import GraphModule
 from typing import Any, Callable, Dict, List, Tuple, Type
 import torch

torch/fx/passes/runtime_assert.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import logging
 import operator

torch/fx/passes/split_module.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import inspect
 from typing import Any, Callable, Dict, List, Optional, Set

torch/fx/passes/split_utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 from dataclasses import dataclass, field

torch/fx/passes/splitter_base.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import argparse
 import copy

torch/fx/passes/tools_common.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import List, Tuple, Union, Dict, Any, Set, Mapping, Optional
 import collections

torch/fx/passes/utils/common.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from typing import Dict, Tuple
 

torch/fx/passes/utils/fuser_utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import copy
 from queue import SimpleQueue

torch/fx/passes/utils/source_matcher_utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 from dataclasses import dataclass, field
 from torch.fx.graph import Graph

torch/fx/subgraph_rewriter.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 from .graph_module import GraphModule
 from .graph import Graph
 from .node import Node

torch/fx/traceback.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import traceback
 from contextlib import contextmanager

torch/jit/_decompositions.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import torch
 from torch import Tensor

torch/masked/_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-decorators
 # mypy: allow-untyped-defs
 import warnings
 from typing import Any, List, Optional, Tuple, TYPE_CHECKING, Union