Let exp support complex types on CUDA and enable device/dtype in complex tests (pytorch#39087)

xuhdev · facebook-github-bot · commit a303fd2ea679 · 2020-06-30T10:50:40.000-07:00
Summary: Pull Request resolved: pytorch#39087 Differential Revision: D22169697 Pulled By: anjali411 fbshipit-source-id: 4866b7be6742508cc40540ed1ac811f005531d8b
diff --git a/aten/src/ATen/native/cuda/UnaryOpsKernel.cu b/aten/src/ATen/native/cuda/UnaryOpsKernel.cu
@@ -7,6 +7,7 @@
 #include <ATen/native/DispatchStub.h>
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/cuda/Math.cuh>
+#include <c10/util/complex.h>
 
 namespace at { namespace native {
 
@@ -25,7 +26,7 @@ void bitwise_not_kernel_cuda(TensorIterator& iter) {
 }
 
 void exp_kernel_cuda(TensorIterator& iter) {
-  AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "exp_cuda", [&]() {
+  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "exp_cuda", [&]() {
     AT_SKIP_BFLOAT16_IF_NOT_ROCM(scalar_t, "exp_cuda", [&] {
       gpu_kernel(iter, []GPU_LAMBDA(scalar_t a) -> scalar_t {
         return ::exp(a);
diff --git a/test/test_complex.py b/test/test_complex.py
@@ -1,36 +1,27 @@
-import math
 import torch
-from torch.testing._internal.common_utils import TestCase, run_tests, TEST_NUMPY
-import unittest
-
-if TEST_NUMPY:
-    import numpy as np
+from torch.testing._internal.common_device_type import instantiate_device_type_tests, dtypes
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 devices = (torch.device('cpu'), torch.device('cuda:0'))
 
 
 class TestComplexTensor(TestCase):
-    def test_to_list_with_complex_64(self):
+    @dtypes(*torch.testing.get_all_complex_dtypes())
+    def test_to_list(self, device, dtype):
         # test that the complex float tensor has expected values and
         # there's no garbage value in the resultant list
-        self.assertEqual(torch.zeros((2, 2), dtype=torch.complex64).tolist(), [[0j, 0j], [0j, 0j]])
-
-    @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
-    def test_exp(self):
-        def exp_fn(dtype):
-            a = torch.tensor(1j, dtype=dtype) * torch.arange(18) / 3 * math.pi
-            expected = np.exp(a.numpy())
-            actual = torch.exp(a)
-            self.assertEqual(actual, torch.from_numpy(expected))
+        self.assertEqual(torch.zeros((2, 2), device=device, dtype=dtype).tolist(), [[0j, 0j], [0j, 0j]])
 
-        exp_fn(torch.complex64)
-        exp_fn(torch.complex128)
-
-    def test_dtype_inference(self):
+    @dtypes(torch.float32, torch.float64)
+    def test_dtype_inference(self, device, dtype):
         # issue: https://github.com/pytorch/pytorch/issues/36834
-        torch.set_default_dtype(torch.double)
-        x = torch.tensor([3., 3. + 5.j])
-        self.assertEqual(x.dtype, torch.cdouble)
+        default_dtype = torch.get_default_dtype()
+        torch.set_default_dtype(dtype)
+        x = torch.tensor([3., 3. + 5.j], device=device)
+        torch.set_default_dtype(default_dtype)
+        self.assertEqual(x.dtype, torch.cdouble if dtype == torch.float64 else torch.cfloat)
+
+instantiate_device_type_tests(TestComplexTensor, globals())
 
 if __name__ == '__main__':
     run_tests()
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -11063,6 +11063,95 @@ def test_exponential(self, device, dtype):
         with self.assertRaises(RuntimeError):
             torch.empty((1,), device=device, dtype=dtype).exponential_(-0.5)
 
+    @unittest.skipIf(not TEST_NUMPY, "Numpy not found")
+    @dtypes(*(torch.testing.get_all_fp_dtypes(include_half=False) +
+              torch.testing.get_all_complex_dtypes()))
+    @dtypesIfCUDA(*(torch.testing.get_all_fp_dtypes(include_half=True) +
+                    torch.testing.get_all_complex_dtypes()))
+    def test_exp(self, device, dtype):
+        for v in (2, -2) + ((1j, 1 + 1j) if dtype.is_complex else ()):
+            if dtype == torch.bfloat16:
+                # Currently multiply a bfloat16 type with floating-point causes error:
+                #   RuntimeError: dtype != ScalarType::Undefined INTERNAL ASSERT FAILED at
+                #   "/pytorch/aten/src/ATen/native/TensorIterator.cpp":125, please report a bug to PyTorch.
+                # We skip bfloat16 for now, but we should fix it. https://github.com/pytorch/pytorch/issues/40580
+                if self.device_type == 'cpu' or self.device_type == 'cuda':
+                    with self.assertRaises(RuntimeError):
+                        torch.tensor(v, dtype=dtype, device=device) * torch.arange(18, device=device)
+                    return
+                elif self.device_type == 'xla':
+                    # Error:
+                    # Traceback (most recent call last):
+                    # File "/opt/conda/lib/python3.6/site-packages/torch/testing/_internal/common_device_type.py",
+                    # line 241, in instantiated_test
+                    #   result = test(self, device_arg, dtype)
+                    # File "/var/lib/jenkins/workspace/xla/test/../../test/test_torch.py", line 11062, in test_exp
+                    #   self.compare_with_numpy(torch.exp, np.exp, a)
+                    # File "/opt/conda/lib/python3.6/site-packages/torch/testing/_internal/common_utils.py", line 878,
+                    # in compare_with_numpy
+                    #   a = tensor_like.detach().cpu().numpy()
+                    # TypeError: Got unsupported ScalarType BFloat16
+                    return
+
+            a = torch.tensor(v, dtype=dtype, device=device) * torch.arange(18, device=device) / 3 * math.pi
+            a = a.to(dtype)
+            self.compare_with_numpy(torch.exp, np.exp, a)
+
+            if dtype.is_complex:
+                inf_real_zero_imag_in = torch.tensor(complex(float('inf'), 0), device=device, dtype=dtype)
+                inf_real_zero_imag_out = torch.exp(inf_real_zero_imag_in).item()
+                self.assertTrue(math.isinf(inf_real_zero_imag_out.real))
+                if self.device_type == 'cpu':
+                    if not IS_WINDOWS:  # Windows tests don't show some bugs consistently
+                        # This is incorrect. It should be zero. Need fix!
+                        # https://github.com/pytorch/pytorch/issues/40590
+                        self.assertNotEqual(inf_real_zero_imag_out.imag, 0)
+                        # This is incorrect. They should equal. Need fix!
+                        # https://github.com/pytorch/pytorch/issues/40590
+                        with self.assertRaises(AssertionError):
+                            self.compare_with_numpy(torch.exp, np.exp, inf_real_zero_imag_in)
+                else:
+                    self.assertEqual(inf_real_zero_imag_out.imag, 0, atol=0, rtol=0)
+                    self.compare_with_numpy(torch.exp, np.exp, inf_real_zero_imag_in)
+
+                zero_real_inf_imag_in = torch.tensor(complex(0, float('inf')), device=device, dtype=dtype)
+                zero_real_inf_imag_out = torch.exp(zero_real_inf_imag_in).item()
+                self.assertTrue(math.isnan(zero_real_inf_imag_out.real))
+                self.assertTrue(math.isnan(zero_real_inf_imag_out.imag))
+                # Ensure we are notified when NumPy changes its behavior
+                self.compare_with_numpy(torch.exp, np.exp, zero_real_inf_imag_in)
+
+                inf_real_imag_in = torch.tensor(complex(float('inf'), float('inf')), device=device, dtype=dtype)
+                inf_real_imag_out = torch.exp(inf_real_imag_in).item()
+                if self.device_type == 'cpu':
+                    if not IS_WINDOWS:  # Windows tests don't show some bugs consistently
+                        # This is incorrect. Need fix! https://github.com/pytorch/pytorch/issues/40590
+                        with self.assertRaises(AssertionError):
+                            self.compare_with_numpy(torch.exp, np.exp, inf_real_imag_in)
+                else:
+                    self.assertTrue(math.isinf(inf_real_imag_out.real))
+                    self.assertTrue(math.isnan(inf_real_imag_out.imag))
+                    self.compare_with_numpy(torch.exp, np.exp, inf_real_imag_in)
+
+                inf_real_nan_imag_in = torch.tensor(complex(float('inf'), float('nan')), device=device, dtype=dtype)
+                inf_real_nan_imag_out = torch.exp(inf_real_nan_imag_in).item()
+                if self.device_type == 'cpu':
+                    if not IS_WINDOWS:  # Windows tests don't show some bugs consistently
+                        # This is incorrect. It should be inf. Need fix! https://github.com/pytorch/pytorch/issues/40590
+                        with self.assertRaises(AssertionError):
+                            self.compare_with_numpy(torch.exp, np.exp, inf_real_nan_imag_in)
+                else:
+                    self.assertTrue(math.isinf(inf_real_nan_imag_out.real))
+                    self.assertTrue(math.isnan(inf_real_nan_imag_out.imag))
+                    self.compare_with_numpy(torch.exp, np.exp, inf_real_nan_imag_in)
+
+                nan_real_inf_imag_in = torch.tensor(complex(float('nan'), float('inf')), device=device, dtype=dtype)
+                nan_real_inf_imag_out = torch.exp(nan_real_inf_imag_in).item()
+                self.assertTrue(math.isnan(nan_real_inf_imag_out.real))
+                self.assertTrue(math.isnan(nan_real_inf_imag_out.imag))
+                # Ensure we are notified when NumPy changes its behavior
+                self.compare_with_numpy(torch.exp, np.exp, nan_real_inf_imag_in)
+
     @skipIfNoSciPy
     @dtypes(*torch.testing.get_all_fp_dtypes())
     def test_uniform_kstest(self, device, dtype):
