[BE] [3/3] Rewrite super() calls in test (pytorch#94592)

Rewrite Python built-in class `super()` calls. Only non-semantic changes should be applied.

- pytorch#94587
- pytorch#94588
- pytorch#94592

Also, methods with only a `super()` call are removed:

```diff
class MyModule(nn.Module):
-   def __init__(self):
-       super().__init__()
-
    def forward(self, ...):
        ...
```

Some cases that change the semantics should be kept unchanged. E.g.:

https://github.com/pytorch/pytorch/blob/f152a79be9612b824e1672b8f8cb88a414ce4c12/caffe2/python/net_printer.py#L184-L190

https://github.com/pytorch/pytorch/blob/f152a79be9612b824e1672b8f8cb88a414ce4c12/test/test_jit_fuser_te.py#L2628-L2635

Pull Request resolved: pytorch#94592
Approved by: https://github.com/ezyang, https://github.com/seemethere

Author: XuehaiPan

test/bottleneck_test/test_cuda.py

@@ -6,7 +6,7 @@
 
 class Model(nn.Module):
     def __init__(self):
-        super(Model, self).__init__()
+        super().__init__()
         self.linear = nn.Linear(20, 20)
 
     def forward(self, input):

test/cpp/jit/test_exception.cpp

@@ -113,7 +113,7 @@ TEST(TestException, TestCustomException) {
   py::exec(R"PY(
   class SimpleValueError(ValueError):
     def __init__(self, message):
-      super(SimpleValueError, self).__init__(message)
+      super().__init__(message)
   )PY");
 
   std::string pythonCode = R"PY(

test/cpp/jit/test_lite_interpreter.cpp

@@ -1157,7 +1157,7 @@ TEST(RunTimeTest, ParseOperator) {
 
   // class Add(torch.nn.Module):
   //     def __init__(self):
-  //         super(Add, self).__init__()
+  //         super().__init__()
 
   //     def forward(self, a, b):
   //         return a + b

test/cpp/jit/tests_setup.py

@@ -26,7 +26,7 @@ class EvalModeForLoadedModule(FileSetup):
     def setup(self):
         class Model(torch.jit.ScriptModule):
             def __init__(self):
-                super(Model, self).__init__()
+                super().__init__()
                 self.dropout = torch.nn.Dropout(0.1)
 
             @torch.jit.script_method

test/cpp/lite_interpreter_runtime/test_lite_interpreter_runtime.cpp

@@ -21,21 +21,21 @@ TEST(RunTimeTest, LoadAndForward) {
   //  sequence.ptl source code:
   //  class A(torch.nn.Module):
   //    def __init__(self):
-  //      super(A, self).__init__()
+  //      super().__init__()
   //
   //    def forward(self, x):
   //      return x + 1
   //
   //  class B(torch.nn.Module):
   //    def __init__(self):
-  //      super(B, self).__init__()
+  //      super().__init__()
   //
   //    def forward(self, x):
   //      return x + 2
   //
   //  class C(torch.nn.Module):
   //    def __init__(self):
-  //      super(C, self).__init__()
+  //      super().__init__()
   //      self.A0 = A()
   //      self.B0 = B()
   //

test/cpp_api_parity/sample_module.py

@@ -13,7 +13,7 @@
 
 class SampleModule(torch.nn.Module):
     def __init__(self, has_parity, has_submodule):
-        super(SampleModule, self).__init__()
+        super().__init__()
         self.has_parity = has_parity
         if has_submodule:
             self.submodule = SampleModule(self.has_parity, False)

test/create_dummy_torchscript_model.py

@@ -7,7 +7,7 @@
 class NeuralNetwork(nn.Module):
 
     def __init__(self):
-        super(NeuralNetwork, self).__init__()
+        super().__init__()
         self.flatten = nn.Flatten()
         self.linear_relu_stack = nn.Sequential(
             nn.Linear(28 * 28, 512),

test/custom_backend/backend.py

@@ -43,9 +43,6 @@ class Model(torch.nn.Module):
     and executing in C++.
     """
 
-    def __init__(self):
-        super(Model, self).__init__()
-
     def forward(self, a, b):
         return (a + b, a - b)
 

test/custom_operator/model.py

@@ -19,7 +19,7 @@ def get_custom_op_library_path():
 
 class Model(torch.jit.ScriptModule):
     def __init__(self):
-        super(Model, self).__init__()
+        super().__init__()
         self.p = torch.nn.Parameter(torch.eye(5))
 
     @torch.jit.script_method

test/distributed/_composable/test_replicate.py

@@ -14,7 +14,7 @@
 
 class Net(nn.Module):
     def __init__(self):
-        super(Net, self).__init__()
+        super().__init__()
         self.fc1 = nn.Linear(2, 10, bias=False)
         self.fc2 = nn.Linear(10, 50, bias=False)
         self.fc3 = nn.Linear(50, 4, bias=False)

test/distributed/_shard/sharded_optim/test_sharded_optim.py

@@ -29,7 +29,7 @@
 
 class MyShardedModel(torch.nn.Module):
     def __init__(self, spec=None, group=None):
-        super(MyShardedModel, self).__init__()
+        super().__init__()
         # Use same seed.
         torch.manual_seed(0)
         self.param = torch.nn.Parameter(torch.rand(5, 10))
@@ -47,7 +47,7 @@ def forward(self, input):
 
 class MyShardedLinear(torch.nn.Module):
     def __init__(self, rank=None):
-        super(MyShardedLinear, self).__init__()
+        super().__init__()
         # Use same seed.
         torch.manual_seed(0)
         self.linear1 = torch.nn.Linear(17, 12)

test/distributed/_tensor/test_dtensor.py

@@ -18,7 +18,7 @@
 
 class DummyMLP(torch.nn.Module):
     def __init__(self, device):
-        super(DummyMLP, self).__init__()
+        super().__init__()
         self.net1 = torch.nn.Linear(5, 1024, device=device)
         self.relu = torch.nn.ReLU()
         self.net2 = torch.nn.Linear(1024, 4, device=device)

test/distributed/algorithms/ddp_comm_hooks/test_ddp_hooks.py

@@ -43,7 +43,7 @@ def gpus_for_rank(world_size):
 
 class Task(nn.Module):
     def __init__(self):
-        super(Task, self).__init__()
+        super().__init__()
         torch.manual_seed(0)
         self.p = nn.Parameter(torch.randn(40, 20))
 
@@ -62,7 +62,7 @@ def forward(self, x, rank):
 
 class DistributedDataParallelCommHookTest(MultiProcessTestCase):
     def setUp(self):
-        super(DistributedDataParallelCommHookTest, self).setUp()
+        super().setUp()
         self._spawn_processes()
 
     def tearDown(self):

test/distributed/algorithms/quantization/test_quantization.py

@@ -43,12 +43,12 @@ def _build_tensor(size, value=None, dtype=torch.float, device_id=None):
     class DistQuantizationTests(MultiProcessTestCase):
 
         def setUp(self):
-            super(DistQuantizationTests, self).setUp()
+            super().setUp()
             self._spawn_processes()
             torch.backends.cudnn.flags(enabled=True, allow_tf32=False).__enter__()
 
         def tearDown(self):
-            super(DistQuantizationTests, self).tearDown()
+            super().tearDown()
             try:
                 os.remove(self.file_name)
             except OSError:

test/distributed/algorithms/test_join.py

@@ -83,7 +83,7 @@ class AllReducer(Joinable):
     per-iteration collective communication.
     """
     def __init__(self, device, process_group):
-        super(AllReducer, self).__init__()
+        super().__init__()
         self.device = device
         self.process_group = process_group
         self.post_hook_tensor = torch.tensor([BEFORE_CONSTANT], device=self.device)
@@ -139,7 +139,7 @@ def find_common_rank(self, rank, to_consider):
 class TestJoin(MultiProcessTestCase):
     r"""Test cases for the generic join context."""
     def setUp(self):
-        super(TestJoin, self).setUp()
+        super().setUp()
         os.environ["WORLD_SIZE"] = str(self.world_size)
         os.environ["BACKEND"] = BACKEND
         self._spawn_processes()

test/distributed/checkpoint/test_2d_fsdp_dt_checkpoint.py

@@ -39,7 +39,7 @@
 
 class SimpleModel(torch.nn.Module):
     def __init__(self):
-        super(SimpleModel, self).__init__()
+        super().__init__()
         self.net1 = torch.nn.Linear(5, 8)
         self.relu = torch.nn.ReLU()
         self.net2 = torch.nn.Linear(8, 4)

test/distributed/checkpoint/test_checkpoint.py

@@ -185,7 +185,7 @@ def _fail_rank_async(self, name, result=None):
 
 class FaultyStorageWriter(TestStorageBase, StorageWriter):
     def __init__(self, fail_conf):
-        super(FaultyStorageWriter, self).__init__(fail_conf)
+        super().__init__(fail_conf)
 
     def set_up_storage_writer(self, is_coordinator: bool) -> None:
         self._fail_rank("fail_set_up_storage_writer")
@@ -212,7 +212,7 @@ def finish(
 
 class FaultyStorageReader(TestStorageBase, StorageReader):
     def __init__(self, metadata, fail_conf):
-        super(FaultyStorageReader, self).__init__(fail_conf)
+        super().__init__(fail_conf)
         self.metadata = metadata
 
     def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None:

test/distributed/checkpoint/test_file_system_checkpoint.py

@@ -101,7 +101,7 @@ def __init__(
         self,
         spec: ShardingSpec,
     ) -> None:
-        super(MyShardedModel3, self).__init__()
+        super().__init__()
         self.sharded_tensor: ShardedTensor = sharded_tensor.rand(
             spec, 10, 20, init_rrefs=False
         )

test/distributed/checkpoint/test_file_system_checkpoint_cpu.py

@@ -100,7 +100,7 @@ def __init__(
         self,
         spec: ShardingSpec,
     ) -> None:
-        super(MyShardedModel3, self).__init__()
+        super().__init__()
         self.sharded_tensor: ShardedTensor = sharded_tensor.rand(
             spec, 10, 20, init_rrefs=False
         )

test/distributed/fsdp/test_checkpoint_wrapper.py

@@ -22,9 +22,6 @@
 
 
 class CheckpointWrapperTest(TestCase):
-    def setUp(self):
-        super().setUp()
-
     def test_load_activation_checkpointed_module(self):
         lin = nn.Linear(10, 10, bias=False)
         lin = checkpoint_wrapper(

test/distributed/fsdp/test_fsdp_mixed_precision.py

@@ -664,7 +664,7 @@ def test_grads_reduced_precision(self):
     def test_mp_batchnorm(self, convert_sync_bn):
         class BatchNormNet(nn.Module):
             def __init__(self, affine=True):
-                super(BatchNormNet, self).__init__()
+                super().__init__()
                 self.fc1 = nn.Linear(2, 40, bias=False)
                 self.bn = nn.BatchNorm1d(4, affine=affine)
                 self.fc2 = nn.Linear(40, 4, bias=False)

test/distributed/fsdp/test_fsdp_optim_state.py

@@ -286,7 +286,7 @@ def param_group1(self) -> List[torch.nn.Parameter]:
 
 class TestFSDPOptimState(FSDPTest):
     def __init__(self, *args, **kwargs):
-        super(TestFSDPOptimState, self).__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
         self._model_class = {
             _ModelClass.NESTED: self._init_nested_model,
             _ModelClass.TRANSFORMER: self._init_transformer_model,
@@ -1655,7 +1655,7 @@ def forward(self, x):
     def test_with_empty_optimizer_state(self):
         class TestDummyModel(torch.nn.Module):
             def __init__(self):
-                super(TestDummyModel, self).__init__()
+                super().__init__()
                 torch.manual_seed(0)
                 self.net1 = nn.Sequential(nn.Linear(8, 16), nn.ReLU())
                 self.net2 = nn.Sequential(nn.Linear(16, 32), nn.ReLU())

test/distributed/optim/test_named_optimizer.py

@@ -28,7 +28,7 @@ def _run_model_training(model_optim_lists):
 
 class TestDummyModel(torch.nn.Module):
     def __init__(self):
-        super(TestDummyModel, self).__init__()
+        super().__init__()
         torch.manual_seed(0)
         self.net1 = nn.Sequential(nn.Linear(8, 16), nn.ReLU())
         self.net2 = nn.Sequential(nn.Linear(16, 32), nn.ReLU())

test/distributed/optim/test_zero_redundancy_optimizer.py

@@ -64,7 +64,7 @@ def _get_backend_for_tests():
 @unittest.skipIf(TEST_WITH_ASAN or TEST_WITH_DEV_DBG_ASAN, "CUDA + ASAN does not work.")
 class TestZeroRedundancyOptimizer(common_distributed.MultiProcessTestCase):
     def setUp(self):
-        super(TestZeroRedundancyOptimizer, self).setUp()
+        super().setUp()
         os.environ["WORLD_SIZE"] = str(self.world_size)
         self._spawn_processes()
 

test/distributed/rpc/test_share_memory.py

@@ -53,9 +53,6 @@ def worker_fn(m):
     pass
 
 class TestRPCPickler(TestCase):
-    def setUp(self):
-        super().setUp()
-
     def test_case(self):
         os.environ['MASTER_ADDR'] = 'localhost'
         os.environ['MASTER_PORT'] = '29500'

test/distributed/tensor/parallel/test_2d_parallel.py

@@ -29,7 +29,7 @@
 
 class SimpleModel(torch.nn.Module):
     def __init__(self):
-        super(SimpleModel, self).__init__()
+        super().__init__()
         self.net1 = torch.nn.Linear(5, 8)
         self.relu = torch.nn.ReLU()
         self.net2 = torch.nn.Linear(8, 4)

test/distributed/tensor/parallel/test_parallelize_api.py

@@ -26,7 +26,7 @@
 
 class MLPModule(torch.nn.Module):
     def __init__(self, device):
-        super(MLPModule, self).__init__()
+        super().__init__()
         torch.manual_seed(5)
         self.net1 = torch.nn.Linear(10, 16, device=device)
         self.relu = torch.nn.ReLU()

test/distributed/tensor/parallel/test_tp_examples.py

@@ -20,7 +20,7 @@
 
 class MLPModule(torch.nn.Module):
     def __init__(self, device):
-        super(MLPModule, self).__init__()
+        super().__init__()
         torch.manual_seed(5)
         self.net1 = torch.nn.Linear(10, 16, device=device)
         self.relu = torch.nn.ReLU()