Fix rsample bug (jankrepl#106)

deepdow/layers/allocate.py

@@ -332,7 +332,7 @@ def forward(self, matrix, rets=None, **kwargs):
         portfolios = []  # n_portfolios elements of (n_samples, n_assets)
 
         for _ in range(self.n_portfolios):
-            draws = dist.sample((n_draws,))  # (n_draws, n_samples, n_assets)
+            draws = dist.rsample((n_draws,))  # (n_draws, n_samples, n_assets)
             rets_ = draws.mean(dim=0) if rets is not None else None  # (n_samples, n_assets)
             covmat_ = CovarianceMatrix(sqrt=self.uses_sqrt)(draws.permute(1, 0, 2))  # (n_samples, n_assets, ...)
 

tests/test_layers.py

@@ -412,7 +412,7 @@ def test_basic(self, dtype_device, allocator_class, random_state):
         single_ = torch.rand(n_assets, n_assets, dtype=dtype, device=device)
         single = single_ @ single_.t()
         covmat = torch.stack([single for _ in range(n_samples)], dim=0)
-        rets = torch.rand(n_samples, n_assets, dtype=dtype, device=device)
+        rets = torch.rand(n_samples, n_assets, dtype=dtype, device=device, requires_grad=True)
 
         if allocator_class.__name__ == 'AnalyticalMarkowitz':
             allocator = allocator_class()
@@ -440,6 +440,16 @@ def test_basic(self, dtype_device, allocator_class, random_state):
         else:
             assert torch.allclose(weights_1, weights_2)
 
+        # Make sure one can run backward pass (just sum the weights to get a scalar)
+        some_loss = weights_1.sum()
+
+        assert rets.grad is None
+
+        some_loss.backward()
+
+        assert rets.grad is not None
+        assert single.grad is None
+
 
 class TestRNN:
     @pytest.mark.parametrize('bidirectional', [True, False], ids=['bidir', 'onedir'])