[distributions] Always enable grad when calculating lazy_property (#7708)

ssnl · web-flow · commit c946db16ecfa · 2018-05-24T11:22:39.000-04:00
* Always enable grad when calculating lazy_property

* Add test with MultiVariableNormal
diff --git a/test/test_distributions.py b/test/test_distributions.py
@@ -53,7 +53,7 @@
                                             SoftmaxTransform,
                                             StickBreakingTransform,
                                             identity_transform)
-from torch.distributions.utils import _finfo, probs_to_logits, softmax
+from torch.distributions.utils import _finfo, probs_to_logits, softmax, lazy_property
 
 TEST_NUMPY = True
 try:
@@ -690,6 +690,31 @@ def test_enumerate_support_type(self):
                 except NotImplementedError:
                     pass
 
+    def test_lazy_property_grad(self):
+        x = torch.randn(1, requires_grad=True)
+
+        class Dummy(object):
+            @lazy_property
+            def y(self):
+                return x + 1
+
+        def test():
+            x.grad = None
+            Dummy().y.backward()
+            self.assertEqual(x.grad, torch.ones(1))
+
+        test()
+        with torch.no_grad():
+            test()
+
+        mean = torch.randn(2)
+        cov = torch.eye(2, requires_grad=True)
+        distn = MultivariateNormal(mean, cov)
+        with torch.no_grad():
+            distn.scale_tril
+        distn.scale_tril.sum().backward()
+        self.assertIsNotNone(cov.grad)
+
     def test_has_examples(self):
         distributions_with_examples = set(e.Dist for e in EXAMPLES)
         for Dist in globals().values():
diff --git a/torch/distributions/utils.py b/torch/distributions/utils.py
@@ -182,6 +182,7 @@ def __init__(self, wrapped):
     def __get__(self, instance, obj_type=None):
         if instance is None:
             return self
-        value = self.wrapped(instance)
+        with torch.enable_grad():
+            value = self.wrapped(instance)
         setattr(instance, self.wrapped.__name__, value)
         return value
