Setter for real and imag tensor attributes

anjali411 · anjali411 · commit ff7b4415678d · 2020-06-19T14:35:52.000-07:00
ghstack-source-id: 2afdbd9 Pull Request resolved: #39860
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_torch.py b/test/test_torch.py
@@ -18142,6 +18142,30 @@ 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(*torch.testing.get_all_complex_dtypes())
+    def test_set_real_imag(self, device, dtype):
+        x = torch.randn(10, dtype=dtype)
+
+        for d in torch.testing.get_all_dtypes():
+            new_real = _make_tensor((10,), d, device)
+            new_imag = _make_tensor((10,), d, device)
+
+            if d.is_complex:
+                regex = "Casting complex values to real discards the imaginary part"
+                with self.assertWarnsRegex(UserWarning, regex):
+                    x.real = new_real
+                with self.assertWarnsRegex(UserWarning, regex):
+                    x.imag = new_imag
+                self.assertEqualIgnoreType(x.real, new_real.real)
+                self.assertEqualIgnoreType(x.imag, new_imag.real)
+
+            else:
+                x.real = new_real
+                x.imag = new_imag
+                self.assertEqualIgnoreType(x.real, new_real)
+                self.assertEqualIgnoreType(x.imag, new_imag)
+
     def test_diagonal_view(self, device) -> None:
         t = torch.ones((5, 5), device=device)
         v = torch.diagonal(t)
@@ -18464,7 +18488,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 not 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
@@ -534,6 +534,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[] = {
@@ -563,8 +583,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`	`}`