bind implementation still using set_x as fallback yet

sbi/inference/posteriors/base_posterior.py

@@ -93,7 +93,7 @@ def potential(
                 This can be helpful for e.g. sensitivity analysis, but increases memory
                 consumption.
         """
-        self.potential_fn.set_x(self._x_else_default_x(x))
+        self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x))
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
         return self.potential_fn(
@@ -294,7 +294,7 @@ def map(
             )
 
         if self._map is None or force_update:
-            self.potential_fn.set_x(self.default_x)
+            self.potential_fn = self.potential_fn.bind(self.default_x)
             self._map = self._calculate_map(
                 num_iter=num_iter,
                 num_to_optimize=num_to_optimize,

sbi/inference/posteriors/ensemble_posterior.py

@@ -398,7 +398,7 @@ def map(
             return log_weights, maps
 
         else:
-            self.potential_fn.set_x(self._x_else_default_x(x))
+            self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x))
 
             if init_method == "posterior":
                 inits = self.sample((num_init_samples,), self._x_else_default_x(x))

sbi/inference/posteriors/importance_posterior.py

@@ -128,7 +128,7 @@ def log_prob(
             `len($\theta$)`-shaped log-probability.
         """
         x = self._x_else_default_x(x)
-        self.potential_fn.set_x(x)
+        self.potential_fn = self.potential_fn.bind(x)
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
 
@@ -212,7 +212,7 @@ def sample(
 
         method = self.method if method is None else method
 
-        self.potential_fn.set_x(self._x_else_default_x(x))
+        self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x))
 
         if method == "sir":
             return self._sir_sample(

sbi/inference/posteriors/rejection_posterior.py

@@ -120,7 +120,7 @@ def log_prob(
         )
         warn("The log-probability is unnormalized!", stacklevel=2)
 
-        self.potential_fn.set_x(self._x_else_default_x(x))
+        self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x))
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
         return self.potential_fn(
@@ -174,7 +174,7 @@ def sample(
             Samples from posterior.
         """
         num_samples = torch.Size(sample_shape).numel()
-        self.potential_fn.set_x(self._x_else_default_x(x))
+        self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x))
 
         potential = partial(self.potential_fn, track_gradients=True)
 

sbi/inference/posteriors/vi_posterior.py

@@ -1229,7 +1229,7 @@ def _compute_amortized_elbo_loss(self, x_batch: Tensor, n_particles: int) -> Ten
         x_expanded = x_batch.repeat(n_particles, 1)
 
         # Set x_o for batched evaluation (x_is_iid=False: each θ paired with its x)
-        self.potential_fn.set_x(x_expanded, x_is_iid=False)
+        self.potential_fn = self.potential_fn.bind(x_expanded, x_is_iid=False)
         log_potential_flat = self.potential_fn(theta_flat)
 
         # Reshape: (n_particles * batch_size,) -> (n_particles, batch_size)

sbi/inference/potentials/base_potential.py

@@ -49,11 +49,17 @@ def x_is_iid(self) -> bool:
             return self._x_is_iid
         else:
             raise ValueError(
-                "No observed data is available. Use `potential_fn.set_x(x_o)`."
+                "No observed data is available. Use `potential_fn.bind(x_o)`."
             )
 
     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."""
+        # import warnings
+        # warnings.warn(
+        #     "set_x() is deprecated, use .bind(x) instead for stateless binding",
+        #     FutureWarning,
+        #     stacklevel=2,
+        # )
         if x_o is not None:
             x_o = process_x(x_o).to(self.device)
         self._x_o = x_o
@@ -66,7 +72,7 @@ def x_o(self) -> Tensor:
             return self._x_o
         else:
             raise ValueError(
-                "No observed data is available. Use `potential_fn.set_x(x_o)`."
+                "No observed data is available. Use `potential_fn.bind(x_o)`."
             )
 
     @x_o.setter
@@ -82,6 +88,20 @@ def return_x_o(self) -> Optional[Tensor]:
         """
         return self._x_o
 
+    def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "BasePotential":
+        """Create new potential with x bound, without mutable state.
+
+        Args:
+            x_o: Observed data to bind.
+            x_is_iid: Whether x represents iid observations.
+
+        Returns:
+            New potential instance with x bound.
+
+        Subclasses must implement this method.
+        """
+        raise NotImplementedError(f"{self.__class__.__name__} must implement bind()")
+
     def to(self, device: Union[str, torch.device]) -> "BasePotential":
         """Move prior and x_o to the given device.
 
@@ -154,3 +174,12 @@ def __call__(self, theta, track_gradients: bool = True):
         """
         with torch.set_grad_enabled(track_gradients):
             return self.potential_fn(theta, self.x_o)
+
+    def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "CustomPotentialWrapper":
+        """Create new potential with x bound, without mutable state."""
+        return CustomPotentialWrapper(
+            potential_fn=self.potential_fn,
+            prior=self.prior,
+            x_o=x_o,
+            device=self.device,
+        )

sbi/inference/potentials/likelihood_based_potential.py

@@ -94,6 +94,17 @@ def to(self, device: Union[str, torch.device]) -> "LikelihoodBasedPotential":
         self.likelihood_estimator.to(device)
         return self
 
+    def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "LikelihoodBasedPotential":
+        """Create new potential with x bound, without mutable state."""
+        bound = LikelihoodBasedPotential(
+            likelihood_estimator=self.likelihood_estimator,
+            prior=self.prior,
+            x_o=None,
+            device=self.device,
+        )
+        bound.set_x(x_o, x_is_iid=x_is_iid)
+        return bound
+
     def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
         r"""Returns the potential $\log(p(x_o|\theta)p(\theta))$.
 

sbi/inference/potentials/posterior_based_potential.py

@@ -106,6 +106,15 @@ def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = False):
         """
         super().set_x(x_o, x_is_iid=x_is_iid)
 
+    def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential":
+        """Create new potential with x bound, without mutable state."""
+        return PosteriorBasedPotential(
+            posterior_estimator=self.posterior_estimator,
+            prior=self.prior,
+            x_o=x_o,
+            device=self.device,
+        )
+
     def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
         r"""Returns the potential for posterior-based methods.
 

sbi/inference/potentials/ratio_based_potential.py

@@ -82,6 +82,17 @@ def to(self, device: Union[str, torch.device]) -> "RatioBasedPotential":
         self.ratio_estimator.to(device)
         return self
 
+    def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "RatioBasedPotential":
+        """Create new potential with x bound, without mutable state."""
+        bound = RatioBasedPotential(
+            ratio_estimator=self.ratio_estimator,
+            prior=self.prior,
+            x_o=None,
+            device=self.device,
+        )
+        bound.set_x(x_o, x_is_iid=x_is_iid)
+        return bound
+
     def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
         r"""Returns the potential for likelihood-ratio-based methods.
 

sbi/utils/conditional_density_utils.py

@@ -412,15 +412,23 @@ def x_is_iid(self) -> bool:
             return self._x_is_iid
         else:
             raise ValueError(
-                "No observed data is available. Use `potential_fn.set_x(x_o)`."
+                "No observed data is available. Use `potential_fn.bind(x_o)`."
             )
 
     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.set_x(x_o, x_is_iid=x_is_iid)
+        self.potential_fn = self.potential_fn.bind(x_o, x_is_iid=x_is_iid)
+
+    def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "ConditionedPotential":
+        """Create new potential with x bound, without mutable state."""
+        return ConditionedPotential(
+            potential_fn=self.potential_fn.bind(x_o, x_is_iid),
+            condition=self.condition,
+            dims_to_sample=self.dims_to_sample,
+        )
 
     @property
     def x_o(self) -> Tensor: