Repository: sbi-dev/sbi

Issue:
torch NotImplementedError
First of all, thanks for writing this great package! I just updated to a new environment, and am running `sbi==0.15.1` with `torch==1.8.1`. Now when I run `posterior.sample` I get the following error: 

```
NotImplementedError                       Traceback (most recent call last)
<ipython-input-13-cfda2fd907d3> in <module>
----> 1 posterior_samples = posterior.sample((200,), x=sw.obs, sample_with_mcmc = False)

~/anaconda3/envs/starwave-env/lib/python3.8/site-packages/sbi/inference/posteriors/direct_posterior.py in sample(self, sample_shape, x, show_progress_bars, sample_with_mcmc, mcmc_method, mcmc_parameters, rejection_sampling_parameters)
    354         else:
    355             # Rejection sampling.
--> 356             samples, _ = utils.sample_posterior_within_prior(
    357                 self.net,
    358                 self._prior,

~/anaconda3/envs/starwave-env/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     25         def decorate_context(*args, **kwargs):
     26             with self.__class__():
---> 27                 return func(*args, **kwargs)
     28         return cast(F, decorate_context)
     29 

~/anaconda3/envs/starwave-env/lib/python3.8/site-packages/sbi/utils/sbiutils.py in sample_posterior_within_prior(posterior_nn, prior, x, num_samples, show_progress_bars, warn_acceptance, sample_for_correction_factor, max_sampling_batch_size)
    206             .cpu()  # Move to cpu to evaluate under prior.
    207         )
--> 208         are_within_prior = within_support(prior, candidates)
    209         samples = candidates[are_within_prior]
    210         accepted.append(samples)

~/anaconda3/envs/starwave-env/lib/python3.8/site-packages/sbi/utils/sbiutils.py in within_support(distribution, samples)
    490         Tensor of bools indicating whether each sample was within the support.
    491     """
--> 492     if hasattr(distribution, "support"):
    493         sample_check = distribution.support.check(samples)
    494 

~/anaconda3/envs/starwave-env/lib/python3.8/site-packages/torch/distributions/distribution.py in support(self)
    105         representing this distribution's support.
    106         """
--> 107         raise NotImplementedError
    108 
    109     @property

NotImplementedError: 
```

It looks like sbi is trying to run `if hasattr(distribution, "support"):` which is breaking torch? Am I doing something wrong on my end? Any advice would be appreciated! 


Discussion:
Hi @vedantchandra and thanks for reporting this. 

Yes, since the new version of `torch` we are having issues with the `support` attribute. The check we are currently performing seems to not work as expected, we will look into this. 

To help us fix this: what kind of prior are you using? 

Resolving patch:
diff --git a/sbi/utils/sbiutils.py b/sbi/utils/sbiutils.py
--- a/sbi/utils/sbiutils.py
+++ b/sbi/utils/sbiutils.py
@@ -478,29 +478,33 @@ def within_support(distribution: Any, samples: Tensor) -> Tensor:
     returns whether it is finite or not (this hanldes e.g. `NeuralPosterior`). Only
     checking whether the log-probabilty is not `-inf` will not work because, as of
     torch v1.8.0, a `torch.distribution` will throw an error at `log_prob()` when the
