Setter for real and imag tensor attributes (#39860)

anjali411 · facebook-github-bot · commit 11a74a58c86c · 2020-06-29T15:44:55.000-07:00
Summary: Pull Request resolved: #39860 Test Plan: Imported from OSS Reviewed By: mruberry Differential Revision: D22163234 Pulled By: anjali411 fbshipit-source-id: 35b4aa16499341edff1a4be4076539ac7c74f5be
diff --git a/aten/src/ATen/native/Copy.cpp b/aten/src/ATen/native/Copy.cpp
@@ -161,6 +161,10 @@ static Tensor & copy_impl(Tensor & self, const Tensor & src, bool non_blocking)
     return self;
   }
 
+  if(!self.is_complex() && src.is_complex()) {
+    TORCH_WARN_ONCE("Casting complex values to real discards the imaginary part");
+  }
+
   copy_stub(device_type, iter, non_blocking);
   return self;
 }
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -2666,8 +2666,9 @@ def test_profiler(self):
         self.assertFalse(torch.autograd._profiler_enabled())
 
         last_end = 0
-        names = ['is_complex', 'mul', 'to', 'empty_strided', 'copy_', 'empty', 'is_complex',
-                 'add', 'to', 'empty_strided', 'copy_', 'empty']
+        names = ['is_complex', 'mul', 'to', 'empty_strided', 'copy_', 'is_complex',
+                 'is_complex', 'empty', 'is_complex', 'add', 'to', 'empty_strided',
+                 'copy_', 'is_complex', 'is_complex', 'empty']
         top_level_names = ['is_complex', 'mul', 'is_complex', 'add']
         top_level_iter = iter(top_level_names)
         self.assertEqual(len(p.function_events), len(names))
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -18291,6 +18291,26 @@ def compare_with_numpy(contiguous_input=True):
         self.assertEqual(a[5:].real, a.real[5:])
         self.assertEqual(a[5:].imag, a.imag[5:])
 
+    @onlyOnCPUAndCUDA
+    @dtypes(*product(torch.testing.get_all_complex_dtypes(), torch.testing.get_all_dtypes()))
+    @suppress_warnings
+    def test_set_real_imag(self, device, dtypes):
+        x = torch.randn(10, dtype=dtypes[0], device=device)
+
+        new_real = _make_tensor((10,), dtypes[1], device)
+        new_imag = _make_tensor((10,), dtypes[1], device)
+
+        x.real = new_real
+        x.imag = new_imag
+
+        if dtypes[1].is_complex:
+            self.assertEqual(x.real, new_real.real, exact_dtype=False)
+            self.assertEqual(x.imag, new_imag.real, exact_dtype=False)
+
+        else:
+            self.assertEqual(x.real, new_real, exact_dtype=False)
+            self.assertEqual(x.imag, new_imag, exact_dtype=False)
+
     def test_diagonal_view(self, device) -> None:
         t = torch.ones((5, 5), device=device)
         v = torch.diagonal(t)
@@ -18615,7 +18635,7 @@ def _make_tensor(shape, dtype, device, fill_ones=False) -> torch.Tensor:
         return torch.ones(*shape, dtype=_convert_t(dtype, device), device=device)
 
     # Returns a tensor with random integer values
-    if not dtype.is_floating_point:
+    if not (dtype.is_floating_point or dtype.is_complex):
         t = torch.randint(0, 10, shape, device=device)
         if dtype != torch.uint8:
             t = t - 5  # generate negative values also
diff --git a/torch/csrc/autograd/python_variable.cpp b/torch/csrc/autograd/python_variable.cpp
@@ -542,6 +542,26 @@ PyObject *THPVariable_get_imag(THPVariable* self, void *unused)
   END_HANDLE_TH_ERRORS
 }
 
+int THPVariable_set_real(THPVariable *self, THPVariable *real, void *unused)
+{
+  HANDLE_TH_ERRORS
+  auto& self_ = self->cdata;
+  auto self_real = at::real(self_);
+  self_real.copy_(real->cdata);
+  return 0;
+  END_HANDLE_TH_ERRORS_RET(-1)
+}
+
+int THPVariable_set_imag(THPVariable* self, THPVariable *imag, void *unused)
+{
+  HANDLE_TH_ERRORS
+  auto& self_ = self->cdata;
+  auto self_imag = at::imag(self_);
+  self_imag.copy_(imag->cdata);
+  return 0;
+  END_HANDLE_TH_ERRORS_RET(-1)
+}
+
 // properties are registered here because we are currently only able to bind them
 // manually. TODO: make declarable in native_functions
 static struct PyGetSetDef THPVariable_properties[] = {
@@ -572,8 +592,8 @@ static struct PyGetSetDef THPVariable_properties[] = {
   {"device", (getter)THPVariable_device, nullptr, nullptr, nullptr},
   {"ndim", (getter)THPVariable_get_ndim, nullptr, nullptr, nullptr},
   {"names", (getter)THPVariable_get_names, (setter)THPVariable_set_names, nullptr, nullptr},
-  {"real", (getter)THPVariable_get_real, nullptr, nullptr, nullptr},
-  {"imag", (getter)THPVariable_get_imag, nullptr, nullptr, nullptr},
+  {"real", (getter)THPVariable_get_real, (setter)THPVariable_set_real, nullptr, nullptr},
+  {"imag", (getter)THPVariable_get_imag, (setter)THPVariable_set_imag, nullptr, nullptr},
   {nullptr}
 };
 

Original file line number	Diff line number	Diff line change
`@@ -161,6 +161,10 @@ static Tensor & copy_impl(Tensor & self, const Tensor & src, bool non_blocking)`
`161`	`161`	`return self;`
`162`	`162`	`}`
`163`	`163`
	`164`	`+ if(!self.is_complex() && src.is_complex()) {`
	`165`	`+ TORCH_WARN_ONCE("Casting complex values to real discards the imaginary part");`
	`166`	`+ }`
	`167`	`+`
`164`	`168`	`copy_stub(device_type, iter, non_blocking);`
`165`	`169`	`return self;`
`166`	`170`	`}`