Add None return type to init (pytorch#132335)

Pull Request resolved: pytorch#132335
Approved by: https://github.com/albanD

Author: oulgen

test/dynamo/test_minifier.py

@@ -119,7 +119,7 @@ def test_cpu_cuda_module_after_dynamo(self):
         backend_name = "relu_compile_error_TESTING_ONLY"
         run_code = f"""\
 class CpuCudaModule(torch.nn.Module):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.m_x = torch.nn.Linear(20, 20).cuda()
         self.m_y = torch.nn.Linear(20, 20)
@@ -149,7 +149,7 @@ def inner(x1, y1):
             res.minifier_module(),
             """\
 class Repro(torch.nn.Module):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.G__mod___m_x = Linear(in_features=20, out_features=20, bias=True).cuda()
         self.G__mod___m_y = Linear(in_features=20, out_features=20, bias=True)
@@ -204,7 +204,7 @@ def inner(x):
             res.repro_module(),
             """\
 class Repro(torch.nn.Module):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
     def forward(self, x_19):

test/inductor/test_minifier.py

@@ -122,7 +122,7 @@ def inner(x):
             res.repro_module(),
             """\
 class Repro(torch.nn.Module):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
     def forward(self, arg0_1):
@@ -138,7 +138,7 @@ def forward(self, arg0_1):
             res.repro_module(),
             """\
 class Repro(torch.nn.Module):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
     def forward(self, arg0_1):

torch/_classes.py

@@ -19,7 +19,7 @@ def __getattr__(self, attr):
 class _Classes(types.ModuleType):
     __file__ = "_classes.py"
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("torch.classes")
 
     def __getattr__(self, name):

torch/_decomp/decompositions_for_rng.py

@@ -71,7 +71,7 @@ class PhiloxState:
     trace time.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.reset()
 
     def reset(self):

torch/_dynamo/backends/distributed.py

@@ -247,7 +247,7 @@ def run_node(self, n: Node) -> Any:
             # This gives us the appropriately strided outputs here which will reflect runtime strides.
 
             class FakeifyFirstAOTInvocationGuard:
-                def __init__(self):
+                def __init__(self) -> None:
                     self.tc = torch._guards.TracingContext.try_get()
                     assert self.tc
                     torch._guards.TracingContext.try_get().fakify_first_call = True

torch/_dynamo/code_context.py

@@ -5,7 +5,7 @@
 
 
 class CodeContextDict:
-    def __init__(self):
+    def __init__(self) -> None:
         self.code_context = ExactWeakKeyDictionary()
 
     def has_context(self, code: types.CodeType):

torch/_dynamo/debug_utils.py

@@ -170,7 +170,7 @@ def convert(gm):
             """
             from torch.nn import *
             class Repro(torch.nn.Module):
-                def __init__(self):
+                def __init__(self) -> None:
                     super().__init__()
             """
         )
@@ -491,7 +491,7 @@ def _mk_defaulter(d: T) -> Callable[[Optional[T]], T]:
 
 
 class NopInputReader:
-    def __init__(self):
+    def __init__(self) -> None:
         self.total = 0
 
     def storage(self, storage_hash, nbytes, *, device=None, dtype_hint=None):

torch/_dynamo/eval_frame.py

@@ -497,7 +497,7 @@ def _fn(*args, **kwargs):
                         wrapper function.
 
                         >> class CallableClass:
-                        >>     def __init__(self):
+                        >>     def __init__(self) -> None:
                         >>         super().__init__()
                         >>         self.relu = torch.nn.ReLU()
                         >>
@@ -578,7 +578,7 @@ def __reduce__(self):
 
 
 class RunOnlyContext(_TorchDynamoContext):
-    def __init__(self):
+    def __init__(self) -> None:
         # cudagraph trees relies on generation increment
         def on_enter():
             torch._dynamo.mutation_guard.GenerationTracker.generation += 1
@@ -590,7 +590,7 @@ def __reduce__(self):
 
 
 class DisableContext(_TorchDynamoContext):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(callback=None)
 
     def __call__(self, fn):

torch/_dynamo/exc.py

@@ -74,7 +74,7 @@ def __init__(self, name):
 
 
 class ResetRequired(TorchDynamoException):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             textwrap.dedent(
                 """

torch/_dynamo/profiler.py

@@ -92,7 +92,7 @@ def print_missing(stack):
 class Profiler:
     unique_graphs = 0
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.prof = torch.profiler.profile(
             activities=[torch.profiler.ProfilerActivity.CPU],
             with_stack=should_print_missing(),

torch/_dynamo/variables/base.py

@@ -70,7 +70,7 @@ class MutableLocal(MutableLocalBase):
     state.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(MutableLocalSource.Local)
 
     def __hash__(self):

torch/_dynamo/variables/builder.py

@@ -274,7 +274,7 @@ def __eq__(self, other):
 
 
 class BackwardStateGraphArg(GraphArg):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             source=None,
             _example=BackwardState(),
@@ -2646,7 +2646,7 @@ class SourcelessBuilder:
     if/else type->VariableTracker trees that were cropping up all over dynamo.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         raise AssertionError("Use SourcelessBuilder.create()")
 
     @staticmethod

torch/_export/db/examples/class_method.py

@@ -10,7 +10,7 @@ class ClassMethod(torch.nn.Module):
     def method(cls, x):
         return x + 1
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.linear = torch.nn.Linear(4, 2)
 

torch/_export/db/examples/cond_branch_class_method.py

@@ -26,7 +26,7 @@ class CondBranchClassMethod(torch.nn.Module):
     NOTE: If the `pred` is test on a dim with batch size < 2, it will be specialized.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.subm = MySubModule()
 

torch/_export/db/examples/model_attr_mutation.py

@@ -8,7 +8,7 @@ class ModelAttrMutation(torch.nn.Module):
     Attribute mutation is not supported.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.attr_list = [torch.randn(3, 2), torch.randn(3, 2)]
 

torch/_export/db/examples/scalar_output.py

@@ -11,7 +11,7 @@ class ScalarOutput(torch.nn.Module):
     Returning scalar values from the graph is supported, in addition to Tensor
     outputs. Symbolic shapes are captured and rank is specialized.
     """
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
 
     def forward(self, x):

torch/_export/db/examples/specialized_attribute.py

@@ -11,7 +11,7 @@ class SpecializedAttribute(torch.nn.Module):
     Model attributes are specialized.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.a = "moo"
         self.b = 4

torch/_export/passes/lift_constants_pass.py

@@ -24,7 +24,7 @@ class ConstantAttrMap(collections.abc.MutableMapping):
     if that's the case).
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         # Underlying dict that we use to implement this mapping.
         self._constant_attrs: Dict[
             Union[int, torch.Tensor, FakeScriptObject], List[Any]

torch/_export/serde/serialize.py

@@ -1413,7 +1413,7 @@ class Result:
         constants: Dict[str, Union[torch.Tensor, FakeScriptObject, torch.ScriptObject]]
         example_inputs: Optional[Tuple[Tuple[torch.Tensor, ...], Dict[str, Any]]]
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.serialized_name_to_node: Dict[str, torch.fx.Node] = {}
         self.serialized_name_to_meta: Dict[str, MetaType] = {}
         self.graph = torch.fx.Graph()

torch/_functorch/_aot_autograd/schemas.py

@@ -602,7 +602,7 @@ class SubclassMeta:
     # Optional field because we don't compute for inference graphs
     grad_input_metas: Optional[List[Union[int, SubclassCreationMeta]]] = None
 
-    def __init__(self):
+    def __init__(self) -> None:
         # The fields in this class get set after its construction.
         pass
 

torch/_functorch/aot_autograd.py

@@ -878,7 +878,7 @@ def functional_call(named_params, named_buffers, *args, **kwargs):
     )
 
     class AOTModule(nn.Module):
-        def __init__(self):
+        def __init__(self) -> None:
             super().__init__()
             self.orig_module = mod
 

torch/_functorch/autograd_function.py

@@ -30,7 +30,7 @@
 # We do this by using creating a custom HigherOrderOperator that only functorch
 # dispatches specially.
 class CustomFunctionHigherOrderOperator(HigherOrderOperator):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("custom_function_call")
 
     def __call__(self, autograd_function, *args, **kwargs):
@@ -713,7 +713,7 @@ def new_forward(ctx, *args, **kwargs):
 
 
 class AutogradFunctionApply(HigherOrderOperator):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("autograd_function_apply")
 
     def __call__(self, fwd, bwd, *fwd_args, **fwd_kwargs):

torch/_guards.py

@@ -427,7 +427,7 @@ def __eq__(self, other):
 
 
 class ModuleContext(Checkpointable[ModuleContextCheckpointState]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.nn_modules: Dict[str, Any] = {}
 
     def copy_graphstate(self):
@@ -476,7 +476,7 @@ class GlobalContext(Checkpointable[GlobalContextCheckpointState]):
         "autocast_cache_enabled",
     }
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.global_state: Dict[str, Tuple[Callable, ...]] = {}
 
     def copy_graphstate(self):
@@ -544,7 +544,7 @@ def remove_guards_with_source(self, source):
 
 
 class GuardsContext(Checkpointable[GuardsCheckpointState]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.dynamo_guards: GuardsSet = GuardsSet()
         self.aotautograd_guards: List[GuardEnvExpr] = []
 

torch/_higher_order_ops/auto_functionalize.py

@@ -54,7 +54,7 @@ class AutoFunctionalized(HigherOrderOperator):
     underscore is to prevent collisions with kwarg names in **kwargs.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("auto_functionalized")
 
     def __call__(

torch/_higher_order_ops/effects.py

@@ -55,7 +55,7 @@ class WithEffects(HigherOrderOperator):
     per "effect type", which are enumerated in the _EffectType enum.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("with_effects")
 
     def __call__(

torch/_higher_order_ops/flex_attention.py

@@ -38,7 +38,7 @@ def __torch_function__(self, func, types, args, kwargs=None):
 
 
 class FlexAttentionHOP(HigherOrderOperator):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("flex_attention")
 
     def __call__(
@@ -74,7 +74,7 @@ def __call__(
 
 
 class FlexAttentionBackwardHOP(HigherOrderOperator):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("flex_attention_backward")
 
     def __call__(

torch/_higher_order_ops/out_dtype.py

@@ -45,7 +45,7 @@ class OutDtypeOperator(HigherOrderOperator):
         3. Cast the output to `out_dtype`
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("out_dtype")
         # TODO(ydwu4): Subclassing HigherOrderOperator causes __module__ to
         # become different (torch._higher_order_ops.out_dtype) which will result

torch/_higher_order_ops/triton_kernel_wrap.py

@@ -519,7 +519,7 @@ def identify_mutated_tensors(kernel, kwargs):
 
 # Used for wrapping a Triton Kernel
 class TritonKernelWrapperMutation(HigherOrderOperator):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("triton_kernel_wrapper_mutation")
 
 
@@ -528,7 +528,7 @@ def __init__(self):
 
 # Used for wrapping a Triton Kernel in a functional manner
 class TritonKernelWrapperFunctional(HigherOrderOperator):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("triton_kernel_wrapper_functional")
 
 

torch/_higher_order_ops/while_loop.py

@@ -18,7 +18,7 @@
 
 
 class WhileLoopOp(HigherOrderOperator):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__("while_loop")
 
     def __call__(