Improve complex support in common_nn test machinery (#50593)

peterbell10 · facebook-github-bot · commit 47f0bda3ef8d · 2021-01-22T09:44:45.000-08:00
Summary: Pull Request resolved: #50593 There are no equivalent to torch.FloatTensor, torch.cuda.FloatTensor for complex types. So `get_gpu_type` and `get_cpu_type` are broken for complex tensors. Also found a few places that explicitly cast inputs to floating point types, which would drop the imaginary component before running the test. Test Plan: Imported from OSS Reviewed By: ngimel Differential Revision: D25954050 Pulled By: mruberry fbshipit-source-id: 1fa8e5af233aa095c839d5e2f860564baaf92aef
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -15885,8 +15885,14 @@ def create_nn_module(*args, **kwargs):
             return module(*args)
 
         # Set up inputs from tuple of sizes or constructor fn
+        dtype = torch.double
         if 'input_fn' in kwargs:
             input = kwargs['input_fn']()
+            if isinstance(input, Tensor):
+                input = (input,)
+
+            if all(tensor.is_complex() for tensor in input):
+                dtype = torch.cdouble
         else:
             input = (kwargs['input_size'],)
 
@@ -15903,7 +15909,7 @@ def create_nn_module(*args, **kwargs):
         if 'extra_args' in kwargs:
             input = input + kwargs['extra_args']
 
-        args_variable, kwargs_variable = create_input(input)
+        args_variable, kwargs_variable = create_input(input, dtype=dtype)
         f_args_variable = deepcopy(unpack_variables(args_variable))
 
         # Check against Python module as reference
diff --git a/torch/testing/_internal/common_nn.py b/torch/testing/_internal/common_nn.py
@@ -4973,7 +4973,7 @@ def test_cuda(self, test_case):
         # Run backwards on CPU and GPU and compare results
         for _ in range(5):
             cpu_gradOutput = cpu_output.clone().normal_()
-            gpu_gradOutput = cpu_gradOutput.type('torch.cuda.FloatTensor')
+            gpu_gradOutput = cpu_gradOutput.type_as(gpu_output)
             cpu_gradInput = test_case._backward(cpu_module, cpu_input_tuple, cpu_output, cpu_gradOutput)
             gpu_gradInput = test_case._backward(gpu_module, gpu_input_tuple, gpu_output, gpu_gradOutput)
             # TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
@@ -5047,6 +5047,7 @@ def __init__(self, *args, **kwargs):
         self.check_inplace = kwargs.get('check_inplace', False)
         self.check_gradgrad = kwargs.get('check_gradgrad', True)
         self.skip_double = kwargs.get('skip_double', False)
+        self.skip_half = kwargs.get('skip_half', False)
         self.with_tf32 = kwargs.get('with_tf32', False)
         self.tf32_precision = kwargs.get('tf32_precision', 0.001)
         self.test_cpu = kwargs.get('test_cpu', True)
@@ -5136,15 +5137,32 @@ def assert_module_parameters_are(tensor_type, device_id=None):
                 assert_module_parameters_are(torch.cuda.FloatTensor, 1)  # type: ignore[attr-defined]
         else:
             # check that float()/double() casters work correctly
+            def to_type(tensor, real, complex):
+                if tensor.is_complex():
+                    return tensor.to(complex)
+                elif tensor.is_floating_point():
+                    return tensor.to(real)
+                else:
+                    return tensor
+
+            def to_half(x):
+                # TODO: torch.complex32 when properly supported
+                return to_type(x, torch.float16, None)
+
+            def to_single(x):
+                return to_type(x, torch.float32, torch.complex64)
+
+            def to_double(x):
+                return to_type(x, torch.float64, torch.complex128)
 
             # to float
-            input_tuple = tuple(t.float() if not isinstance(t, torch.LongTensor) else t for t in input_tuple)
+            input_tuple = tuple(to_single(t) for t in input_tuple)
             module.float()
             module(*input_tuple)
             assert_module_parameters_are(torch.FloatTensor)
 
             # and back to double
-            input_tuple = tuple(t.double() if not isinstance(t, torch.LongTensor) else t for t in input_tuple)
+            input_tuple = tuple(to_double(t) for t in input_tuple)
             module.double()
             module(*input_tuple)
             assert_module_parameters_are(torch.DoubleTensor)
@@ -5154,8 +5172,7 @@ def assert_module_parameters_are(tensor_type, device_id=None):
                 # and that float() casts parameters correctly
 
                 # to GPU0
-                input_tuple = tuple(
-                    t.float().cuda() if not isinstance(t, torch.LongTensor) else t.cuda() for t in input_tuple)
+                input_tuple = tuple(to_single(t).cuda() for t in input_tuple)
                 module.float().cuda()
                 module(*input_tuple)
                 assert_module_parameters_are(torch.cuda.FloatTensor, 0)  # type: ignore[attr-defined]
@@ -5189,18 +5206,17 @@ def assert_module_parameters_are(tensor_type, device_id=None):
 
                 if not self.skip_double:
                     # test double()
-                    input_tuple = tuple(
-                        t.double().cuda() if not isinstance(t, torch.LongTensor) else t.cuda() for t in input_tuple)
+                    input_tuple = tuple(to_double(t).cuda() for t in input_tuple)
                     module.double().cuda()
                     module(*input_tuple)
                     assert_module_parameters_are(torch.cuda.DoubleTensor, 0)  # type: ignore[attr-defined]
 
                 # test half()
-                input_tuple = tuple(
-                    t.half().cuda() if not isinstance(t, torch.LongTensor) else t.cuda() for t in input_tuple)
-                module.half().cuda()
-                module(*input_tuple)
-                assert_module_parameters_are(torch.cuda.HalfTensor, 0)  # type: ignore[attr-defined]
+                if not self.skip_half:
+                    input_tuple = tuple(to_half(t).cuda() for t in input_tuple)
+                    module.half().cuda()
+                    module(*input_tuple)
+                    assert_module_parameters_are(torch.cuda.HalfTensor, 0)  # type: ignore[attr-defined]
         torch.set_num_threads(num_threads)
 
     def _get_target(self):
diff --git a/torch/testing/_internal/jit_metaprogramming_utils.py b/torch/testing/_internal/jit_metaprogramming_utils.py
@@ -527,8 +527,14 @@ def try_get_nn_module_compiled_mod_and_inputs(*args, **kwargs):
         constructor_args = kwargs.get('constructor_args', ())
 
     # Set up inputs from tuple of sizes or constructor fn
+    input_dtype = torch.double
     if 'input_fn' in kwargs:
         input = kwargs['input_fn']()
+        if isinstance(input, torch.Tensor):
+            input = (input,)
+
+        if all(tensor.is_complex() for tensor in input):
+            input_dtype = torch.cdouble
     else:
         input = (kwargs['input_size'],)
 
@@ -543,7 +549,7 @@ def try_get_nn_module_compiled_mod_and_inputs(*args, **kwargs):
             input = (input,)
         input = input + (kwargs['target_fn'](),)
 
-    args_variable, kwargs_variable = create_input(input)
+    args_variable, kwargs_variable = create_input(input, dtype=input_dtype)
     f_args_variable = deepcopy(unpack_variables(args_variable))
     out_var = deepcopy(f_args_variable)
 
