Update on "torch.sgn for complex tensors"

anjali411 · anjali411 · commit 26b367ee5202 · 2020-09-03T11:55:16.000-07:00
resolves #36323 by adding `torch.sgn` for complex tensors. `torch.sgn` returns `x/abs(x)` for `x != 0` and returns `0 + 0j` for `x==0` This PR doesn't test the correctness of the gradients. It will be done as a part of auditing all the ops in future once we decide the autograd behavior (JAX vs TF) and add gradchek. Differential Revision: [D23460526](https://our.internmc.facebook.com/intern/diff/D23460526) [ghstack-poisoned]
diff --git a/aten/src/ATen/cpu/vec256/vec256_base.h b/aten/src/ATen/cpu/vec256/vec256_base.h
@@ -241,7 +241,7 @@ struct Vec256 {
   }
 
   template <typename other_t_sgn = T,
-            typename std::enable_if<c10::is_complex_t<other_t_sgn>::value, int>::type = 0>
+            typename std::enable_if<c10::is_complex<other_t_sgn>::value, int>::type = 0>
   Vec256<T> sgn() const {
     return map(at::native::sgn_impl);
   }
diff --git a/torch/csrc/autograd/FunctionsManual.cpp b/torch/csrc/autograd/FunctionsManual.cpp
@@ -774,9 +774,13 @@ Tensor trace_backward(const Tensor & grad, IntArrayRef sizes) {
 }
 
 Tensor sgn_backward(Tensor result, Tensor grad, Tensor self) {
-  // [grad / abs(z) - Re(grad/self) * result
-  auto abs = at::abs(self);
-  return at::where(abs == 0.0, at::zeros({}, grad.options()), (grad/abs - (at::real(grad/self) * result)));
+  if (self.is_complex()) {
+    // vjp = grad / abs(self) - Re(grad/self) * result
+    auto abs = at::abs(self);
+    return at::where(abs == 0.0, at::zeros({}, grad.options()), (grad/abs - (at::real(grad/self) * result)));
+  } else {
+    return at::zeros_like(grad, at::MemoryFormat::Preserve);
+  }
 }
 
 Tensor var_backward(const Tensor & grad, const Tensor & self, bool unbiased) {

Original file line number	Diff line number	Diff line change
`@@ -241,7 +241,7 @@ struct Vec256 {`
`241`	`241`	`}`
`242`	`242`
`243`	`243`	`template <typename other_t_sgn = T,`
`244`		`- typename std::enable_if<c10::is_complex_t<other_t_sgn>::value, int>::type = 0>`
	`244`	`+ typename std::enable_if<c10::is_complex<other_t_sgn>::value, int>::type = 0>`
`245`	`245`	`Vec256<T> sgn() const {`
`246`	`246`	`return map(at::native::sgn_impl);`
`247`	`247`	`}`