fix map error handling and tests.

Author: janfb

sbi/inference/posteriors/base_posterior.py

@@ -52,9 +52,8 @@ def __init__(
                 stacklevel=2,
             )
 
-        if not isinstance(potential_fn, BasePotential) and not isinstance(
-            potential_fn, BasePotential
-        ):
+        # Wrap as `CallablePotentialWrapper` if `potential_fn` is a Callable.
+        if not isinstance(potential_fn, BasePotential):
             kwargs_of_callable = list(inspect.signature(potential_fn).parameters.keys())
             for key in ["theta", "x_o"]:
                 assert key in kwargs_of_callable, (

sbi/utils/sbiutils.py

@@ -945,13 +945,13 @@ def gradient_ascent(
             try:
                 optimize_inits.requires_grad_(False)  # type: ignore
                 gradient = potential_fn.gradient(optimize_inits)
-            except NotImplementedError:
+            except (NotImplementedError, AttributeError):
                 optimize_inits.requires_grad_(True)  # type: ignore
                 probs = potential_fn(optimize_inits).squeeze()
                 loss = probs.sum()
                 loss.backward()
                 gradient = optimize_inits.grad
-                assert gradient is Tensor, "Gradient must be a tensor."
+                assert isinstance(gradient, Tensor), "Gradient must be a tensor."
 
             # Update the parameters with gradient descent.
             # See https://discuss.pytorch.org/t/updatation-of-parameters-without-using-optimizer-step/34244/2

tests/linearGaussian_npse_test.py

@@ -206,7 +206,8 @@ def test_npse_iid_inference(num_trials):
 
 @pytest.mark.slow
 @pytest.mark.xfail(
-    raises=AssertionError, reason="MAP optimization via score not working accurately."
+    raises=NotImplementedError,
+    reason="MAP optimization via score not working accurately.",
 )
 def test_npse_map():
     num_dim = 2

tests/posterior_nn_test.py

@@ -32,7 +32,13 @@
     (
         0,
         1,
-        pytest.param(2, marks=pytest.mark.xfail(raises=AssertionError)),
+        pytest.param(
+            2,
+            marks=pytest.mark.xfail(
+                raises=AssertionError,
+                reason=".log_prob() supports only batch size 1 for x_o.",
+            ),
+        ),
     ),
 )
 def test_log_prob_with_different_x(snpe_method: type, x_o_batch_dim: bool):

tests/sbc_test.py

@@ -13,10 +13,7 @@
 from sbi.analysis import sbc_rank_plot
 from sbi.diagnostics import check_sbc, get_nltp, run_sbc
 from sbi.inference import NPSE, SNLE, SNPE
-from sbi.simulators import linear_gaussian
-from sbi.simulators.linear_gaussian import (
-    linear_gaussian,
-)
+from sbi.simulators.linear_gaussian import linear_gaussian
 from sbi.utils import BoxUniform, MultipleIndependent
 from tests.test_utils import PosteriorPotential, TractablePosterior