Repository: pytorch/pytorch

Issue:
[distributions] rsample().detach() and sample() yields different gradients
I have a MultivariateNormal distribution with loc defined as output of neural net (given input) and diagonal covariance matrix with trainable parameters (but does not depending on some input). 
If I sample via `distr.rsample().detach()` and optimise sum of log_probs, `.backward()` provides correct gradients w.r.t. both loc and cov matrix params. But if I sample via `distr.sample()`, `.backward()` sets `None` gradients for cov matrix params.

Here is a minimal reproducing code:
```
import torch
from torch.distributions import MultivariateNormal

torch.manual_seed(124)
inp = torch.randn(3)
W = torch.randn((2, 3), requires_grad=True)
loc = W@inp
cov = torch.tensor([[1., 0.], [0., 2.]], requires_grad=True)
distr = MultivariateNormal(loc, cov)

y = distr.rsample().detach()  # difference here
loss = distr.log_prob(y)
loss.backward()
W.grad, cov.grad
```

yields 

```
(tensor([[-0.1991, -1.0776, -0.6339],
         [ 0.3455,  1.8700,  1.1000]]), 
 tensor([[-0.2678,  0.0000],
         [-0.8057,  0.4491]]))
```

while 

```
import torch
from torch.distributions import MultivariateNormal

torch.manual_seed(124)
inp = torch.randn(3)
W = torch.randn((2, 3), requires_grad=True)
loc = W@inp
cov = torch.tensor([[1., 0.], [0., 2.]], requires_grad=True)
distr = MultivariateNormal(loc, cov)

x = distr.sample()  # difference here
loss = distr.log_prob(x)
loss.backward()
W.grad, cov.grad
```
yields
```
(tensor([[-0.1991, -1.0776, -0.6339],
         [ 0.3455,  1.8700,  1.1000]]), 
 None)
```

I am observing this problem both on macOS and Ubuntu systems.


Discussion:
Thanks for reporting. I know the problem. The `MultivariateNormal` uses a lazily calculated property `scale_tril` to compute everything, including the samples, log prob, etc. The `sample()` func in distributions is just `rsample()` wrapped in `torch.no_grad` (see [here](https://github.com/pytorch/pytorch/blob/master/torch/distributions/distribution.py#L96-L97)). So when you do `distr.sample()`, it first calculates `scale_tril` **in a `no_grad` block**, and subsequent computations, which all use `scale_tril`, lose history.

Therefore, if you do 
```py
import torch
from torch.distributions import MultivariateNormal

torch.manual_seed(124)
inp = torch.randn(3)
W = torch.randn((2, 3), requires_grad=True)
loc = W@inp
cov = torch.tensor([[1., 0.], [0., 2.]], requires_grad=True)
distr = MultivariateNormal(loc, cov)
distr.scale_tril  # <------ force it to calculate scale_tril here

x = distr.sample()  # difference here
loss = distr.log_prob(x)
loss.backward()
W.grad, cov.grad
```

It will give correct gradients. 

Fix incoming.

Resolving patch:
diff --git a/torch/distributions/utils.py b/torch/distributions/utils.py
--- 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


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- 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():