prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN 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.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `pyro/distributions/rejector.py`
  Content:
  ```
  1 from __future__ import absolute_import, division, print_function
  2 
  3 import torch
  4 from pyro.distributions.distribution import Distribution
  5 from pyro.distributions.score_parts import ScoreParts
  6 from pyro.distributions.util import copy_docs_from
  7 
  8 
  9 @copy_docs_from(Distribution)
  10 class Rejector(Distribution):
  11     """
  12     Rejection sampled distribution given an acceptance rate function.
  13 
  14     :param Distribution propose: A proposal distribution that samples batched
  15         propsals via `propose()`.
  16     :param callable log_prob_accept: A callable that inputs a batch of
  17         proposals and returns a batch of log acceptance probabilities.
  18     :param log_scale: Total log probability of acceptance.
  19     """
  20     stateful = True
  21     reparameterized = True
  22 
  23     def __init__(self, propose, log_prob_accept, log_scale):
  24         self.propose = propose
  25         self.log_prob_accept = log_prob_accept
  26         self._log_scale = log_scale
  27 
  28         # These LRU(1) caches allow work to be shared across different method calls.
  29         self._log_prob_accept_cache = None, None
  30         self._propose_batch_log_pdf_cache = None, None
  31 
  32     def _log_prob_accept(self, x):
  33         if x is not self._log_prob_accept_cache[0]:
  34             self._log_prob_accept_cache = x, self.log_prob_accept(x) - self._log_scale
  35         return self._log_prob_accept_cache[1]
  36 
  37     def _propose_batch_log_pdf(self, x):
  38         if x is not self._propose_batch_log_pdf_cache[0]:
  39             self._propose_batch_log_pdf_cache = x, self.propose.log_prob(x)
  40         return self._propose_batch_log_pdf_cache[1]
  41 
  42     def sample(self, sample_shape=torch.Size()):
  43         if sample_shape:
  44             raise ValueError("Arbitrary `sample_shape` not supported by Rejector class.")
  45         # Implements parallel batched accept-reject sampling.
  46         x = self.propose()
  47         log_prob_accept = self.log_prob_accept(x)
  48         probs = torch.exp(log_prob_accept).clamp_(0.0, 1.0)
  49         done = torch.bernoulli(probs).byte()
  50         while not done.all():
  51             proposed_x = self.propose()
  52             log_prob_accept = self.log_prob_accept(proposed_x)
  53             prob_accept = torch.exp(log_prob_accept).clamp_(0.0, 1.0)
  54             accept = torch.bernoulli(prob_accept).byte() & ~done
  55             if accept.any():
  56                 x[accept] = proposed_x[accept]
  57                 done |= accept
  58         return x
  59 
  60     def log_prob(self, x):
  61         return self._propose_batch_log_pdf(x) + self._log_prob_accept(x)
  62 
  63     def score_parts(self, x):
  64         score_function = self._log_prob_accept(x)
  65         log_pdf = self.log_prob(x)
  66         return ScoreParts(log_pdf, score_function, log_pdf)
  67
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: 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 @@
       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 @@
           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
  
verification_info: {"golden_diff": "diff --git a/pyro/distributions/rejector.py b/pyro/distributions/rejector.py\n--- a/pyro/distributions/rejector.py\n+++ b/pyro/distributions/rejector.py\n@@ -12,7 +12,8 @@\n     Rejection sampled distribution given an acceptance rate function.\n \n     :param Distribution propose: A proposal distribution that samples batched\n-        propsals via `propose()`.\n+        proposals via ``propose()``. :meth:`sample` supports a ``sample_shape``\n+        arg only if ``propose()`` supports a ``sample_shape`` arg.\n     :param callable log_prob_accept: A callable that inputs a batch of\n         proposals and returns a batch of log acceptance probabilities.\n     :param log_scale: Total log probability of acceptance.\n@@ -40,15 +41,13 @@\n         return self._propose_batch_log_pdf_cache[1]\n \n     def sample(self, sample_shape=torch.Size()):\n-        if sample_shape:\n-            raise ValueError(\"Arbitrary `sample_shape` not supported by Rejector class.\")\n         # Implements parallel batched accept-reject sampling.\n-        x = self.propose()\n+        x = self.propose(sample_shape) if sample_shape else self.propose()\n         log_prob_accept = self.log_prob_accept(x)\n         probs = torch.exp(log_prob_accept).clamp_(0.0, 1.0)\n         done = torch.bernoulli(probs).byte()\n         while not done.all():\n-            proposed_x = self.propose()\n+            proposed_x = self.propose(sample_shape) if sample_shape else self.propose()\n             log_prob_accept = self.log_prob_accept(proposed_x)\n             prob_accept = torch.exp(log_prob_accept).clamp_(0.0, 1.0)\n             accept = torch.bernoulli(prob_accept).byte() & ~done\n", "issue": "Support arbitrary sample_shape in Rejector\nThis will require the `proposal` to accept an optional `sample_shape` arg and the `log_prob_accept()` method to support broadcasting.\n", "before_files": [{"content": "from __future__ import absolute_import, division, print_function\n\nimport torch\nfrom pyro.distributions.distribution import Distribution\nfrom pyro.distributions.score_parts import ScoreParts\nfrom pyro.distributions.util import copy_docs_from\n\n\n@copy_docs_from(Distribution)\nclass Rejector(Distribution):\n    \"\"\"\n    Rejection sampled distribution given an acceptance rate function.\n\n    :param Distribution propose: A proposal distribution that samples batched\n        propsals via `propose()`.\n    :param callable log_prob_accept: A callable that inputs a batch of\n        proposals and returns a batch of log acceptance probabilities.\n    :param log_scale: Total log probability of acceptance.\n    \"\"\"\n    stateful = True\n    reparameterized = True\n\n    def __init__(self, propose, log_prob_accept, log_scale):\n        self.propose = propose\n        self.log_prob_accept = log_prob_accept\n        self._log_scale = log_scale\n\n        # These LRU(1) caches allow work to be shared across different method calls.\n        self._log_prob_accept_cache = None, None\n        self._propose_batch_log_pdf_cache = None, None\n\n    def _log_prob_accept(self, x):\n        if x is not self._log_prob_accept_cache[0]:\n            self._log_prob_accept_cache = x, self.log_prob_accept(x) - self._log_scale\n        return self._log_prob_accept_cache[1]\n\n    def _propose_batch_log_pdf(self, x):\n        if x is not self._propose_batch_log_pdf_cache[0]:\n            self._propose_batch_log_pdf_cache = x, self.propose.log_prob(x)\n        return self._propose_batch_log_pdf_cache[1]\n\n    def sample(self, sample_shape=torch.Size()):\n        if sample_shape:\n            raise ValueError(\"Arbitrary `sample_shape` not supported by Rejector class.\")\n        # Implements parallel batched accept-reject sampling.\n        x = self.propose()\n        log_prob_accept = self.log_prob_accept(x)\n        probs = torch.exp(log_prob_accept).clamp_(0.0, 1.0)\n        done = torch.bernoulli(probs).byte()\n        while not done.all():\n            proposed_x = self.propose()\n            log_prob_accept = self.log_prob_accept(proposed_x)\n            prob_accept = torch.exp(log_prob_accept).clamp_(0.0, 1.0)\n            accept = torch.bernoulli(prob_accept).byte() & ~done\n            if accept.any():\n                x[accept] = proposed_x[accept]\n                done |= accept\n        return x\n\n    def log_prob(self, x):\n        return self._propose_batch_log_pdf(x) + self._log_prob_accept(x)\n\n    def score_parts(self, x):\n        score_function = self._log_prob_accept(x)\n        log_pdf = self.log_prob(x)\n        return ScoreParts(log_pdf, score_function, log_pdf)\n", "path": "pyro/distributions/rejector.py"}], "after_files": [{"content": "from __future__ import absolute_import, division, print_function\n\nimport torch\nfrom pyro.distributions.distribution import Distribution\nfrom pyro.distributions.score_parts import ScoreParts\nfrom pyro.distributions.util import copy_docs_from\n\n\n@copy_docs_from(Distribution)\nclass Rejector(Distribution):\n    \"\"\"\n    Rejection sampled distribution given an acceptance rate function.\n\n    :param Distribution propose: A proposal distribution that samples batched\n        proposals via ``propose()``. :meth:`sample` supports a ``sample_shape``\n        arg only if ``propose()`` supports a ``sample_shape`` arg.\n    :param callable log_prob_accept: A callable that inputs a batch of\n        proposals and returns a batch of log acceptance probabilities.\n    :param log_scale: Total log probability of acceptance.\n    \"\"\"\n    stateful = True\n    reparameterized = True\n\n    def __init__(self, propose, log_prob_accept, log_scale):\n        self.propose = propose\n        self.log_prob_accept = log_prob_accept\n        self._log_scale = log_scale\n\n        # These LRU(1) caches allow work to be shared across different method calls.\n        self._log_prob_accept_cache = None, None\n        self._propose_batch_log_pdf_cache = None, None\n\n    def _log_prob_accept(self, x):\n        if x is not self._log_prob_accept_cache[0]:\n            self._log_prob_accept_cache = x, self.log_prob_accept(x) - self._log_scale\n        return self._log_prob_accept_cache[1]\n\n    def _propose_batch_log_pdf(self, x):\n        if x is not self._propose_batch_log_pdf_cache[0]:\n            self._propose_batch_log_pdf_cache = x, self.propose.log_prob(x)\n        return self._propose_batch_log_pdf_cache[1]\n\n    def sample(self, sample_shape=torch.Size()):\n        # Implements parallel batched accept-reject sampling.\n        x = self.propose(sample_shape) if sample_shape else self.propose()\n        log_prob_accept = self.log_prob_accept(x)\n        probs = torch.exp(log_prob_accept).clamp_(0.0, 1.0)\n        done = torch.bernoulli(probs).byte()\n        while not done.all():\n            proposed_x = self.propose(sample_shape) if sample_shape else self.propose()\n            log_prob_accept = self.log_prob_accept(proposed_x)\n            prob_accept = torch.exp(log_prob_accept).clamp_(0.0, 1.0)\n            accept = torch.bernoulli(prob_accept).byte() & ~done\n            if accept.any():\n                x[accept] = proposed_x[accept]\n                done |= accept\n        return x\n\n    def log_prob(self, x):\n        return self._propose_batch_log_pdf(x) + self._log_prob_accept(x)\n\n    def score_parts(self, x):\n        score_function = self._log_prob_accept(x)\n        log_pdf = self.log_prob(x)\n        return ScoreParts(log_pdf, score_function, log_pdf)\n", "path": "pyro/distributions/rejector.py"}]}