Add max_sampling_time support to rejection samplers and corresponding tests (#1705)

* Add max_sampling_time timeout support to rejection samplers

* Add timeout tests for rejection and accept_reject samplers

* Add Licence header to new test file

* refine timeout RuntimeError messages per reviewer suggestion

* Add timeout handling and reject_outside_prior option to rejection sampling

* add reject_outside_prior_support logic and update rejection sampling tests

* warn if sampling without rejection

* add tests for warnings

* improve docstrings.

---------

Co-authored-by: Jan <jan.boelts@mailbox.org>

Author: satwiksps

sbi/inference/posteriors/direct_posterior.py

@@ -19,7 +19,7 @@
 )
 from sbi.samplers.rejection import rejection
 from sbi.sbi_types import Shape
-from sbi.utils.sbiutils import within_support
+from sbi.utils.sbiutils import warn_if_outside_prior_support, within_support
 from sbi.utils.torchutils import ensure_theta_batched
 from sbi.utils.user_input_checks import check_prior
 
@@ -139,6 +139,8 @@ def sample(
         max_sampling_batch_size: int = 10_000,
         sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ) -> Tensor:
         r"""Return samples from posterior distribution $p(\theta|x)$.
 
@@ -149,6 +151,14 @@ def sample(
             sample_with: This argument only exists to keep backward-compatibility with
                 `sbi` v0.17.2 or older. If it is set, we instantly raise an error.
             show_progress_bars: Whether to show sampling progress monitor.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the neural density estimator without rejection, which is
+                faster but may include samples outside the prior support.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised. Only
+                applies when `reject_outside_prior=True` (no effect otherwise since
+                direct sampling is fast).
         """
         num_samples = torch.Size(sample_shape).numel()
         x = self._x_else_default_x(x)
@@ -177,15 +187,25 @@ def sample(
                 f"`.build_posterior(sample_with={sample_with}).`"
             )
 
-        samples = rejection.accept_reject_sample(
-            proposal=self.posterior_estimator.sample,
-            accept_reject_fn=lambda theta: within_support(self.prior, theta),
-            num_samples=num_samples,
-            show_progress_bars=show_progress_bars,
-            max_sampling_batch_size=max_sampling_batch_size,
-            proposal_sampling_kwargs={"condition": x},
-            alternative_method="build_posterior(..., sample_with='mcmc')",
-        )[0]  # [0] to return only samples, not acceptance probabilities.
+        if reject_outside_prior:
+            # Normal rejection behavior.
+            samples = rejection.accept_reject_sample(
+                proposal=self.posterior_estimator.sample,
+                accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                num_samples=num_samples,
+                show_progress_bars=show_progress_bars,
+                max_sampling_batch_size=max_sampling_batch_size,
+                proposal_sampling_kwargs={"condition": x},
+                alternative_method="build_posterior(..., sample_with='mcmc')",
+                max_sampling_time=max_sampling_time,
+            )[0]
+        else:
+            # Bypass rejection sampling entirely.
+            samples = self.posterior_estimator.sample(
+                torch.Size([num_samples]),
+                condition=x,
+            )
+            warn_if_outside_prior_support(self.prior, samples[:, 0])
 
         return samples[:, 0]  # Remove batch dimension.
 
@@ -195,6 +215,8 @@ def sample_batched(
         x: Tensor,
         max_sampling_batch_size: int = 10_000,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ) -> Tensor:
         r"""Given a batch of observations [x_1, ..., x_B] this function samples from
         posteriors $p(\theta|x_1)$, ... ,$p(\theta|x_B)$, in a batched (i.e. vectorized)
@@ -207,6 +229,13 @@ def sample_batched(
                 `batch_dim` corresponds to the number of observations to be drawn.
             max_sampling_batch_size: Maximum batch size for rejection sampling.
             show_progress_bars: Whether to show sampling progress monitor.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the neural density estimator without rejection, which is
+                faster but may include samples outside the prior support.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised. Only
+                applies when `reject_outside_prior=True`.
 
         Returns:
             Samples from the posteriors of shape (*sample_shape, B, *input_shape)
@@ -242,15 +271,25 @@ def sample_batched(
             )
             max_sampling_batch_size = capped
 
-        samples = rejection.accept_reject_sample(
-            proposal=self.posterior_estimator.sample,
-            accept_reject_fn=lambda theta: within_support(self.prior, theta),
-            num_samples=num_samples,
-            show_progress_bars=show_progress_bars,
-            max_sampling_batch_size=max_sampling_batch_size,
-            proposal_sampling_kwargs={"condition": x},
-            alternative_method="build_posterior(..., sample_with='mcmc')",
-        )[0]
+        if reject_outside_prior:
+            # Normal rejection behavior.
+            samples = rejection.accept_reject_sample(
+                proposal=self.posterior_estimator.sample,
+                accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                num_samples=num_samples,
+                show_progress_bars=show_progress_bars,
+                max_sampling_batch_size=max_sampling_batch_size,
+                proposal_sampling_kwargs={"condition": x},
+                alternative_method="build_posterior(..., sample_with='mcmc')",
+                max_sampling_time=max_sampling_time,
+            )[0]
+        else:
+            # Bypass rejection sampling entirely.
+            samples = self.posterior_estimator.sample(
+                torch.Size([num_samples]),
+                condition=x,
+            )
+            warn_if_outside_prior_support(self.prior, samples)
 
         return samples
 

sbi/inference/posteriors/rejection_posterior.py

@@ -137,6 +137,8 @@ def sample(
         m: Optional[float] = None,
         sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ):
         r"""Return samples from posterior $p(\theta|x)$ via rejection sampling.
 
