Added autograd tests, disabled jit autograd tests for complex and added a separate list for tests for complex dtype only (#38399)

anjali411 · facebook-github-bot · commit 1676c7d61845 · 2020-05-13T16:41:09.000-07:00
Summary: Pull Request resolved: #38399 Test Plan: Imported from OSS Differential Revision: D21555941 Pulled By: anjali411 fbshipit-source-id: ea9f5a76590c5bab3df6a540617b074238bfb535
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -4185,11 +4185,18 @@ def run_functional_checks(test_case, test_name, name, apply_fn, run_grad_checks,
         test_case.assertEqualTypeString(self_variable, self_variable.grad)
         test_case.assertEqual(self_variable.size(), self_variable.grad.size())
 
+# this list corresponds to ops which have separate tests defined for complex dtypes in
+# common_methods_invocations.py
+# test for these ops with 'complex' in variant should only run for complex and
+# the tests for these ops which do not have 'complex' in variant should not run for complex
+# and only run for floating point
+separate_complex_tests = ['log', 'log10', 'log1p', 'log2', 'reciprocal']
+
 # white list for complex
 complex_list = ['t', 'view', 'reshape', 'reshape_as', 'view_as', 'zero_', 'clone',
                 'tril', 'triu', 'fill_', 'eq_', 'ne_', 'permute', 'squeeze', 'unsqueeze',
                 'chunk', 'split', 'split_with_sizes', 'resize', 'resize_as', 'sin', 'cos',
-                '__rmul__', '__rdiv__']
+                '__rmul__', '__rdiv__', 'transpose', 'round'] + separate_complex_tests
 
 def add_test(
         name,
@@ -4206,17 +4213,29 @@ def add_test(
     if variant_name != '':
         basic_test_name += '_' + variant_name
 
+    if name in separate_complex_tests and 'complex' in variant_name:
+        run_only_complex = True
+    else:
+        run_only_complex = False
+
     for dtype in [torch.double, torch.cdouble]:
         for dim_perm in product([-1, 1], repeat=len(dim_args_idx)):
             test_name = basic_test_name
             new_args = [arg * dim_perm[dim_args_idx.index(i)] if i in dim_args_idx else arg for i, arg in enumerate(args)]
             test_name = basic_test_name + ''.join('_neg' + str(i) for i, idx in enumerate(dim_perm) if idx < 0)
+
             if dtype.is_complex:
                 # TODO: remove this. this is temporary while we ramp up the complex support.
                 if name in complex_list and 'scalar' not in test_name and 'constant' not in test_name:
-                    test_name = test_name + '_complex'
+                    if name in separate_complex_tests and 'complex' not in variant_name:
+                        continue
+                    if not run_only_complex:
+                        test_name = test_name + '_complex'
                 else:
                     continue
+            elif run_only_complex:
+                continue
+
             new_args = tuple(new_args)
 
             # for-loop bodies don't define scopes, so we have to save the variables
@@ -4239,8 +4258,6 @@ def check(name):
                         output_variable = getattr(self_variable, name)(*args_variable, **kwargs_variable)
                         output_tensor = getattr(self_tensor, name)(*args_tensor, **kwargs_variable)
                         if not isinstance(output_tensor, torch.Tensor) and not istuple(output_tensor):
-                            # TODO: I'm not sure why we insert an outer dimension
-                            # here, seems a bit strange
                             if dtype.is_complex:
                                 output_tensor = torch.tensor((output_tensor, ), dtype=torch.cfloat, device=device)
                             else:
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -17948,6 +17948,12 @@ def add_autograd_test(
         skipTestIf=(),
         output_process_fn=lambda x: x,
         kwargs=None):
+
+    # Disable complex tests
+    # TODO: Add complex support for jit
+    if 'complex' in variant_name:
+        return
+
     basic_test_name = 'test_' + name
     if variant_name != '':
         basic_test_name += '_' + variant_name
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
@@ -176,8 +176,8 @@ def method_tests():
         ('transpose', (1, 2, 3), (1, 2), 'dim', (True,), [0, 1]),
         ('transpose', (), (0, 0), 'scalar', (True,)),
         ('transpose', (1,), (0, 0), '1d', (True,)),
-        ('transpose', torch.rand(L, L), (0, 1), '2d', (True,)),
-        ('transpose', torch.rand(S, S, S), (2, 0), '3d', (True,)),
+        ('transpose', (L, L), (0, 1), '2d', (True,)),
+        ('transpose', (S, S, S), (2, 0), '3d', (True,)),
         ('t', (1, 2), NO_ARGS, '', (True,)),
         ('view', (S, S, S), (S * S, S), '', (True,)),
         ('view', (S, S, S), (torch.Size([S * S, S]),), 'size', (True,)),
@@ -238,6 +238,12 @@ def method_tests():
         ('log1p', uniform_scalar(requires_grad=True), NO_ARGS, 'scalar', (True,)),
         ('log2', torch.rand(S, S, S) + 1e-2, NO_ARGS, '', (True,)),
         ('log2', uniform_scalar(1e-2, requires_grad=True), NO_ARGS, 'scalar', (True,)),
+        ('log', torch.randn(S, S, S, dtype=torch.cfloat) + 1e-2, NO_ARGS, 'complex', (True,)),
+        ('log', uniform_scalar(1e-2j, requires_grad=True), NO_ARGS, 'complex_scalar', (True,)),
+        ('log10', torch.randn(S, S, S, dtype=torch.cfloat) + 1e-2, NO_ARGS, 'complex', (True,)),
+        ('log10', uniform_scalar(1e-2j, requires_grad=True), NO_ARGS, 'complex_scalar', (True,)),
+        ('log2', torch.randn(S, S, S, dtype=torch.cfloat) + 1e-2, NO_ARGS, 'complex', (True,)),
+        ('log2', uniform_scalar(1e-2j, requires_grad=True), NO_ARGS, 'complex_scalar', (True,)),
         ('tanh', (S, S, S), NO_ARGS, '', (True,)),
         ('tanh', (), NO_ARGS, 'scalar', (True,)),
         ('sigmoid', (S, S, S), NO_ARGS, '', (True,)),
@@ -273,6 +279,8 @@ def method_tests():
         ('atan2', (S, 1, S), ((S, S),), 'broadcast_all'),
         ('reciprocal', torch.rand(S, S, S) + 0.1, NO_ARGS, '', (True,)),
         ('reciprocal', uniform_scalar(0.1, requires_grad=True), NO_ARGS, 'scalar', (True,)),
+        ('reciprocal', torch.randn(S, S, S, dtype=torch.cdouble) + 0.1, NO_ARGS, 'complex', (True,)),
+        ('reciprocal', uniform_scalar(0.1j), NO_ARGS, 'complex_scalar', (True,)),
         ('round', (S, S, S), NO_ARGS, '', (True,)),
         ('round', (), NO_ARGS, 'scalar', (True,)),
         ('sign', (S, S, S), NO_ARGS),
