refactor: remove deprecated x_shape where not needed. (#1271)

sbi-dev · Sep 11, 2024 · 299854e · 299854e
1 parent dbcd06b
commit 299854e
Show file tree

Hide file tree

Showing 5 changed files with 1 addition and 10 deletions.
diff --git a/sbi/inference/posteriors/ensemble_posterior.py b/sbi/inference/posteriors/ensemble_posterior.py
@@ -96,7 +96,6 @@ def __init__(
             potential_fn=potential_fn,
             theta_transform=theta_transform,
             device=device,
-            x_shape=None,
         )
 
     def ensure_same_device(self, posteriors: List) -> str:

diff --git a/sbi/inference/trainers/fmpe/fmpe.py b/sbi/inference/trainers/fmpe/fmpe.py
@@ -23,7 +23,6 @@
     npe_msg_on_invalid_x,
     validate_theta_and_x,
     warn_if_zscoring_changes_data,
-    x_shape_from_simulation,
 )
 from sbi.utils.sbiutils import mask_sims_from_prior
 
@@ -199,7 +198,6 @@ def train(
                 theta[self.train_indices].to("cpu"),
                 x[self.train_indices].to("cpu"),
             )
-            self._x_shape = x_shape_from_simulation(x.to("cpu"))
 
             del theta, x
 

diff --git a/sbi/inference/trainers/npse/npse.py b/sbi/inference/trainers/npse/npse.py
@@ -26,7 +26,6 @@
     test_posterior_net_for_multi_d_x,
     validate_theta_and_x,
     warn_if_zscoring_changes_data,
-    x_shape_from_simulation,
 )
 from sbi.utils.sbiutils import ImproperEmpirical, mask_sims_from_prior
 
@@ -282,7 +281,6 @@ def default_calibration_kernel(x):
                 theta[self.train_indices].to("cpu"),
                 x[self.train_indices].to("cpu"),
             )
-            self._x_shape = x_shape_from_simulation(x.to("cpu"))
 
             test_posterior_net_for_multi_d_x(
                 self._neural_net,

diff --git a/sbi/inference/trainers/nre/nre_base.py b/sbi/inference/trainers/nre/nre_base.py
@@ -21,7 +21,6 @@
     check_estimator_arg,
     check_prior,
     clamp_and_warn,
-    x_shape_from_simulation,
 )
 from sbi.utils.torchutils import repeat_rows
 
@@ -203,7 +202,6 @@ def train(
                 theta[self.train_indices].to("cpu"),
                 x[self.train_indices].to("cpu"),
             )
-            self._x_shape = x_shape_from_simulation(x.to("cpu"))
             del x, theta
         self._neural_net.to(self._device)
 

diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -219,7 +219,6 @@ def __init__(
         potential_fn: Callable,
         theta_transform: Optional[TorchTransform] = None,
         device: Optional[str] = "cpu",
-        x_shape: Optional[torch.Size] = None,
     ):
         """
         Args:
@@ -228,10 +227,9 @@ def __init__(
                 Allows to perform, e.g. MCMC in unconstrained space.
             device: Training device, e.g., "cpu", "cuda" or "cuda:0". If None,
                 `potential_fn.device` is used.
-            x_shape: Shape of the observed data.
         """
         assert isinstance(potential_fn, PosteriorPotential)
-        super().__init__(potential_fn, theta_transform, device, x_shape)
+        super().__init__(potential_fn, theta_transform, device)
 
     def sample(
         self,