move NaN check to posterior.sample level, update tests, fix rejection sampling warning

Author: janfb

sbi/inference/posteriors/vector_field_posterior.py

@@ -340,6 +340,12 @@ def _sample_via_diffusion(
         # Concatenate all batches and ensure we return exactly the requested number
         samples = torch.cat(all_samples, dim=0)[:total_samples_needed]
 
+        if torch.isnan(samples).all():
+            raise RuntimeError(
+                "All samples NaN after diffusion sampling. "
+                "This may indicate numerical instability in the vector field."
+            )
+
         return samples
 
     def sample_via_ode(

sbi/samplers/rejection/rejection.py

@@ -327,7 +327,7 @@ def accept_reject_sample(
             max(int(1.5 * num_remaining / max(min_acceptance_rate, 1e-12)), 100),
         )
         if (
-            num_samples_possible > 1000
+            num_samples_possible > (sampling_batch_size - 1)
             and min_acceptance_rate < warn_acceptance
             and not leakage_warning_raised
         ):

sbi/samplers/score/diffuser.py

@@ -174,13 +174,6 @@ def run(
             if save_intermediate:
                 intermediate_samples.append(samples)
 
-        # Check for NaN values after predictor
-        if torch.isnan(samples).any():
-            raise RuntimeError(
-                "NaN values detected after diffusion sampling "
-                "This may indicate numerical instability in the vector field."
-            )
-
         if save_intermediate:
             return torch.cat(intermediate_samples, dim=0)
         else:

tests/linearGaussian_vector_field_test.py

@@ -373,7 +373,11 @@ def test_vector_field_iid_inference(
         num_trials: The number of trials to run.
     """
 
-    if prior_type == "uniform" and iid_method in ["gauss", "auto_gauss", "jac_gauss"]:
+    if (
+        vector_field_type == "fmpe"
+        and prior_type == "uniform"
+        and iid_method in ["gauss", "auto_gauss", "jac_gauss"]
+    ):
         # TODO: Predictor produces NaNs for these cases, see #1656
         pytest.skip("Known issue of IID methods with uniform priors, see #1656.")