Repository: sbi-dev/sbi

Issue:
Mask and Tensor Dimension Mismatch, Inconsistent Error
Hello. I've been toying around with SBI on a toy model of a two-body gravitation simulation. I've got a pretty straightforward simulator, but when I try to run inference on it, I get a torch error that my tensor and mask have mismatched dimensions. The _really_ weird thing is that depending on how many parameters I put in to the simulator wrapper, I seem to get the error for different reasons. Attached is the MWE with the simulator class. When I run this script as-is, I get the following error message:

> Running 4 simulations in 4 batches.: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 308.52it/s]
Traceback (most recent call last):
  File ".\sbi_gravity_mwe.py", line 72, in <module>
    posterior = infer(sim_wrapper, prior, method='SNRE', num_simulations=4, num_workers=4)
  File "C:\Users\Rainier\AppData\Local\Programs\Python\Python38\lib\site-packages\sbi\inference\base.py", line 72, in infer
    posterior = infer_(num_simulations=num_simulations)
  File "C:\Users\Rainier\AppData\Local\Programs\Python\Python38\lib\site-packages\sbi\inference\snre\snre_b.py", line 126, in __call__
    return super().__call__(**kwargs)
  File "C:\Users\Rainier\AppData\Local\Programs\Python\Python38\lib\site-packages\sbi\inference\snre\snre_base.py", line 149, in __call__
    theta, x, _ = self._get_from_data_bank(self._round, exclude_invalid_x, False)
  File "C:\Users\Rainier\AppData\Local\Programs\Python\Python38\lib\site-packages\sbi\inference\base.py", line 238, in _get_from_data_bank
    return theta[is_valid_x], x[is_valid_x], prior_masks[is_valid_x]
IndexError: The shape of the mask [4, 2] at index 1 does not match the shape of the indexed tensor [4, 200, 2] at index 1

(Running with more simulations just increases the first dimension, from e.g. 4 to 300 or what have you.)

However, if we un-do the comments on lines 59, 60, 66 and 67 (now giving in 6 parameters), the error is different:

> Running 4 simulations in 4 batches.: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 286.97it/s]
Traceback (most recent call last):
  File ".\sbi_gravity_mwe.py", line 72, in <module>
    posterior = infer(sim_wrapper, prior, method='SNRE', num_simulations=4, num_workers=4)
  File "C:\Users\Rainier\AppData\Local\Programs\Python\Python38\lib\site-packages\sbi\inference\base.py", line 72, in infer
    posterior = infer_(num_simulations=num_simulations)
  File "C:\Users\Rainier\AppData\Local\Programs\Python\Python38\lib\site-packages\sbi\inference\snre\snre_b.py", line 126, in __call__
    return super().__call__(**kwargs)
  File "C:\Users\Rainier\AppData\Local\Programs\Python\Python38\lib\site-packages\sbi\inference\snre\snre_base.py", line 149, in __call__
    theta, x, _ = self._get_from_data_bank(self._round, exclude_invalid_x, False)
  File "C:\Users\Rainier\AppData\Local\Programs\Python\Python38\lib\site-packages\sbi\inference\base.py", line 238, in _get_from_data_bank
    return theta[is_valid_x], x[is_valid_x], prior_masks[is_valid_x]
IndexError: The shape of the mask [4, 2] at index 1 does not match the shape of the indexed tensor [4, 6] at index 1

Notice that the indexed tensor shape has lost an entire dimension.(!) Am I missing some detail of how I am feeding my parameters to the model wrapper, or how the wrapper is set up? Or is something deeper going on?

This is running on the latest version of sbi, 12.2. Python version 3.8.4 on Windows 10 Home. 

