Test set x deprecation

sbi/inference/posteriors/ensemble_posterior.py

@@ -1,6 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+import warnings
 from typing import List, Optional, Tuple, Union
 
 import torch
@@ -479,14 +480,21 @@ def allow_iid_x(self) -> bool:
         )
 
     def set_x(self, x_o: Optional[Tensor]):
-        """Check the shape of the observed data and, if valid, set it."""
-        if x_o is not None:
-            x_o = process_x(x_o).to(  # type: ignore
-                self.device
-            )
-        self._x_o = x_o
-        for comp_potential in self.potential_fns:
-            comp_potential.set_x(x_o)
+        """Check the shape of the observed data and, if valid, set it.
+
+        DEPRECATED: Use bind() instead. This method delegates to bind() internally.
+        It will be removed in a future release.
+        """
+
+        warnings.warn(
+            "set_x() is deprecated and will be removed in a future release. "
+            "Use bind() instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        bound = self.bind(x_o)
+        self._x_o = bound._x_o
+        self.potential_fns = bound.potential_fns
 
     def bind(
         self,

sbi/inference/potentials/base_potential.py

@@ -56,11 +56,20 @@ def x_is_iid(self) -> bool:
             )
 
     def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True):
-        """Check the shape of the observed data and, if valid, set it."""
-        if x_o is not None:
-            x_o = process_x(x_o).to(self.device)
-        self._x_o = x_o
-        self._x_is_iid = x_is_iid
+        """Check the shape of the observed data and, if valid, set it.
+
+        DEPRECATED: Use bind() instead. This method delegates to bind() internally.
+        """
+        import warnings
+
+        warnings.warn(
+            "set_x() is deprecated, use bind() instead",
+            FutureWarning,
+            stacklevel=2,
+        )
+        bound = self.bind(x_o, x_is_iid=x_is_iid)
+        self._x_o = bound._x_o
+        self._x_is_iid = bound._x_is_iid
 
     @property
     def x_o(self) -> Tensor:

sbi/inference/potentials/posterior_based_potential.py

@@ -115,13 +115,22 @@ def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential"
         return bound
 
     def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = False):
+        """Check the shape of the observed data and, if valid, set it.
+
+        DEPRECATED: Use bind() instead. This method delegates to bind() internally.
+        It will be removed in a future release.
         """
-        Check the shape of the observed data and, if valid, set it.
-        """
-        if x_o is not None:
-            x_o = process_x(x_o).to(self.device)
-        self._x_o = x_o
-        self._x_is_iid = x_is_iid
+        import warnings
+
+        warnings.warn(
+            "set_x() is deprecated and will be removed in a future release. "
+            "Use bind() instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        bound = self.bind(x_o, x_is_iid=x_is_iid)
+        self._x_o = bound._x_o
+        self._x_is_iid = bound._x_is_iid
 
     def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
         r"""Returns the potential for posterior-based methods.

sbi/inference/potentials/vector_field_potential.py

@@ -1,6 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+import warnings
 from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 import torch