@@ -147,6 +149,14 @@ def sample(
             sample_with: This argument only exists to keep backward-compatibility with
                 `sbi` v0.17.2 or older. If it is set, we instantly raise an error.
             show_progress_bars: Whether to show sampling progress monitor.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the proposal without rejection, which is faster but may
+                include samples outside the prior support.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised. Only
+                applies when `reject_outside_prior=True` (no effect otherwise since
+                direct sampling from the proposal is fast).
 
         Returns:
             Samples from posterior.
@@ -180,18 +190,29 @@ def sample(
         )
         m = self.m if m is None else m
 
-        samples, _ = rejection_sample(
-            potential,
-            proposal=self.proposal,
-            num_samples=num_samples,
-            show_progress_bars=show_progress_bars,
-            warn_acceptance=0.01,
-            max_sampling_batch_size=max_sampling_batch_size,
-            num_samples_to_find_max=num_samples_to_find_max,
-            num_iter_to_find_max=num_iter_to_find_max,
-            m=m,
-            device=self._device,
-        )
+        if reject_outside_prior:
+            samples, _ = rejection_sample(
+                potential,
+                proposal=self.proposal,
+                num_samples=num_samples,
+                show_progress_bars=show_progress_bars,
+                warn_acceptance=0.01,
+                max_sampling_batch_size=max_sampling_batch_size,
+                num_samples_to_find_max=num_samples_to_find_max,
+                num_iter_to_find_max=num_iter_to_find_max,
+                m=m,
+                max_sampling_time=max_sampling_time,
+                device=self._device,
+            )
+        else:
+            # Bypass rejection sampling entirely.
+            samples = self.proposal.sample((num_samples,))
+            warn(
+                "Samples drawn with reject_outside_prior=False are taken directly "
+                "from the proposal without rejection sampling. These samples may lie "
+                "outside the prior support, which could lead to incorrect inference.",
+                stacklevel=2,
+            )
 
         return samples.reshape((*sample_shape, -1))
 

sbi/inference/posteriors/vector_field_posterior.py

@@ -25,7 +25,11 @@
 from sbi.samplers.score.predictors import Predictor
 from sbi.sbi_types import Shape
 from sbi.utils import check_prior
-from sbi.utils.sbiutils import gradient_ascent, within_support
+from sbi.utils.sbiutils import (
+    gradient_ascent,
+    warn_if_outside_prior_support,
+    within_support,
+)
 from sbi.utils.torchutils import ensure_theta_batched
 
 
@@ -158,6 +162,8 @@ def sample(
         max_sampling_batch_size: int = 10_000,
         sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ) -> Tensor:
         r"""Return samples from posterior distribution $p(\theta|x)$.
 
@@ -195,6 +201,14 @@ def sample(
                 use the 'sde' sampling method, the vector field estimator must support
                 it and have the SCORE_DEFINED class attribute set to True.
             show_progress_bars: Whether to show a progress bar during sampling.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the ODE/SDE sampler without rejection, which is faster but
+                may include samples outside the prior support.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised. Only
+                applies when `reject_outside_prior=True` (no effect otherwise since
+                direct sampling does not use rejection).
         """
 
         if sample_with is None:
@@ -213,13 +227,18 @@ def sample(
         num_samples = torch.Size(sample_shape).numel()
 
         if sample_with == "ode":
-            samples, _ = rejection.accept_reject_sample(
-                proposal=self.sample_via_ode,
-                accept_reject_fn=lambda theta: within_support(self.prior, theta),
-                num_samples=num_samples,
-                show_progress_bars=show_progress_bars,
-                max_sampling_batch_size=max_sampling_batch_size,
-            )
+            if reject_outside_prior:
+                samples, _ = rejection.accept_reject_sample(
+                    proposal=self.sample_via_ode,
+                    accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                    num_samples=num_samples,
+                    show_progress_bars=show_progress_bars,
+                    max_sampling_batch_size=max_sampling_batch_size,
+                    max_sampling_time=max_sampling_time,
+                )
+            else:
+                # Bypass rejection sampling entirely.
+                samples = self.sample_via_ode(torch.Size([num_samples]))
         elif sample_with == "sde":
             proposal_sampling_kwargs = {
                 "predictor": predictor,
@@ -231,19 +250,30 @@ def sample(
                 "max_sampling_batch_size": max_sampling_batch_size,
                 "show_progress_bars": show_progress_bars,
             }
-            samples, _ = rejection.accept_reject_sample(
-                proposal=self._sample_via_diffusion,
-                accept_reject_fn=lambda theta: within_support(self.prior, theta),
-                num_samples=num_samples,
-                show_progress_bars=show_progress_bars,
-                max_sampling_batch_size=max_sampling_batch_size,
-                proposal_sampling_kwargs=proposal_sampling_kwargs,
-            )
+            if reject_outside_prior:
+                samples, _ = rejection.accept_reject_sample(
+                    proposal=self._sample_via_diffusion,
+                    accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                    num_samples=num_samples,
+                    show_progress_bars=show_progress_bars,
+                    max_sampling_batch_size=max_sampling_batch_size,
+                    proposal_sampling_kwargs=proposal_sampling_kwargs,
+                    max_sampling_time=max_sampling_time,
+                )
+            else:
+                # Bypass rejection sampling entirely.
+                samples = self._sample_via_diffusion(
+                    (num_samples,),
+                    **proposal_sampling_kwargs,
+                )
         else:
             raise ValueError(
                 f"Expected sample_with to be 'ode' or 'sde', but got {sample_with}."
             )
 
+        if not reject_outside_prior:
+            warn_if_outside_prior_support(self.prior, samples)
+
         samples = samples.reshape(
             sample_shape + self.vector_field_estimator.input_shape
         )
@@ -427,6 +457,8 @@ def sample_batched(
         ts: Optional[Tensor] = None,
         max_sampling_batch_size: int = 10000,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ) -> Tensor:
         r"""Given a batch of observations [x_1, ..., x_B] this function samples from
         posteriors $p(\theta|x_1)$, ... ,$p(\theta|x_B)$, in a batched (i.e. vectorized)
@@ -449,6 +481,13 @@ def sample_batched(
                 linear grid between t_max and t_min is used.
             max_sampling_batch_size: Maximum batch size for sampling.
             show_progress_bars: Whether to show sampling progress monitor.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the ODE/SDE sampler without rejection, which is faster but
+                may include samples outside the prior support.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised. Only
+                applies when `reject_outside_prior=True`.
 
         Returns:
             Samples from the posteriors of shape (*sample_shape, B, *input_shape)
@@ -477,14 +516,19 @@ def sample_batched(
             max_sampling_batch_size = capped
 
         if self.sample_with == "ode":
-            samples, _ = rejection.accept_reject_sample(
-                proposal=self.sample_via_ode,
-                accept_reject_fn=lambda theta: within_support(self.prior, theta),
-                num_samples=num_samples,
-                num_xos=batch_size,
-                show_progress_bars=show_progress_bars,
-                max_sampling_batch_size=max_sampling_batch_size,
-            )
+            if reject_outside_prior:
+                samples, _ = rejection.accept_reject_sample(
+                    proposal=self.sample_via_ode,
+                    accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                    num_samples=num_samples,
+                    num_xos=batch_size,
+                    show_progress_bars=show_progress_bars,
+                    max_sampling_batch_size=max_sampling_batch_size,
+                    max_sampling_time=max_sampling_time,
+                )
+            else:
+                # Bypass rejection sampling.
+                samples = self.sample_via_ode(torch.Size([num_samples]))
             samples = samples.reshape(
                 sample_shape + batch_shape + self.vector_field_estimator.input_shape
             )
@@ -499,19 +543,29 @@ def sample_batched(
                 "max_sampling_batch_size": max_sampling_batch_size,
                 "show_progress_bars": show_progress_bars,
             }
-            samples, _ = rejection.accept_reject_sample(
-                proposal=self._sample_via_diffusion,
-                accept_reject_fn=lambda theta: within_support(self.prior, theta),
-                num_samples=num_samples,
-                num_xos=batch_size,
-                show_progress_bars=show_progress_bars,
-                max_sampling_batch_size=max_sampling_batch_size,
-                proposal_sampling_kwargs=proposal_sampling_kwargs,
-            )
+            if reject_outside_prior:
+                samples, _ = rejection.accept_reject_sample(
+                    proposal=self._sample_via_diffusion,
+                    accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                    num_samples=num_samples,
+                    num_xos=batch_size,
+                    show_progress_bars=show_progress_bars,
+                    max_sampling_batch_size=max_sampling_batch_size,
+                    proposal_sampling_kwargs=proposal_sampling_kwargs,
+                    max_sampling_time=max_sampling_time,
+                )
+            else:
+                # Bypass rejection sampling.
+                samples = self._sample_via_diffusion(
+                    (num_samples,), **proposal_sampling_kwargs
+                )
             samples = samples.reshape(
                 sample_shape + batch_shape + self.vector_field_estimator.input_shape
             )
 
+        if not reject_outside_prior:
+            warn_if_outside_prior_support(self.prior, samples)
+
         return samples
 
     def map(