[sbi_gravity_mwe.txt](https://github.com/mackelab/sbi/files/5235307/sbi_gravity_mwe.txt)




Discussion:
Hi Rainier, 
and thanks for pointing us to this issue. Before we dig deeper, what is the output dimension of your simulator? 
I might be missing something, but the [4, 200, 2] shape of the tensor suggests that it is 2D? 
Thanks for the MWE. It shows that your simulator output for a single parameter vector has shape `[200, 2]`. However,`SBI` does not support multi-dimensional `x` (yet), and this seems to lead to this inconsistent behavior. 

One possibility is to just flatten the output of the simulator, i.e., return shape `(400,)`. Then `sbi` will adapt the input layer shape of the density estimator accordingly and things should work without bugs. However, it could be challenging to use such a high dimensional input.  

Therefore, another possibility is to use another neural network to learn informative summary statistics. To this end, your simulator would still return the flattened version of the output, `(400,)`, but you pass a PyTorch `nn.Module` of your choice to `sbi`. 
E.g, the network could have several linear or convolutional layers and would reduce the dimensionality of the input from `(400,)` to `(20,)` or so. 
This _embedding_ net is prepended to the density estimator neural network and its parameters are trained jointly with the density estimator. [Here](https://github.com/mackelab/sbi/blob/main/tutorials/05_embedding_net.ipynb) is an example with an simulator that outputs images. 
Thanks @janfb, it appears to work with a smaller number of points from the trajectory and just calling `flatten` on the simulator output.

However, I'm still perplexed why this behavior seemed to be inconsistent before...
I was able to reproduce the error you described above for the first case. This is because `sbi` is not able to handle multidimensional input, e.g., the function that creates the mask of valid `x`s assumes 1D `x` leading to the additional dimension in the mask and to the error. 
I could not reproduce the second case you described because it seems the line numbers don't match. 

In any case, `sbi` should be able to also handle multidimensional input and I will create a separate issue for that. Thanks for pointing us to this problem! 

Resolving patch:
diff --git a/sbi/inference/posteriors/base_posterior.py b/sbi/inference/posteriors/base_posterior.py
--- a/sbi/inference/posteriors/base_posterior.py
+++ b/sbi/inference/posteriors/base_posterior.py
@@ -9,18 +9,14 @@
     Dict,
     List,
     Optional,
-    Sequence,
     Tuple,
-    TypeVar,
-    Union,
-    cast,
 )
 
 import numpy as np
 import torch
 from pyro.infer.mcmc import HMC, NUTS
 from pyro.infer.mcmc.api import MCMC
-from torch import Tensor, log
+from torch import Tensor
 from torch import multiprocessing as mp
 from torch import nn
 
diff --git a/sbi/user_input/user_input_checks.py b/sbi/user_input/user_input_checks.py
--- a/sbi/user_input/user_input_checks.py
+++ b/sbi/user_input/user_input_checks.py
@@ -257,14 +257,6 @@ def check_for_possibly_batched_x_shape(x_shape):
             respectively, without the user needing to wrap or unsqueeze them.
             """
         )
-    # Warn on multidimensional data with `batch_size` one.
-    elif inferred_data_ndim > 1 and inferred_batch_shape == 1:
-        warnings.warn(
-            f"""Beware: The observed data (x_o) you passed was interpreted to have
-            matrix shape: {inferred_data_shape}. The current implementation of sbi
-            might not provide stable support for this and result in shape mismatches.
-            """
-        )
     else:
         pass
 
diff --git a/sbi/utils/sbiutils.py b/sbi/utils/sbiutils.py
--- a/sbi/utils/sbiutils.py
+++ b/sbi/utils/sbiutils.py
@@ -4,15 +4,11 @@
 import logging
 from typing import (
     Any,
-    Callable,
     Dict,
     List,
-    Optional,
     Sequence,
     Tuple,
-    TypeVar,
     Union,
-    cast,
 )
 
 import torch
@@ -21,10 +17,21 @@
 from torch import nn as nn
 from torch import ones, zeros
 from tqdm.auto import tqdm
+import warnings
 
 
 def x_shape_from_simulation(batch_x: Tensor) -> torch.Size:
-    assert batch_x.ndim == 2, "Simulated data must be a batch with two dimensions."
+    ndims = batch_x.ndim
+    assert ndims >= 2, "Simulated data must be a batch with at least two dimensions."
+
+    # Warn in case of multi-dimensional x.
+    if ndims > 2:
+        warnings.warn(
+            f"""The simulated data x has {ndims-1} dimensions. With default settings, 
+            sbi cannot deal with multidimensional simulations. Make sure to use an
+            embedding net that reduces the dimensionality, e.g., a CNN in case of
+            images, or change the simulator to return one-dimensional x."""
+        )
     return batch_x[0].unsqueeze(0).shape
 
 
@@ -76,11 +83,11 @@ def standardizing_transform(
 class Standardize(nn.Module):
     def __init__(self, mean: Union[Tensor, float], std: Union[Tensor, float]):
         super(Standardize, self).__init__()
-        mean, std = map(torch.as_tensor, (mean, std))                
+        mean, std = map(torch.as_tensor, (mean, std))
         self.mean = mean
         self.std = std
-        self.register_buffer("_mean", mean)  
-        self.register_buffer("_std", std)              
+        self.register_buffer("_mean", mean)
+        self.register_buffer("_std", std)
 
     def forward(self, tensor):
         return (tensor - self.mean) / self.std
@@ -111,7 +118,8 @@ def standardizing_net(batch_t: Tensor, min_std: float = 1e-7) -> nn.Module:
             f"""Using a one-dimensional batch will instantiate a Standardize transform 
             with (mean, std) parameters which are not representative of the data. We allow
             this behavior because you might be loading a pre-trained. If this is not the case, 
-            please be sure to use a larger batch.""")    
+            please be sure to use a larger batch."""
+        )
 
     return Standardize(t_mean, t_std)
 
