Repository: sbi-dev/sbi

Issue:
Segfault on PyTorch 1.10 when data contains NaNs
# Overview 

The title of this issue is my own interpretation of what is happening, after I experimented with a few versions of the package.

On multiple machines and clusters where `sbi` is installed together with `torch` version 1.10.0, all the training from `sbi` fails with a `Segmentation fault (core dumped)` message.

The same training (a snippet of a typical script is below) is successful when the version of pytorch is 1.9.0 or 1.9.0.post2 
I tried to investigate further what is happening during training (e.g. with a debugger, trying to step into the `.train()` method produces a segmentation fault):
*  I have tested that setting the `max_num_epochs=1` gives rise to the same error, so the problem is at the beginning of training
* The minimal example in `sbi/examples/minimal.py` works as expected. This indicates something is particularly defective in my datasets. So I checked the operations done on the data...
* The `inference.get_simulations()` command fails in the same manner with a segfault. Since that is what happens at the beginning of the `train` method, I guess there should be a clue of what is going on. My simulations have several `torch.tensor(float("inf"))` values (around 20-30%) which should be excluded from training: maybe the handling of `inf` is wrong.
* `handle_invalid_x` works as expected and catches all the `inf`
* `warn_if_zscoring_changes_data(x)` causes a Segmentation Fault. Is it supposed to be applied only to `valid x` values?
* This does not crash the code: 
```python
        is_valid_x, num_nans, num_infs = handle_invalid_x(x, exclude_invalid_x)
        warn_if_zscoring_changes_data(x[is_validₓ])
``` 
* `torch.unique(zx, dim=0).numel()` this crashes the code is `x` has `nan` or `inf`

## Snippet of code

```python
inference = SNPE(prior=prior_sbi)
inference = inference.append_simulations(theta, x)
density_estimator = inference.train(show_train_summary=True)
```

# Summary

`sbi` is installed following the instructions on the website (but also independently installed with an environment.yml file similar to the one in this repo). When the version of pytorch installed is 1.10.0, the training of an `PosteriorEstimator` object like `SNPE()` will result in segmentation fault.
The culprit seems to be `warn_if_zscoring_changes_data(x)` in the `get_simulations()` of an estimator object. This might have something to do with the breaking changes in `torch==1.10.0` https://github.com/pytorch/pytorch/releases/tag/v1.10.0 or it should be called on valid x values.

**Note**: currently the installation will result in getting `setuptools==60.2.0` which has a different way of using `LooseVersion` and it will results in this https://stackoverflow.com/questions/70520120/attributeerror-module-setuptools-distutils-has-no-attribute-version bug when calling `sbi.utils` 



Discussion:
Thanks a lot for the comprehensive summary! We'll look into this.
Thanks for reporting this! 

I've opened an issue upstream reporting the segfault (https://github.com/pytorch/pytorch/issues/71089). Filtering out NaNs before calling `torch.unique` sounds like a viable option as a workaround.
The issue got fixed in the meantime -- we should be able to close this once 1.10.2 is released

Resolving patch:
diff --git a/sbi/inference/base.py b/sbi/inference/base.py
--- a/sbi/inference/base.py
+++ b/sbi/inference/base.py
@@ -252,7 +252,7 @@ def get_simulations(
         # Check for NaNs in simulations.
         is_valid_x, num_nans, num_infs = handle_invalid_x(x, exclude_invalid_x)
         # Check for problematic z-scoring
-        warn_if_zscoring_changes_data(x)
+        warn_if_zscoring_changes_data(x[is_valid_x])
         if warn_on_invalid:
             warn_on_invalid_x(num_nans, num_infs, exclude_invalid_x)
             warn_on_invalid_x_for_snpec_leakage(


Test patch: