Add .bind() method to potentials

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

@@ -317,7 +317,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))
 
@@ -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))
@@ -488,6 +488,28 @@ def set_x(self, x_o: Optional[Tensor]):
         for comp_potential in self.potential_fns:
             comp_potential.set_x(x_o)
 
+    def bind(
+        self,
+        x_o: Tensor,
+        x_is_iid: bool = True,
+        **kwargs,
+    ) -> "EnsemblePotential":
+        """Return a new EnsemblePotential with x_o bound to all component potentials."""
+        if x_o is not None:
+            x_o = process_x(x_o).to(self.device)
+
+        bound_potentials = [
+            p.bind(x_o, x_is_iid=x_is_iid, **kwargs) for p in self.potential_fns
+        ]
+
+        return EnsemblePotential(
+            potential_fns=bound_potentials,
+            weights=self._weights,
+            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/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/mcmc_posterior.py

@@ -2,7 +2,6 @@
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 import inspect
 import warnings
-from copy import deepcopy
 from functools import partial
 from math import ceil
 from typing import Any, Callable, Dict, Literal, Optional, Union
@@ -237,7 +236,7 @@ def log_prob(
         warn("The log-probability is unnormalized!", stacklevel=2)
 
         x = self._x_else_default_x(x)
-        self.potential_fn.set_x(x, x_is_iid=True)
+        self.potential_fn = self.potential_fn.bind(x, x_is_iid=True)
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
         return self.potential_fn(
@@ -291,7 +290,7 @@ def sample(
         """
 
         x = self._x_else_default_x(x)
-        self.potential_fn.set_x(x, x_is_iid=True)
+        self.potential_fn = self.potential_fn.bind(x, x_is_iid=True)
 
         # Replace arguments that were not passed with their default.
         method = self.method if method is None else method
@@ -451,7 +450,7 @@ def sample_batched(
         # in the order of the observations.
         x_ = x.repeat_interleave(num_chains, dim=0)
 
-        self.potential_fn.set_x(x_, x_is_iid=False)
+        self.potential_fn = self.potential_fn.bind(x_, x_is_iid=False)
         self.potential_ = self._prepare_potential(method)  # type: ignore
 
         # For each observation in the batch, we have num_chains independent chains.
@@ -662,25 +661,25 @@ def _get_initial_params_batched(
             Tensor: initial parameters, one for each chain
         """
 
-        potential_ = deepcopy(self.potential_fn)
+        potential_ = self.potential_fn
         initial_params = []
-        init_fn = self._build_mcmc_init_fn(
-            self.proposal,
-            potential_fn=potential_,
-            transform=self.theta_transform,
-            init_strategy=init_strategy,  # type: ignore
-            **kwargs,
-        )
         for xi in x:
-            # Build init function
-            potential_.set_x(xi)
+            # Build init function with bound potential for this specific xi
+            potential_bound = potential_.bind(xi)
+            init_fn = self._build_mcmc_init_fn(
+                self.proposal,
+                potential_fn=potential_bound,
+                transform=self.theta_transform,
+                init_strategy=init_strategy,  # type: ignore
+                **kwargs,
+            )
 
             # Parallelize inits for resampling or sir.
             if num_workers > 1 and (
                 init_strategy == "resample" or init_strategy == "sir"
             ):
 
-                def seeded_init_fn(seed):
+                def seeded_init_fn(seed, init_fn=init_fn):
                     torch.manual_seed(seed)
                     return init_fn()
 

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/vector_field_posterior.py

@@ -235,7 +235,7 @@ def sample(
         x = self._x_else_default_x(x)
         x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
         is_iid = x.shape[0] > 1
-        self.potential_fn.set_x(
+        self.potential_fn = self.potential_fn.bind(
             x,
             x_is_iid=is_iid,
             iid_method=iid_method or self.potential_fn.iid_method,
@@ -457,7 +457,9 @@ def log_prob(
         x = self._x_else_default_x(x)
         x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
         is_iid = x.shape[0] > 1
-        self.potential_fn.set_x(x, x_is_iid=is_iid, **(ode_kwargs or {}))
+        self.potential_fn = self.potential_fn.bind(
+            x, x_is_iid=is_iid, **(ode_kwargs or {})
+        )
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
         return self.potential_fn(
@@ -521,7 +523,7 @@ def sample_batched(
         condition_dim = len(self.vector_field_estimator.condition_shape)
         batch_shape = x.shape[:-condition_dim]
         batch_size = batch_shape.numel()
-        self.potential_fn.set_x(x)
+        self.potential_fn = self.potential_fn.bind(x)
 
         max_sampling_batch_size = (
             self.max_sampling_batch_size
@@ -661,7 +663,9 @@ def map(
 
         if self._map is None or force_update:
             # rebuild coarse flow fast for MAP optimization.
-            self.potential_fn.set_x(self.default_x, atol=1e-2, rtol=1e-3, exact=True)
+            self.potential_fn = self.potential_fn.bind(
+                self.default_x, atol=1e-2, rtol=1e-3, exact=True
+            )
             callable_potential_fn = CallableDifferentiablePotentialFunction(
                 self.potential_fn
             )

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,7 +49,7 @@ 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):
@@ -66,9 +66,23 @@ 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)`."
             )
 
+    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 return_x_o(self) -> Optional[Tensor]:
         """Return the observed data at which the potential is evaluated.
 
@@ -149,3 +163,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

@@ -100,6 +100,17 @@ def to(self, device: Union[str, torch.device]) -> "PosteriorBasedPotential":
         self.posterior_estimator.to(device)
         return self
 
+    def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential":
+        """Create new potential with x bound, without mutable state."""
+        bound = PosteriorBasedPotential(
+            posterior_estimator=self.posterior_estimator,
+            prior=self.prior,
+            x_o=None,
+            device=self.device,
+        )
+        bound.set_x(x_o, x_is_iid=x_is_iid)
+        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.

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/inference/potentials/vector_field_potential.py

@@ -61,6 +61,7 @@ def __init__(
         self.vector_field_estimator.eval()
         self.iid_method = iid_method
         self.iid_params = iid_params
+        self.neural_ode_backend = neural_ode_backend
 
         neural_ode_kwargs = neural_ode_kwargs or {}
         self.neural_ode = build_neural_ode(
@@ -128,6 +129,40 @@ def set_x(
         elif self._x_o is not None:
             self.flows = self.rebuild_flows_for_batch(**ode_kwargs)
 
+    def bind(
+        self,
+        x_o: Tensor,
+        x_is_iid: bool = False,
+        iid_method: Optional[str] = None,
+        iid_params: Optional[Dict[str, Any]] = None,
+        guidance_method: Optional[str] = None,
+        guidance_params: Optional[Dict[str, Any]] = None,
+        **ode_kwargs,
+    ) -> "VectorFieldBasedPotential":
+        """Create new potential with x bound, without mutable state."""
+        bound = VectorFieldBasedPotential(
+            vector_field_estimator=self.vector_field_estimator,
+            prior=self.prior,
+            x_o=None,
+            device=self.device,
+            iid_method=self.iid_method,
+            iid_params=self.iid_params,
+            neural_ode_backend=self.neural_ode_backend,
+        )
+        bound.neural_ode.params.update(self.neural_ode.params)
+        # Transfer learned neural ODE parameters (e.g., from training) to the bound
+        # potential. This preserves the trained flow model when conditioning on new x.
+        bound.set_x(
+            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,
+        )
+        return bound
+
     def __call__(
         self,
         theta: Tensor,