fix: cap max_sampling_batch_size to prevent excessive memory (#1624)

* add warnings
* review fixes, add test.

Author: janfb

sbi/inference/posteriors/direct_posterior.py

@@ -219,8 +219,8 @@ def sample_batched(
         # throw warning if num_x * num_samples is too large
         if num_xos * num_samples > 2**21:  # 2 million-ish
             warnings.warn(
-                "Note that for batched sampling, the direct posterior sampling "
-                "generates {num_xos} * {num_samples} = {num_xos * num_samples} "
+                f"Note that for batched sampling, the direct posterior sampling "
+                f"generates {num_xos} * {num_samples} = {num_xos * num_samples} "
                 "samples. This can be slow and memory-intensive. Consider "
                 "reducing the number of samples or batch size.",
                 stacklevel=2,
@@ -232,6 +232,16 @@ def sample_batched(
             else max_sampling_batch_size
         )
 
+        # Adjust max_sampling_batch_size to avoid excessive memory usage
+        if max_sampling_batch_size * num_xos > 100_000:
+            capped = max(1, 100_000 // num_xos)
+            warnings.warn(
+                f"Capping max_sampling_batch_size from {max_sampling_batch_size} "
+                f"to {capped} to avoid excessive memory usage.",
+                stacklevel=2,
+            )
+            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),

sbi/inference/posteriors/vector_field_posterior.py

@@ -1,6 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+import warnings
 from typing import Dict, Literal, Optional, Union
 
 import torch
@@ -431,6 +432,16 @@ def sample_batched(
             else max_sampling_batch_size
         )
 
+        # Adjust max_sampling_batch_size to avoid excessive memory usage
+        if max_sampling_batch_size * batch_size > 100_000:
+            capped = max(1, 100_000 // batch_size)
+            warnings.warn(
+                f"Capping max_sampling_batch_size from {max_sampling_batch_size} "
+                f"to {capped} to avoid excessive memory usage.",
+                stacklevel=2,
+            )
+            max_sampling_batch_size = capped
+
         if self.sample_with == "ode":
             samples = rejection.accept_reject_sample(
                 proposal=self.sample_via_ode,

sbi/samplers/rejection/rejection.py

@@ -215,8 +215,8 @@ def accept_reject_sample(
 
     Args:
         proposal: A callable that takes `sample_shape` as arguments (and kwargs as
-        needed). Returns samples from the proposal distribution with shape
-        (*sample_shape, event_dim).
+            needed). Returns samples from the proposal distribution with shape
+            (*sample_shape, event_dim).
         accept_reject_fn: Function that evaluates which samples are accepted or
             rejected. Must take a batch of parameters and return a boolean tensor which
             indicates which parameters get accepted.
@@ -272,7 +272,6 @@ def accept_reject_sample(
     accepted = [[] for _ in range(num_xos)]
     acceptance_rate = torch.full((num_xos,), float("Nan"))
     leakage_warning_raised = False
-    # Ruff suggestion
 
     # To cover cases with few samples without leakage:
     sampling_batch_size = min(num_samples, max_sampling_batch_size)

tests/posterior_nn_test.py

@@ -102,19 +102,19 @@ def test_importance_posterior_sample_log_prob(snplre_method: type):
 
 @pytest.mark.parametrize("snpe_method", [NPE_A, NPE_C])
 @pytest.mark.parametrize("x_o_batch_dim", (0, 1, 2))
-@pytest.mark.parametrize("prior", ("mvn", "uniform"))
+@pytest.mark.parametrize("prior_type", ("mvn", "uniform"))
 def test_batched_sample_log_prob_with_different_x(
     snpe_method: type,
     x_o_batch_dim: bool,
-    prior: str,
+    prior_type: str,
 ):
     num_dim = 2
     num_simulations = 1000
 
     # We also want to test on bounded support! Which will invoke leakage correction.
-    if prior == "mvn":
+    if prior_type == "mvn":
         prior = MultivariateNormal(loc=zeros(num_dim), covariance_matrix=eye(num_dim))
-    elif prior == "uniform":
+    elif prior_type == "uniform":
         prior = Independent(Uniform(-1.0 * ones(num_dim), 1.0 * ones(num_dim)), 1)
     simulator = diagonal_linear_gaussian
 
@@ -131,6 +131,12 @@ def test_batched_sample_log_prob_with_different_x(
     samples = posterior.sample_batched((10,), x_o)
     batched_log_probs = posterior.log_prob_batched(samples, x_o)
 
+    # Test large max_sampling_batch_size to test capping warning.
+    with pytest.warns(UserWarning, match="Capping max_sampling_batch_size"):
+        posterior.sample_batched(
+            (10,), ones(3, num_dim), max_sampling_batch_size=40_000
+        )
+
     assert (
         samples.shape == (10, x_o_batch_dim, num_dim)
         if x_o_batch_dim > 0