Flip default value for mypy disallow_untyped_defs [5/11] (#127842)

See #127836 for details.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/127842
Approved by: https://github.com/oulgen

torch/distributed/_composable/contract.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import uuid
 from collections import OrderedDict
 from functools import wraps

torch/distributed/_composable/fsdp/_fsdp_api.py

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

torch/distributed/_composable/fsdp/_fsdp_common.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import math
 import traceback
 

torch/distributed/_composable/fsdp/_fsdp_param.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import itertools
 from dataclasses import dataclass, field
 from enum import auto, Enum

torch/distributed/_composable/fsdp/_fsdp_param_group.py

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

torch/distributed/_composable/fsdp/_fsdp_state.py

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

torch/distributed/_composable/fsdp/fully_shard.py

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

torch/distributed/_composable/replicate.py

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

torch/distributed/_cuda_p2p/__init__.py

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

torch/distributed/_functional_collectives.py

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

torch/distributed/_functional_collectives_impl.py

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

torch/distributed/_shard/api.py

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

torch/distributed/_shard/common_op_utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import torch
 from torch.utils import _pytree as pytree
 from typing import Optional

torch/distributed/_shard/metadata.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from dataclasses import dataclass
 from typing import List, Union, Optional
 from functools import reduce

torch/distributed/_shard/op_registry_utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import functools
 from inspect import signature
 from .common_op_utils import _basic_validation

torch/distributed/_shard/sharded_optim/api.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from typing import List, Union, Mapping, Dict, Any
 
 import torch.optim as optim

torch/distributed/_shard/sharded_tensor/__init__.py

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

torch/distributed/_shard/sharded_tensor/_ops/_common.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import functools
 from torch.distributed._shard.sharded_tensor import (
     _sharded_op_impl,

torch/distributed/_shard/sharded_tensor/_ops/binary_cmp.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import torch
 import torch.distributed as dist
 import torch.distributed.distributed_c10d as distributed_c10d

torch/distributed/_shard/sharded_tensor/_ops/init.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import torch
 import torch.distributed._shard.sharded_tensor as sharded_tensor
 from torch.distributed._shard.sharded_tensor import (

torch/distributed/_shard/sharded_tensor/_ops/misc_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import torch
 from torch.distributed._shard.sharded_tensor import (
     _sharded_op_impl,

torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import copy
 import torch
 from torch.distributed._shard.sharded_tensor import (

torch/distributed/_shard/sharded_tensor/api.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from __future__ import annotations  # type: ignore[attr-defined]
 from dataclasses import dataclass
 from typing import (

torch/distributed/_shard/sharded_tensor/metadata.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import List

torch/distributed/_shard/sharded_tensor/reshard.py

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

torch/distributed/_shard/sharded_tensor/shard.py

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

torch/distributed/_shard/sharded_tensor/utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import collections.abc
 import copy
 from typing import Optional, List, Sequence, TYPE_CHECKING

torch/distributed/_shard/sharding_spec/_internals.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from typing import List, Optional, Tuple
 
 from torch.distributed._shard.metadata import ShardMetadata

torch/distributed/_shard/sharding_spec/api.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 import functools

torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from dataclasses import dataclass
 import torch
 import torch.distributed._shard.sharded_tensor.metadata as sharded_tensor_meta

torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py

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

torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding.py

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

torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding_bag.py

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

torch/distributed/_spmd/api.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from abc import ABC, abstractmethod
 from contextlib import contextmanager, nullcontext
 from copy import copy

torch/distributed/_spmd/batch_dim_utils.py

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

torch/distributed/_spmd/comm_tensor.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from dataclasses import dataclass
 from functools import partial
 from typing import Any, List, Optional, Tuple

torch/distributed/_spmd/config.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import logging
 import sys
 from types import ModuleType

torch/distributed/_spmd/data_parallel.py

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

torch/distributed/_spmd/distribute.py

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

torch/distributed/_spmd/experimental_ops.py

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

torch/distributed/_spmd/graph_optimization.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 # Owner(s): ["oncall: distributed"]
 import collections
 import itertools

torch/distributed/_spmd/iter_graph_module.py

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

torch/distributed/_state_dict_utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import copy
 import io
 import math

torch/distributed/_tensor/__init__.py

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

torch/distributed/_tensor/_collective_utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import logging
 import math
 from dataclasses import dataclass

torch/distributed/_tensor/_op_schema.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from dataclasses import dataclass
 from functools import cached_property
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

torch/distributed/_tensor/_redistribute.py

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

torch/distributed/_tensor/_sharding_prop.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from functools import lru_cache
 from itertools import chain
 from typing import Callable, cast, Dict, List, Optional, Sequence, Tuple, Union

torch/distributed/_tensor/_tp_conv.py

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

torch/distributed/_tensor/api.py

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

torch/distributed/_tensor/debug/__init__.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from torch.distributed._tensor.api import DTensor
 
 from torch.distributed._tensor.debug.comm_mode import CommDebugMode

torch/distributed/_tensor/debug/_op_coverage.py

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

torch/distributed/_tensor/debug/comm_mode.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from collections import defaultdict
 from typing import Any, Dict
 

torch/distributed/_tensor/debug/visualize_sharding.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from typing import List, Sequence, Tuple
 
 import numpy as np

torch/distributed/_tensor/examples/checkpoint_example.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 """
 The following example contains a simple MLP model that uses
 different DTensor layouts, and use the checkpointing API to

torch/distributed/_tensor/examples/convnext_example.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 """
 The following example demonstrates how to train a ConvNeXt model
 with intermediate activations sharded across mutliple GPUs via DTensor

torch/distributed/_tensor/examples/torchrec_sharding_example.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 """
 The following example demonstrates how to represent torchrec's embedding
 sharding with the DTensor API.

torch/distributed/_tensor/experimental/__init__.py

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

torch/distributed/_tensor/experimental/local_map.py

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

torch/distributed/_tensor/experimental/tp_transform.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import copy
 import operator
 from typing import Any, cast, Dict, List, Optional, Sequence, Tuple

torch/distributed/_tensor/ops/embedding_ops.py

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

torch/distributed/_tensor/ops/math_ops.py

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

torch/distributed/_tensor/ops/tensor_ops.py

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

torch/distributed/_tensor/ops/utils.py

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

torch/distributed/_tensor/ops/view_ops.py

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

torch/distributed/_tensor/placement_types.py

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

torch/distributed/_tensor/random.py

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

torch/distributed/_tools/memory_tracker.py

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

torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import warnings
 from enum import auto, Enum
 from functools import partial

torch/distributed/algorithms/_comm_hooks/default_hooks.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import functools
 import torch
 import torch.distributed as dist

torch/distributed/algorithms/_optimizer_overlap/optimizer_overlap.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from abc import ABC, abstractmethod
 import inspect
 from typing import Dict, Type

torch/distributed/algorithms/_quantization/quantization.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import functools
 import torch
 import torch.distributed as dist

torch/distributed/algorithms/ddp_comm_hooks/__init__.py

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

torch/distributed/algorithms/ddp_comm_hooks/ddp_zero_hook.py

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

torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from typing import Any, Callable, cast, Tuple
 
 import torch

torch/distributed/algorithms/ddp_comm_hooks/optimizer_overlap_hooks.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from typing import Any, Callable, List, no_type_check
 
 import torch

torch/distributed/algorithms/ddp_comm_hooks/post_localSGD_hook.py

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

torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py

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

torch/distributed/algorithms/ddp_comm_hooks/quantization_hooks.py

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

torch/distributed/algorithms/join.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import warnings
 from abc import ABC, abstractmethod
 from types import TracebackType

torch/distributed/algorithms/model_averaging/averagers.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import warnings
 from abc import ABC, abstractmethod
 from typing import Union, Iterable, Dict

torch/distributed/algorithms/model_averaging/hierarchical_model_averager.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 # Copyright 2022 Cruise LLC
 import logging
 import warnings

torch/distributed/algorithms/model_averaging/utils.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 # flake8: noqa C101
 import itertools
 from typing import Union, Iterable, Dict, Iterator

torch/distributed/argparse_util.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# mypy: allow-untyped-defs
 
 # Copyright (c) Facebook, Inc. and its affiliates.
 # All rights reserved.

torch/distributed/autograd/__init__.py

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

torch/distributed/benchmarks/benchmark_ddp_rpc.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import argparse
 import io
 import os

torch/distributed/c10d_logger.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+# mypy: allow-untyped-defs
 
 # Copyright (c) Facebook, Inc. and its affiliates.
 # All rights reserved.

torch/distributed/checkpoint/api.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import traceback as tb
 from typing import Any, Dict, Tuple
 

torch/distributed/checkpoint/default_planner.py

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

torch/distributed/checkpoint/examples/async_checkpointing_example.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 # Owner(s): ["oncall: distributed"]
 
 import os

torch/distributed/checkpoint/examples/fsdp_checkpoint_example.py

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

torch/distributed/checkpoint/examples/stateful_example.py

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

torch/distributed/checkpoint/filesystem.py

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

torch/distributed/checkpoint/format_utils.py

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

torch/distributed/checkpoint/logger.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import functools
 import time
 from typing import Any, Callable, Dict, List, TypeVar

torch/distributed/checkpoint/metadata.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import os
 from dataclasses import dataclass, field
 from enum import Enum

torch/distributed/checkpoint/planner_helpers.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from typing import Any, cast, List
 
 import torch

torch/distributed/checkpoint/resharding.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 from typing import List, Tuple
 
 from torch.distributed.checkpoint.metadata import ChunkStorageMetadata

torch/distributed/checkpoint/state_dict.py

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

torch/distributed/checkpoint/state_dict_loader.py

@@ -1,3 +1,4 @@
+# mypy: allow-untyped-defs
 import os
 import warnings
 from typing import Any, cast, Dict, Optional, Set, Union