add std to autodiff, and mean/var/std to operator set (#17137)

Summary:
supersedes #16684
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17137

Differential Revision: D14096724

Pulled By: wanchaol

fbshipit-source-id: d801d70029a6a1f5851400ff4094c0299c102b2b

diff --git a/torch/csrc/jit/symbolic_script.cpp b/torch/csrc/jit/symbolic_script.cpp
index e401c0e..7a7ab6d 100644 (file)
--- a/torch/csrc/jit/symbolic_script.cpp
+++ b/torch/csrc/jit/symbolic_script.cpp
@@ -294,6 +294,26 @@ const std::vector<std::string> functions = {
 
             return torch.var(self, dim, unbiased, keepdim), backward
 
+        def std_0(self,
+                  unbiased: bool=True):
+            std_out = torch.std(self, unbiased)
+            def backward(grad_output):
+                grad_self = torch.var_backward(grad_output / (std_out * 2), self, unbiased)
+                return grad_self, None
+
+            return std_out, backward
+
+        def std_1(self,
+                  dim: List[int],
+                  unbiased: bool,
+                  keepdim: bool):
+            std_out = torch.std(self, dim, unbiased, keepdim)
+            def backward(grad_output):
+                grad_self = torch.var_backward(grad_output / (std_out * 2), self, dim, unbiased, keepdim)
+                return grad_self, None, None, None
+
+            return std_out, backward
+
         def view(self,
                  size: List[int]):
             self_size = self.size()