@@ -232,6 +240,9 @@ def handle_invalid_x(
 
     batch_size = x.shape[0]
 
+    # Squeeze to cover all dimensions in case of multidimensional x.
+    x = x.reshape(batch_size, -1)
+
     x_is_nan = torch.isnan(x).any(dim=1)
     x_is_inf = torch.isinf(x).any(dim=1)
     num_nans = int(x_is_nan.sum().item())


Test patch:
diff --git a/tests/multidimensional_x_test.py b/tests/multidimensional_x_test.py
new file mode 100644
--- /dev/null
+++ b/tests/multidimensional_x_test.py
@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from torch import zeros
+import torch
+import pytest
+from sbi import utils as utils
+from sbi.inference import SNPE, prepare_for_sbi
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# Minimal 2D simulator.
+def simulator_2d(theta):
+    x = torch.rand(theta.shape[0], 32, 32)
+    return x
+
+
+class CNNEmbedding(nn.Module):
+    """Big CNN embedding net to levarage GPU computation."""
+
+    def __init__(self):
+        super().__init__()
+        # 2D convolutional layer
+        self.conv1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, padding=2)
+        self.conv2 = nn.Conv2d(in_channels=16, out_channels=6, kernel_size=5, padding=2)
+        # Maxpool layer that reduces 32x32 image to 4x4
+        self.pool = nn.MaxPool2d(kernel_size=8, stride=8)
+        # Fully connected layer taking as input the 6 flattened output arrays from the
+        # maxpooling layer
+        self.fc = nn.Linear(in_features=6 * 4 * 4, out_features=8)
+
+    def forward(self, x):
+        x = x.view(-1, 1, 32, 32)
+        x = F.relu(self.conv1(x))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 6 * 4 * 4)
+        x = F.relu(self.fc(x))
+        return x
+
+
+@pytest.mark.parametrize(
+    "embedding",
+    (
+        pytest.param(nn.Identity, marks=pytest.mark.xfail(reason="Invalid embedding.")),
+        pytest.param(CNNEmbedding),
+    ),
+)
+def test_inference_with_2d_x(embedding):
+
+    num_dim = 2
+    num_samples = 10
+    num_simulations = 100
+
+    prior = utils.BoxUniform(zeros(num_dim), torch.ones(num_dim))
+
+    simulator, prior = prepare_for_sbi(simulator_2d, prior)
+
+    theta_o = torch.ones(1, num_dim)
+    x_o = simulator(theta_o)
+
+    infer = SNPE(
+        simulator,
+        prior,
+        show_progress_bars=False,
+        density_estimator=utils.posterior_nn(model="mdn", embedding_net=embedding(),),
+    )
+
+    posterior = infer(
+        num_simulations=num_simulations, training_batch_size=100, max_num_epochs=10
+    ).set_default_x(x_o)
+
+    posterior.log_prob(posterior.sample((num_samples,), show_progress_bars=False))