refactor: fix docstrings and signatures of sample* methods (#1719)

janfb · web-flow · commit 299f3d06fe46 · 2026-01-23T18:37:42.000+01:00
*refactor: fix docstrings and signatures of sample* methods

- remove deprecated args
- make base class consistent
- clarify formulations in docstring.

* fix: remove unused sampling args from vi tests.

* fix pyright issues

* fix fromreview comments
diff --git a/sbi/inference/posteriors/base_posterior.py b/sbi/inference/posteriors/base_posterior.py
@@ -106,21 +106,47 @@ def sample(
         sample_shape: Shape = torch.Size(),
         x: Optional[Tensor] = None,
         show_progress_bars: bool = True,
-        mcmc_method: Optional[str] = None,
-        mcmc_parameters: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
     ) -> Tensor:
-        """See child classes for docstring."""
+        r"""Draw samples from the approximate posterior distribution $p(\theta|x)$.
+
+        Args:
+            sample_shape: Shape of samples to draw.
+            x: Conditioning observation $x_o$. If not provided, uses the default
+                `x` set via `.set_default_x()`.
+            show_progress_bars: Whether to show a progress bar during sampling.
+            **kwargs: Additional keyword arguments passed to the specific
+                posterior's sampling method. See the docstring of the specific
+                posterior class for available options.
+
+        Returns:
+            Samples from the posterior with shape `(*sample_shape, *theta_shape)`.
+        """
         pass
 
     @abstractmethod
     def sample_batched(
         self,
         sample_shape: Shape,
         x: Tensor,
-        max_sampling_batch_size: int = 10_000,
         show_progress_bars: bool = True,
+        **kwargs: Any,
     ) -> Tensor:
-        """See child classes for docstring."""
+        r"""Draw samples from the posteriors for a batch of different xs.
+
+        Given a batch of observations `[x_1, ..., x_B]`, this method samples from
+        posteriors $p(\theta|x_1), \ldots, p(\theta|x_B)$ in a vectorized manner.
+
+        Args:
+            sample_shape: Shape of samples to draw for each observation.
+            x: Batch of observations with shape `(batch_dim, *event_shape_x)`.
+            show_progress_bars: Whether to show a progress bar during sampling.
+            **kwargs: Additional keyword arguments passed to the specific
+                posterior's sampling method.
+
+        Returns:
+            Samples with shape `(*sample_shape, batch_dim, *theta_shape)`.
+        """
         pass
 
     @property
diff --git a/sbi/inference/posteriors/direct_posterior.py b/sbi/inference/posteriors/direct_posterior.py
@@ -137,19 +137,19 @@ def sample(
         sample_shape: Shape = torch.Size(),
         x: Optional[Tensor] = None,
         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)$.
+        r"""Draw samples from the approximate posterior distribution $p(\theta|x)$.
 
         Args:
             sample_shape: Desired shape of samples that are drawn from posterior. If
                 sample_shape is multidimensional we simply draw `sample_shape.numel()`
                 samples and then reshape into the desired shape.
-            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.
+            x: Conditioning observation $x_o$. If not provided, uses the default `x`
+                set via `.set_default_x()`.
+            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
@@ -180,13 +180,6 @@ def sample(
             else max_sampling_batch_size
         )
 
-        if sample_with is not None:
-            raise ValueError(
-                f"You set `sample_with={sample_with}`. As of sbi v0.18.0, setting "
-                f"`sample_with` is no longer supported. You have to rerun "
-                f"`.build_posterior(sample_with={sample_with}).`"
-            )
-
         if reject_outside_prior:
             # Normal rejection behavior.
             samples = rejection.accept_reject_sample(
@@ -218,9 +211,10 @@ def sample_batched(
         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)
-        manner.
+        r"""Draw samples from the posteriors for a batch of different xs.
+
+        Given a batch of observations `[x_1, ..., x_B]`, this method samples from
+        posteriors $p(\theta|x_1), \ldots, p(\theta|x_B)$ in a vectorized manner.
 
         Args:
             sample_shape: Desired shape of samples that are drawn from the posterior
diff --git a/sbi/inference/posteriors/importance_posterior.py b/sbi/inference/posteriors/importance_posterior.py
@@ -190,14 +190,14 @@ def sample(
         method: Optional[str] = None,
         oversampling_factor: int = 32,
         max_sampling_batch_size: int = 10_000,
-        sample_with: Optional[str] = None,
         show_progress_bars: bool = False,
     ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
-        """Return samples from the approximate posterior distribution.
+        """Draw samples from the approximate posterior distribution $p(\theta|x)$.
 
         Args:
             sample_shape: Shape of samples that are drawn from posterior.
-            x: Observed data.
+            x: Conditioning observation $x_o$. If not provided, uses the default `x`
+                set via `.set_default_x()`.
             method: Either of [`sir`|`importance`]. This sets the behavior of the
                 `.sample()` method. With `sir`, approximate posterior samples are
                 generated with sampling importance resampling (SIR). With
@@ -212,13 +212,6 @@ def sample(
 
         method = self.method if method is None else method
 
-        if sample_with is not None:
-            raise ValueError(
-                f"You set `sample_with={sample_with}`. As of sbi v0.18.0, setting "
-                f"`sample_with` is no longer supported. You have to rerun "
-                f"`.build_posterior(sample_with={sample_with}).`"
-            )
-
         self.potential_fn.set_x(self._x_else_default_x(x))
 
         if method == "sir":
@@ -255,8 +248,6 @@ def _importance_sample(
 
         Args:
             sample_shape: Desired shape of samples that are drawn from posterior.
-            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.
 
         Returns:
@@ -286,13 +277,10 @@ def _sir_sample(
             sample_shape: Desired shape of samples that are drawn from posterior. If
                 sample_shape is multidimensional we simply draw `sample_shape.numel()`
                 samples and then reshape into the desired shape.
-            x: Observed data.
-            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.
-            oversampling_factor: Number of proposed samples form which only one is
+            oversampling_factor: Number of proposed samples from which only one is
                 selected based on its importance weight.
-            max_sampling_batch_size: The batchsize of samples being drawn from
-                the proposal at every iteration. Used only in `sir_sample()`.
+            max_sampling_batch_size: The batch size of samples being drawn from
+                the proposal at every iteration.
             show_progress_bars: Whether to show sampling progress monitor.
 
         Returns:
diff --git a/sbi/inference/posteriors/mcmc_posterior.py b/sbi/inference/posteriors/mcmc_posterior.py
@@ -257,29 +257,36 @@ def sample(
         num_chains: Optional[int] = None,
         init_strategy: Optional[str] = None,
         init_strategy_parameters: Optional[Dict[str, Any]] = None,
-        init_strategy_num_candidates: Optional[int] = None,
-        mcmc_parameters: Optional[Dict] = None,
-        mcmc_method: Optional[str] = None,
-        sample_with: Optional[str] = None,
         num_workers: Optional[int] = None,
         mp_context: Optional[str] = None,
         show_progress_bars: bool = True,
     ) -> Tensor:
-        r"""Return samples from posterior distribution $p(\theta|x)$ with MCMC.
-
-        Check the `__init__()` method for a description of all arguments as well as
-        their default values.
+        r"""Draw samples from the approximate posterior distribution $p(\theta|x)$.
 
         Args:
             sample_shape: Desired shape of samples that are drawn from posterior. If
                 sample_shape is multidimensional we simply draw `sample_shape.numel()`
                 samples and then reshape into the desired shape.
-            mcmc_parameters: Dictionary that is passed only to support the API of
-                `sbi` v0.17.2 or older.
-            mcmc_method: This argument only exists to keep backward-compatibility with
-                `sbi` v0.17.2 or older. Please use `method` instead.
-            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.
+            x: Conditioning observation $x_o$. If not provided, uses the default `x`
+                set via `.set_default_x()`.
+            method: MCMC method to use. One of `slice_np`, `slice_np_vectorized`,
+                `hmc_pyro`, `nuts_pyro`, `slice_pymc`, `hmc_pymc`, `nuts_pymc`.
+                If not provided, uses the method specified at initialization.
+            thin: Thinning factor for the chain. If not provided, uses the value
+                specified at initialization.
+            warmup_steps: Number of warmup steps to discard. If not provided, uses
+                the value specified at initialization.
+            num_chains: Number of MCMC chains to run. If not provided, uses the
+                value specified at initialization.
+            init_strategy: Initialization strategy for chains (`proposal`, `sir`,
+                or `resample`). If not provided, uses the value specified at
+                initialization.
+            init_strategy_parameters: Parameters for the initialization strategy.
+                If not provided, uses the value specified at initialization.
+            num_workers: Number of CPU cores for parallelization. If not provided,
+                uses the value specified at initialization.
+            mp_context: Multiprocessing context (`fork` or `spawn`). If not provided,
+                uses the value specified at initialization.
             show_progress_bars: Whether to show sampling progress monitor.
 
         Returns:
@@ -301,46 +308,6 @@ def sample(
             if init_strategy_parameters is None
             else init_strategy_parameters
         )
-        if init_strategy_num_candidates is not None:
-            warn(
-                f"Passing `init_strategy_num_candidates` is deprecated as of sbi \
-                v0.19.0. Instead, use e.g., \
-                `init_strategy_parameters={'num_candidate_samples': 1000}`",
-                stacklevel=2,
-            )
-            self.init_strategy_parameters["num_candidate_samples"] = (
-                init_strategy_num_candidates
-            )
-        if sample_with is not None:
-            raise ValueError(
-                f"You set `sample_with={sample_with}`. As of sbi v0.18.0, setting "
-                "`sample_with` is no longer supported. You have to rerun "
-                f"`.build_posterior(sample_with={sample_with}).`"
-            )
-        if mcmc_method is not None:
-            warn(
-                "You passed `mcmc_method` to `.sample()`. As of sbi v0.18.0, this "
-                "is deprecated and will be removed in a future release. Use `method` "
-                "instead of `mcmc_method`.",
-                stacklevel=2,
-            )
-            method = mcmc_method
-        if mcmc_parameters:
-            warn(
-                "You passed `mcmc_parameters` to `.sample()`. As of sbi v0.18.0, this "
-                "is deprecated and will be removed in a future release. Instead, pass "
-                "the variable to `.sample()` directly, e.g. "
-                "`posterior.sample((1,), num_chains=5)`.",
-                stacklevel=2,
-            )
-        # The following lines are only for backwards compatibility with sbi v0.17.2 or
-        # older.
-        m_p = mcmc_parameters or {}  # define to shorten the variable name
-        method = _maybe_use_dict_entry(method, "mcmc_method", m_p)
-        thin = _maybe_use_dict_entry(thin, "thin", m_p)
-        warmup_steps = _maybe_use_dict_entry(warmup_steps, "warmup_steps", m_p)
-        num_chains = _maybe_use_dict_entry(num_chains, "num_chains", m_p)
-        init_strategy = _maybe_use_dict_entry(init_strategy, "init_strategy", m_p)
         self.potential_ = self._prepare_potential(method)  # type: ignore
 
         initial_params = self._get_initial_params(
@@ -415,9 +382,10 @@ def sample_batched(
         mp_context: Optional[str] = None,
         show_progress_bars: bool = True,
     ) -> 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)
-        manner.
+        r"""Draw samples from the posteriors for a batch of different xs.
+
+        Given a batch of observations `[x_1, ..., x_B]`, this method samples from
+        posteriors $p(\theta|x_1), \ldots, p(\theta|x_B)$ in a vectorized manner.
 
         Check the `__init__()` method for a description of all arguments as well as
         their default values.
@@ -1134,24 +1102,6 @@ def _process_thin_default(thin: int) -> int:
     return thin
 
 
-def _maybe_use_dict_entry(default: Any, key: str, dict_to_check: Dict) -> Any:
-    """Returns `default` if `key` is not in the dict and otherwise the dict entry.
-
-    This method exists only to keep backwards compatibility with `sbi` v0.17.2 or
-    older. It allows passing `mcmc_parameters` to `.sample()`.
-
-    Args:
-        default: The default value if `key` is not in `dict_to_check`.
-        key: The key for which to check in `dict_to_check`.
-        dict_to_check: The dictionary to be checked.
-
-    Returns:
-        The potentially replaced value.
-    """
-    attribute = dict_to_check.get(key, default)
-    return attribute
-
-
 def _num_required_args(func):
     """
     Utility for counting the number of positional args in a function.
diff --git a/sbi/inference/posteriors/rejection_posterior.py b/sbi/inference/posteriors/rejection_posterior.py
@@ -135,19 +135,26 @@ def sample(
         num_samples_to_find_max: Optional[int] = None,
         num_iter_to_find_max: Optional[int] = None,
         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.
+        r"""Draw samples from the approximate posterior via rejection sampling.
 
         Args:
             sample_shape: Desired shape of samples that are drawn from posterior. If
                 sample_shape is multidimensional we simply draw `sample_shape.numel()`
                 samples and then reshape into the desired shape.
-            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.
+            x: Conditioning observation $x_o$. If not provided, uses the default `x`
+                set via `.set_default_x()`.
+            max_sampling_batch_size: Maximum batch size for rejection sampling.
+                If not provided, uses the value specified at initialization.
+            num_samples_to_find_max: Number of samples to find the maximum of the
+                potential function. If not provided, uses the value from initialization.
+            num_iter_to_find_max: Number of optimization iterations to find the
+                maximum. If not provided, uses the value from initialization.
+            m: Multiplier for the proposal distribution. If not provided, uses the
+                value from initialization.
             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
@@ -166,12 +173,6 @@ def sample(
 
         potential = partial(self.potential_fn, track_gradients=True)
 
-        if sample_with is not None:
-            raise ValueError(
-                f"You set `sample_with={sample_with}`. As of sbi v0.18.0, setting "
-                f"`sample_with` is no longer supported. You have to rerun "
-                f"`.build_posterior(sample_with={sample_with}).`"
-            )
         # Replace arguments that were not passed with their default.
         max_sampling_batch_size = (
             self.max_sampling_batch_size
diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py
@@ -332,12 +332,16 @@ def sample(
         self,
         sample_shape: Shape = torch.Size(),
         x: Optional[Tensor] = None,
-        **kwargs,
+        show_progress_bars: bool = True,
     ) -> Tensor:
-        """Samples from the variational posterior distribution.
+        r"""Draw samples from the variational posterior distribution $p(\theta|x)$.
 
         Args:
-            sample_shape: Shape of samples
+            sample_shape: Desired shape of samples that are drawn from the posterior.
+            x: Conditioning observation $x_o$. If not provided, uses the default `x`
+                set via `.set_default_x()`.
+            show_progress_bars: Unused for `VIPosterior` since sampling from the
+                variational distribution is fast. Included for API consistency.
 
         Returns:
             Samples from posterior.
diff --git a/tests/vi_test.py b/tests/vi_test.py
