Fix LKJ sample shape bug

pyro/distributions/lkj.py

@@ -64,7 +64,8 @@ def __init__(self, d, eta, validate_args=None):
         super().__init__(torch.Size(), torch.Size((d, d)), validate_args=validate_args)
 
     def sample(self, sample_shape=torch.Size()):
-        y = self._gen.sample(sample_shape=self.batch_shape + sample_shape).detach()
+        with torch.no_grad():
+            y = self._gen.sample(sample_shape=sample_shape + self.batch_shape)
         z = y.mul(2).add(-1.0)
         return _vector_to_l_cholesky(z)
 

tests/distributions/test_lkj.py

@@ -109,3 +109,17 @@ def test_log_prob_d2(eta):
     tst = test_dist.log_prob(x)
 
     assert_tensors_equal(lp, tst, prec=1e-6)
+
+
+def test_sample_batch():
+    # Regression test for https://github.com/pyro-ppl/pyro/issues/2615
+    dist = LKJCorrCholesky(d=3, eta=torch.ones(())).expand([12])
+    # batch shape and event shape are as you'd expect
+    assert dist.batch_shape == torch.Size([12])
+    assert dist.event_shape == torch.Size([3, 3])
+    # samples have correct shape when sample_shape=()
+    assert dist.shape(()) == torch.Size([12, 3, 3])
+    assert dist.sample().shape == torch.Size([12, 3, 3])
+    # samples had the wrong shape when sample_shape is non-unit
+    assert dist.shape((4,)) == torch.Size([4, 12, 3, 3])
+    assert dist.sample((4,)).shape == torch.Size([4, 12, 3, 3])