Repository: pyro-ppl/pyro

Issue:
Support arbitrary sample_shape in Rejector
This will require the `proposal` to accept an optional `sample_shape` arg and the `log_prob_accept()` method to support broadcasting.


Discussion:

Resolving patch:
diff --git a/pyro/distributions/rejector.py b/pyro/distributions/rejector.py
--- a/pyro/distributions/rejector.py
+++ b/pyro/distributions/rejector.py
@@ -12,7 +12,8 @@ class Rejector(Distribution):
     Rejection sampled distribution given an acceptance rate function.
 
     :param Distribution propose: A proposal distribution that samples batched
-        propsals via `propose()`.
+        proposals via ``propose()``. :meth:`sample` supports a ``sample_shape``
+        arg only if ``propose()`` supports a ``sample_shape`` arg.
     :param callable log_prob_accept: A callable that inputs a batch of
         proposals and returns a batch of log acceptance probabilities.
     :param log_scale: Total log probability of acceptance.
@@ -40,15 +41,13 @@ def _propose_batch_log_pdf(self, x):
         return self._propose_batch_log_pdf_cache[1]
 
     def sample(self, sample_shape=torch.Size()):
-        if sample_shape:
-            raise ValueError("Arbitrary `sample_shape` not supported by Rejector class.")
         # Implements parallel batched accept-reject sampling.
-        x = self.propose()
+        x = self.propose(sample_shape) if sample_shape else self.propose()
         log_prob_accept = self.log_prob_accept(x)
         probs = torch.exp(log_prob_accept).clamp_(0.0, 1.0)
         done = torch.bernoulli(probs).byte()
         while not done.all():
-            proposed_x = self.propose()
+            proposed_x = self.propose(sample_shape) if sample_shape else self.propose()
             log_prob_accept = self.log_prob_accept(proposed_x)
             prob_accept = torch.exp(log_prob_accept).clamp_(0.0, 1.0)
             accept = torch.bernoulli(prob_accept).byte() & ~done


Test patch:
diff --git a/pyro/distributions/testing/rejection_gamma.py b/pyro/distributions/testing/rejection_gamma.py
--- a/pyro/distributions/testing/rejection_gamma.py
+++ b/pyro/distributions/testing/rejection_gamma.py
@@ -28,9 +28,9 @@ def __init__(self, alpha):
         log_scale = self.propose_log_prob(x) + self.log_prob_accept(x) - self.log_prob(x)
         super(RejectionStandardGamma, self).__init__(self.propose, self.log_prob_accept, log_scale)
 
-    def propose(self):
+    def propose(self, sample_shape=torch.Size()):
         # Marsaglia & Tsang's x == Naesseth's epsilon
-        x = self.alpha.new(self.alpha.shape).normal_()
+        x = self.alpha.new(sample_shape + self.alpha.shape).normal_()
         y = 1.0 + self._c * x
         v = y * y * y
         return (self._d * v).clamp_(1e-30, 1e30)
@@ -68,9 +68,7 @@ def __init__(self, alpha, beta):
         self.beta = beta
 
     def sample(self, sample_shape=torch.Size()):
-        if sample_shape:
-            raise ValueError("Arbitrary `sample_shape` not supported by RejectionGamma class.")
-        return self._standard_gamma.sample() / self.beta
+        return self._standard_gamma.sample(sample_shape) / self.beta
 
     def log_prob(self, x):
         return self._standard_gamma.log_prob(x * self.beta) + torch.log(self.beta)
@@ -100,9 +98,7 @@ def __init__(self, alpha, beta, boost=1):
         self._unboost_x_cache = None, None
 
     def sample(self, sample_shape=torch.Size()):
-        if sample_shape:
-            raise ValueError("Arbitrary `sample_shape` not supported by ShapeAugmentedGamma class.")
-        x = self._rejection_gamma.sample()
+        x = self._rejection_gamma.sample(sample_shape)
         boosted_x = x.clone()
         for i in range(self._boost):
             boosted_x *= (1 - x.new(x.shape).uniform_()) ** (1 / (i + self.alpha))
diff --git a/tests/distributions/test_rejector.py b/tests/distributions/test_rejector.py
--- a/tests/distributions/test_rejector.py
+++ b/tests/distributions/test_rejector.py
@@ -9,6 +9,27 @@
 from pyro.distributions.testing.rejection_gamma import RejectionGamma, RejectionStandardGamma, ShapeAugmentedGamma
 from tests.common import assert_equal
 
+SIZES = list(map(torch.Size, [[], [1], [2], [3], [1, 1], [1, 2], [2, 3, 4]]))
+
+
+@pytest.mark.parametrize('sample_shape', SIZES)
+@pytest.mark.parametrize('batch_shape', filter(bool, SIZES))
+def test_rejection_standard_gamma_sample_shape(sample_shape, batch_shape):
+    alphas = Variable(torch.ones(batch_shape))
+    dist = RejectionStandardGamma(alphas)
+    x = dist.sample(sample_shape)
+    assert x.shape == sample_shape + batch_shape
+
+
+@pytest.mark.parametrize('sample_shape', SIZES)
+@pytest.mark.parametrize('batch_shape', filter(bool, SIZES))
+def test_rejection_exponential_sample_shape(sample_shape, batch_shape):
+    rates = Variable(torch.ones(batch_shape))
+    factors = Variable(torch.ones(batch_shape)) * 0.5
+    dist = RejectionExponential(rates, factors)
+    x = dist.sample(sample_shape)
+    assert x.shape == sample_shape + batch_shape
+
 
 def compute_elbo_grad(model, guide, variables):
     x = guide.sample()