@@ -113,6 +114,8 @@ def set_x(
 
         Rebuilds the continuous normalizing flow if the observed data is set.
 
+        DEPRECATED: Use bind() instead. This method delegates to bind() internally.
+        It will be removed in a future release.
         Args:
             x_o: The observed data.
             x_is_iid: Whether the observed data is IID (if batch_dim>1).
@@ -122,18 +125,32 @@ def set_x(
                 `IIDScoreFunction`.
             ode_kwargs: Additional keyword arguments for the neural ODE.
         """
-        if x_o is not None:
-            x_o = process_x(x_o).to(self.device)
-        self._x_o = x_o
-        self._x_is_iid = x_is_iid
-        self.iid_method = iid_method or self.iid_method
-        self.iid_params = iid_params
-        self.guidance_method = guidance_method
-        self.guidance_params = guidance_params
-        if not x_is_iid and (self._x_o is not None):
-            self.flow = self.rebuild_flow(**ode_kwargs)
-        elif self._x_o is not None:
-            self.flows = self.rebuild_flows_for_batch(**ode_kwargs)
+
+        warnings.warn(
+            "set_x() is deprecated and will be removed in a future release. "
+            "Use bind() instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        bound = self.bind(
+            x_o,
+            x_is_iid=x_is_iid,
+            iid_method=iid_method,
+            iid_params=iid_params,
+            guidance_method=guidance_method,
+            guidance_params=guidance_params,
+            **ode_kwargs,
+        )
+        self._x_o = bound._x_o
+        self._x_is_iid = bound._x_is_iid
+        self.iid_method = bound.iid_method
+        self.iid_params = bound.iid_params
+        self.guidance_method = bound.guidance_method
+        self.guidance_params = bound.guidance_params
+        if not x_is_iid and (bound._x_o is not None):
+            self.flow = bound.flow
+        elif bound._x_o is not None:
+            self.flows = bound.flows
 
     def bind(
         self,

sbi/utils/conditional_density_utils.py

@@ -15,7 +15,6 @@
     MixtureDensityEstimator,
 )
 from sbi.utils.torchutils import ensure_theta_batched
-from sbi.utils.user_input_checks import process_x
 
 
 def compute_corrcoeff(probs: Tensor, limits: Tensor):
@@ -418,11 +417,20 @@ def x_is_iid(self) -> bool:
             )
 
     def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True):
-        """Check the shape of the observed data and, if valid, set it."""
-        if x_o is not None:
-            x_o = process_x(x_o).to(self.device)
-        self._x_is_iid = x_is_iid
-        self.potential_fn = self.potential_fn.bind(x_o, x_is_iid=x_is_iid)
+        """Check the shape of the observed data and, if valid, set it.
+
+        DEPRECATED: Use bind() instead. This method delegates to bind() internally.
+        """
+        import warnings
+
+        warnings.warn(
+            "set_x() is deprecated, use bind() instead",
+            FutureWarning,
+            stacklevel=2,
+        )
+        bound = self.bind(x_o, x_is_iid=x_is_iid)
+        self._x_is_iid = bound._x_is_iid
+        self.potential_fn = bound.potential_fn
 
     def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "ConditionedPotential":
         """Create new potential with x bound, without mutable state."""

tests/inference_on_device_test.py

@@ -51,7 +51,7 @@
 from sbi.utils import BoxUniform
 from sbi.utils.sbiutils import seed_all_backends
 from sbi.utils.torchutils import gpu_available, process_device
-from sbi.utils.user_input_checks import validate_theta_and_x
+from sbi.utils.user_input_checks import process_x, validate_theta_and_x
 
 pytestmark = pytest.mark.skipif(
     not gpu_available(), reason="No CUDA or MPS device available."
@@ -399,7 +399,9 @@ def allow_iid_x(self) -> bool:
         def bind(self, x_o: torch.Tensor, x_is_iid: bool = True) -> "FakePotential":
             """Create new potential with x bound, without mutable state."""
             bound = FakePotential(prior=self.prior, device=self.device)
-            bound.set_x(x_o, x_is_iid=x_is_iid)
+            x_o = process_x(x_o).to(self.device)
+            bound._x_o = x_o
+            bound._x_is_iid = x_is_iid
             return bound
 
     potential_fn = FakePotential(
@@ -459,7 +461,9 @@ def allow_iid_x(self) -> bool:
         def bind(self, x_o: torch.Tensor, x_is_iid: bool = True) -> "FakePotential":
             """Create new potential with x bound, without mutable state."""
             bound = FakePotential(prior=self.prior, device=self.device)
-            bound.set_x(x_o, x_is_iid=x_is_iid)
+            x_o = process_x(x_o).to(self.device)
+            bound._x_o = x_o
+            bound._x_is_iid = x_is_iid
             return bound
 
     potential_fn = FakePotential(prior=prior, device=device)

tests/vi_test.py

@@ -34,6 +34,7 @@
 )
 from sbi.utils import MultipleIndependent
 from sbi.utils.metrics import c2st, check_c2st
+from sbi.utils.user_input_checks import process_x
 
 # Supported variational families for VI
 FLOWS = ["maf", "nsf", "naf", "unaf", "nice", "sospf", "gaussian", "gaussian_diag"]
@@ -60,7 +61,9 @@ def allow_iid_x(self) -> bool:
     def bind(self, x_o: torch.Tensor, x_is_iid: bool = True) -> "FakePotential":
         """Create new potential with x bound, without mutable state."""
         bound = FakePotential(prior=self.prior, device=self.device)
-        bound.set_x(x_o, x_is_iid=x_is_iid)
+        x_o = process_x(x_o).to(self.device)
+        bound._x_o = x_o
+        bound._x_is_iid = x_is_iid
         return bound
 
 
@@ -81,7 +84,9 @@ def bind(
         ) -> "TractablePotential":
             """Create new potential with x bound, without mutable state."""
             bound = TractablePotential(prior=self.prior, device=self.device)
-            bound.set_x(x_o, x_is_iid=x_is_iid)
+            x_o = process_x(x_o).to(self.device)
+            bound._x_o = x_o
+            bound._x_is_iid = x_is_iid
             return bound
 
     return TractablePotential(prior=prior)