new ruff

sbi/inference/potentials/score_based_potential.py

@@ -10,6 +10,7 @@
 from zuko.transforms import FreeFormJacobianTransform
 
 from sbi.inference.potentials.base_potential import BasePotential
+from sbi.inference.potentials.score_fn_iid import FNPEScoreFn
 from sbi.neural_nets.estimators.score_estimator import ConditionalScoreEstimator
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
@@ -19,8 +20,6 @@
 from sbi.utils.sbiutils import mcmc_transform, within_support
 from sbi.utils.torchutils import ensure_theta_batched
 
-from sbi.inference.potentials.score_fn_iid import FNPEScoreFn
-
 
 def score_estimator_based_potential(
     score_estimator: ConditionalScoreEstimator,
@@ -250,4 +249,3 @@ def f(t, x):
         exact=exact,
     )
     return transform
-

sbi/inference/potentials/score_fn_iid.py

@@ -1,11 +1,11 @@
-from typing import Callable, Optional, Union
-from sbi.neural_nets.estimators.score_estimator import ConditionalScoreEstimator
-from torch.distributions import Distribution
-
 from abc import abstractmethod
+from typing import Callable, Optional
+
 import torch
 from torch import Tensor
+from torch.distributions import Distribution
 
+from sbi.neural_nets.estimators.score_estimator import ConditionalScoreEstimator
 from sbi.utils.torchutils import ensure_theta_batched
 
 
@@ -131,4 +131,4 @@ def __call__(
         # Accumulate
         score = (1 - N) * prior_score + base_score.sum(-2, keepdim=True)
 
-        return score
+        return score

sbi/inference/potentials/score_utils.py

@@ -1,5 +1,5 @@
-from torch.distributions import Distribution, Independent, Normal, Uniform
 from torch import Tensor
+from torch.distributions import Distribution, Independent, Normal
 
 # Automatic denoising -----------------------------------------------------
 
@@ -129,4 +129,4 @@ def marginalize_gaussian(p: Normal, m: Tensor, s: Tensor) -> Normal:
     var = (m * std0) ** 2 + s**2
     std = var**0.5
 
-    return Normal(mu, std)
+    return Normal(mu, std)

tests/linearGaussian_npse_test.py

@@ -157,12 +157,12 @@ def simulator(theta):
     check_c2st(samples, target_samples, alg="npse_different_dims_and_resume_training")
 
 
-@pytest.mark.xfail(
-    reason="iid_bridge not working.",
-    raises=AssertionError,
-    strict=True,
-    match="Score accumulation*",
-)
+# @pytest.mark.xfail(
+#     reason="iid_bridge not working.",
+#     raises=AssertionError,
+#     strict=True,
+#     match="Score accumulation*",
+# )
 @pytest.mark.parametrize("num_trials", [2, 10])
 def test_npse_iid_inference(num_trials):
     """Test whether NPSE infers well a simple example with available ground truth."""