-    sample is out of the support (see #451). In `prepare_for_sbi()`, we set 
-    `validate_args=False`. Thish would take care of this, but requires running 
+    sample is out of the support (see #451). In `prepare_for_sbi()`, we set
+    `validate_args=False`. This would take care of this, but requires running
     `prepare_for_sbi()` and otherwise throws a cryptic error.
 
     Args:
-        distribution: Distribution under which to evaluate the `samples`.
+        distribution: Distribution under which to evaluate the `samples`, e.g., a
+            PyTorch distribution or NeuralPosterior.
         samples: Samples at which to evaluate.
 
     Returns:
         Tensor of bools indicating whether each sample was within the support.
     """
-    if hasattr(distribution, "support"):
+    # Try to check using the support property, use log prob method otherwise.
+    try:
         sample_check = distribution.support.check(samples)
 
-        # Before torch v1.7.0, `support.check()` returned bools for every element. From
-        # v1.8.0 on, it directly considers all dimensions of a sample. E.g., for a 
-        # single sample in 3D, v1.7.0 would return [[True, True, True]] and v1.8.0 
-        # would return [True].
+        # Before torch v1.7.0, `support.check()` returned bools for every element.
+        # From v1.8.0 on, it directly considers all dimensions of a sample. E.g.,
+        # for a single sample in 3D, v1.7.0 would return [[True, True, True]] and
+        # v1.8.0 would return [True].
         if sample_check.ndim > 1:
             return torch.all(distribution.support.check(samples), axis=1)
         else:
             return distribution.support.check(samples)
-    else:
+    # Falling back to log prob method of either the NeuralPosterior's net, or of a
+    # custom wrapper distribution's.
+    except (NotImplementedError, AttributeError):
         return torch.isfinite(distribution.log_prob(samples))
 
 
diff --git a/sbi/utils/user_input_checks.py b/sbi/utils/user_input_checks.py
--- a/sbi/utils/user_input_checks.py
+++ b/sbi/utils/user_input_checks.py
@@ -19,6 +19,7 @@
     PytorchReturnTypeWrapper,
     ScipyPytorchWrapper,
 )
+from sbi.utils.sbiutils import within_support
 
 
 def process_prior(prior) -> Tuple[Distribution, int, bool]:
@@ -350,10 +351,23 @@ def check_prior_batch_behavior(prior) -> None:
     ), "prior.log_prob must return as many log probs as samples."
 
 
+def check_prior_support(prior):
+    """Check whether prior allows to check for support.
+
+    This either uses the PyTorch support property, or the custom prior .logprob method
+    """
+
+    try:
+        within_support(prior, prior.sample())
+    except NotImplementedError:
+        raise NotImplementedError(
+            """The prior must implement the support property or allow to call
+            .log_prob() outside of support."""
+        )
+
+
 def process_simulator(
-    user_simulator: Callable,
-    prior,
-    is_numpy_simulator: bool,
+    user_simulator: Callable, prior, is_numpy_simulator: bool,
 ) -> Callable:
     """Return a simulator that meets the requirements for usage in sbi.
 
@@ -457,10 +471,7 @@ def process_x(x: Tensor, x_shape: torch.Size) -> Tensor:
     return x
 
 
-def prepare_for_sbi(
-    simulator: Callable,
-    prior,
-) -> Tuple[Callable, Distribution]:
+def prepare_for_sbi(simulator: Callable, prior,) -> Tuple[Callable, Distribution]:
     """Prepare simulator, prior and for usage in sbi.
 
     One of the goals is to allow you to use sbi with inputs computed in numpy.
@@ -502,6 +513,7 @@ def check_sbi_inputs(simulator: Callable, prior: Distribution) -> None:
         prior: prior (Distribution like)
         x_shape: Shape of single simulation output $x$.
     """
+    check_prior_support(prior)
     num_prior_samples = 1
     theta = prior.sample((num_prior_samples,))
     theta_batch_shape, *_ = theta.shape
diff --git a/sbi/utils/user_input_checks_utils.py b/sbi/utils/user_input_checks_utils.py
--- a/sbi/utils/user_input_checks_utils.py
+++ b/sbi/utils/user_input_checks_utils.py
@@ -69,7 +69,8 @@ def _set_mean_and_variance(self):
                 ** 2
             )
             warnings.warn(
-                "Prior is lacking variance attribute, estimating prior variance from samples...",
+                """Prior is lacking variance attribute, estimating prior variance from
+                samples...""",
                 UserWarning,
             )
 
@@ -153,6 +154,10 @@ def mean(self):
     def variance(self):
         return torch.as_tensor(self.prior.variance, dtype=self.return_type)
 
+    @property
+    def support(self):
+        return self.prior.support
+
 
 class MultipleIndependent(Distribution):
     """Wrap a sequence of PyTorch distributions into a joint PyTorch distribution.
@@ -313,6 +318,9 @@ def __init__(self, constraints: Sequence[Constraint], dims_per_constraint: list)
     def check(self, value: Tensor) -> Tensor:
         """Returns a byte tensor of ``sample_shape + batch_shape`` indicating
         whether each event in value satisfies its corresponding constraint."""
+
+        if value.ndim < 2:
+            value = value.unsqueeze(0)
         result = torch.zeros((value.shape[0], len(self.constraints)))
         dim_idx = 0
         # For each constraint, select the corresponding values according to its


Test patch:
diff --git a/tests/linearGaussian_snpe_test.py b/tests/linearGaussian_snpe_test.py
--- a/tests/linearGaussian_snpe_test.py
+++ b/tests/linearGaussian_snpe_test.py
@@ -28,17 +28,10 @@
 
 
 @pytest.mark.parametrize(
-    "num_dim, prior_str",
-    (
-        (2, "gaussian"),
-        (2, "uniform"),
-        (1, "gaussian"),
-    ),
+    "num_dim, prior_str", ((2, "gaussian"), (2, "uniform"), (1, "gaussian"),),
 )
 def test_c2st_snpe_on_linearGaussian(
-    num_dim: int,
-    prior_str: str,
-    set_seed,
+    num_dim: int, prior_str: str, set_seed,
 ):
     """Test whether SNPE C infers well a simple example with available ground truth.
 
@@ -48,6 +41,7 @@ def test_c2st_snpe_on_linearGaussian(
 
     x_o = zeros(1, num_dim)
     num_samples = 1000
+    num_simulations = 2500
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(num_dim)
@@ -71,12 +65,11 @@ def simulator(theta):
         return linear_gaussian(theta, likelihood_shift, likelihood_cov)
 
     simulator, prior = prepare_for_sbi(simulator, prior)
-    inference = SNPE_C(
-        prior,
-        show_progress_bars=False,
-    )
+    inference = SNPE_C(prior, show_progress_bars=False,)
 
-    theta, x = simulate_for_sbi(simulator, prior, 2000, simulation_batch_size=1000)
+    theta, x = simulate_for_sbi(
+        simulator, prior, num_simulations, simulation_batch_size=1000
+    )
     _ = inference.append_simulations(theta, x).train(training_batch_size=100)
     posterior = inference.build_posterior().set_default_x(x_o)
     samples = posterior.sample((num_samples,))
@@ -195,8 +188,7 @@ def simulator(theta):
         pytest.param(
             "snpe_b",
             marks=pytest.mark.xfail(
-                raises=NotImplementedError,
-                reason="""SNPE-B not implemented""",
+                raises=NotImplementedError, reason="""SNPE-B not implemented""",
             ),
         ),
         "snpe_c",
@@ -366,11 +358,7 @@ def simulator(theta):
     net = utils.posterior_nn("maf", hidden_features=20)
 
     simulator, prior = prepare_for_sbi(simulator, prior)
-    inference = SNPE_C(
-        prior,
-        density_estimator=net,
-        show_progress_bars=False,
-    )
+    inference = SNPE_C(prior, density_estimator=net, show_progress_bars=False,)
 
     # We need a pretty big dataset to properly model the bimodality.
     theta, x = simulate_for_sbi(simulator, prior, 10000)
@@ -396,16 +384,8 @@ def simulator(theta):
     density = gaussian_kde(cond_samples.numpy().T, bw_method="scott")
 
     X, Y = np.meshgrid(
-        np.linspace(
-            limits[0][0],
-            limits[0][1],
-            50,
-        ),
-        np.linspace(
-            limits[1][0],
-            limits[1][1],
-            50,
-        ),
+        np.linspace(limits[0][0], limits[0][1], 50,),
+        np.linspace(limits[1][0], limits[1][1], 50,),
     )
     positions = np.vstack([X.ravel(), Y.ravel()])
     sample_kde_grid = np.reshape(density(positions).T, X.shape)
@@ -446,10 +426,7 @@ def simulator(theta):
         return linear_gaussian(theta, likelihood_shift, likelihood_cov)
 
     simulator, prior = prepare_for_sbi(simulator, prior)
-    inference = SNPE_C(
-        prior,
-        show_progress_bars=False,
-    )
+    inference = SNPE_C(prior, show_progress_bars=False,)
     theta, x = simulate_for_sbi(
         simulator, prior, 1000, simulation_batch_size=10, num_workers=6
     )