Cleaned up test_mps.py::test_output*_match (#108092)

vfdev-5 · pytorchmergebot · commit b7624fc91e05 · 2023-08-29T10:46:02.000Z
Description: - cleaned up test_mps.py::test_output_match and test_mps.py::test_output_grad_match tests - removed unused variables and useless brackets - simplified atol/rtol setup if/else code Pull Request resolved: #108092 Approved by: https://github.com/kulinseth
diff --git a/test/test_mps.py b/test/test_mps.py
@@ -10624,14 +10624,11 @@ class TestConsistency(TestCaseMPS):
     @ops(mps_ops_modifier(test_consistency_op_db), allowed_dtypes=MPS_DTYPES)
     def test_output_match(self, device, dtype, op):
         self.assertEqual(device, "cpu")
-        key = op.name + op.variant_test_name
-        run_grad_test = True
 
         def get_samples():
             return op.sample_inputs(device, dtype, requires_grad=(dtype.is_floating_point or dtype.is_complex))
         cpu_samples = get_samples()
 
-        all_backward_pass = True
         for cpu_sample in cpu_samples:
             #
             # Forward check
@@ -10645,7 +10642,7 @@ def get_samples():
             mps_kwargs = mps_sample.kwargs
 
             # for tensor_split(), the second tensor arg ("tensor_indices_or_sections") must be on CPU only
-            if (op.name == "tensor_split" and isinstance(mps_args[1], torch.Tensor)):
+            if op.name == "tensor_split" and isinstance(mps_args[1], torch.Tensor):
                 mps_args[1] = cpu_args[1]
 
             cpu_out = op(*cpu_args, **cpu_kwargs)
@@ -10676,16 +10673,11 @@ def get_samples():
     @ops(mps_ops_grad_modifier(copy.deepcopy(test_consistency_op_db)), allowed_dtypes=MPS_GRAD_DTYPES)
     def test_output_grad_match(self, device, dtype, op):
         self.assertEqual(device, "cpu")
-        key = op.name + op.variant_test_name
-
-        run_grad_test = True
 
         def get_samples():
             return op.sample_inputs(device, dtype, requires_grad=(dtype.is_floating_point or dtype.is_complex))
         cpu_samples = get_samples()
 
-        all_forward_pass = True
-        all_backward_pass = True
         for cpu_sample in cpu_samples:
             #
             # Forward check
@@ -10700,19 +10692,16 @@ def get_samples():
             mps_kwargs = mps_sample.kwargs
 
             # for tensor_split(), the second tensor arg ("tensor_indices_or_sections") must be on CPU only
-            if (op.name == "tensor_split" and isinstance(mps_args[1], torch.Tensor)):
+            if op.name == "tensor_split" and isinstance(mps_args[1], torch.Tensor):
                 mps_args[1] = cpu_args[1]
 
             cpu_out = op(*cpu_args, **cpu_kwargs)
             mps_out = op(*mps_args, **mps_kwargs)
 
-            if (op.name in self.FP32_LOW_PRECISION_LIST) and dtype == torch.float32:
-                atol = 1e-4
-                rtol = 3e-5
-            elif op.name == "nn.functional.conv2d" or op.name == "linalg.multi_dot" and dtype == torch.float32:
+            if op.name in self.FP32_LOW_PRECISION_LIST and dtype == torch.float32:
                 atol = 1e-4
                 rtol = 3e-5
-            elif (op.name in self.FP16_LOW_PRECISION_LIST) and dtype == torch.float16:
+            elif op.name in self.FP16_LOW_PRECISION_LIST and dtype == torch.float16:
                 atol = 1e-2
                 rtol = 1e-2
             elif (op.name == "masked.mean"):
@@ -10732,7 +10721,6 @@ def get_samples():
 
             self.assertEqual(cpu_out, mps_out, atol=atol, rtol=rtol)
 
-
             #
             # Backward check
             #
