From 89bc9a2a61eec0708b4210a18d1244929881599a Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 07:26:34 +0200 Subject: [PATCH 01/14] Add .bind() method to potentials and fix _base_recursor scoping - Add bind() method to BasePotential (abstract) and CustomPotentialWrapper - Add bind() to LikelihoodBasedPotential, PosteriorBasedPotential, RatioBasedPotential - Add bind() to VectorFieldBasedPotential with full parameter support - Add bind() to EnsemblePotential in ensemble_posterior.py - Add bind() to ConditionedPotential in conditional_density_utils.py - Fix _base_recursor to use local variable _active_holder to avoid Python scoping issues with the _active parameter - Update posteriors to use bind() instead of set_x() for immutability - Fix PytorchReturnTypeWrapper, MultipleIndependent, OneDimPriorWrapper init order and add __deepcopy__ methods for proper deepcopy support - Update error messages to reference bind() instead of set_x() - Add bind() to test FakePotential and TractablePotential classes --- sbi/inference/posteriors/base_posterior.py | 4 +- .../posteriors/ensemble_posterior.py | 26 +++++++++- .../posteriors/importance_posterior.py | 4 +- sbi/inference/posteriors/mcmc_posterior.py | 29 ++++++------ .../posteriors/rejection_posterior.py | 4 +- sbi/inference/posteriors/vi_posterior.py | 2 +- sbi/inference/potentials/base_potential.py | 27 ++++++++++- .../potentials/likelihood_based_potential.py | 11 +++++ .../potentials/posterior_based_potential.py | 11 +++++ .../potentials/ratio_based_potential.py | 11 +++++ .../potentials/vector_field_potential.py | 32 +++++++++++++ sbi/utils/conditional_density_utils.py | 12 ++++- sbi/utils/torchutils.py | 17 +++---- sbi/utils/user_input_checks_utils.py | 47 ++++++++++++++----- tests/inference_on_device_test.py | 12 +++++ tests/vi_test.py | 14 ++++++ 16 files changed, 215 insertions(+), 48 deletions(-) diff --git a/sbi/inference/posteriors/base_posterior.py b/sbi/inference/posteriors/base_posterior.py index 599622904..f5111cecb 100644 --- a/sbi/inference/posteriors/base_posterior.py +++ b/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, diff --git a/sbi/inference/posteriors/ensemble_posterior.py b/sbi/inference/posteriors/ensemble_posterior.py index 755f1cebe..fa4015f92 100644 --- a/sbi/inference/posteriors/ensemble_posterior.py +++ b/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. diff --git a/sbi/inference/posteriors/importance_posterior.py b/sbi/inference/posteriors/importance_posterior.py index cfbc3d2f8..fdc98a9a1 100644 --- a/sbi/inference/posteriors/importance_posterior.py +++ b/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( diff --git a/sbi/inference/posteriors/mcmc_posterior.py b/sbi/inference/posteriors/mcmc_posterior.py index f90db67d6..fc1d254cf 100644 --- a/sbi/inference/posteriors/mcmc_posterior.py +++ b/sbi/inference/posteriors/mcmc_posterior.py @@ -2,7 +2,6 @@ # under the Apache License Version 2.0, see 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() diff --git a/sbi/inference/posteriors/rejection_posterior.py b/sbi/inference/posteriors/rejection_posterior.py index 7b05160a2..92e173129 100644 --- a/sbi/inference/posteriors/rejection_posterior.py +++ b/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) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index e070d904f..d10162722 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/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) diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index d24b05b94..f75f3e050 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/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,7 +66,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 @@ -74,6 +74,20 @@ def x_o(self, x_o: Optional[Tensor]) -> None: """Check the shape of the observed data and, if valid, set it.""" self.set_x(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. @@ -154,3 +168,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, + ) diff --git a/sbi/inference/potentials/likelihood_based_potential.py b/sbi/inference/potentials/likelihood_based_potential.py index d2907430c..991254489 100644 --- a/sbi/inference/potentials/likelihood_based_potential.py +++ b/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))$. diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py index eda384626..91a4af0e3 100644 --- a/sbi/inference/potentials/posterior_based_potential.py +++ b/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. diff --git a/sbi/inference/potentials/ratio_based_potential.py b/sbi/inference/potentials/ratio_based_potential.py index 844405e75..917fede47 100644 --- a/sbi/inference/potentials/ratio_based_potential.py +++ b/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. diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index b7ac4ba45..fbac610fd 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/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,37 @@ 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.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, diff --git a/sbi/utils/conditional_density_utils.py b/sbi/utils/conditional_density_utils.py index 0d48efedb..a42e6b4ed 100644 --- a/sbi/utils/conditional_density_utils.py +++ b/sbi/utils/conditional_density_utils.py @@ -412,7 +412,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): @@ -420,7 +420,15 @@ def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): 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: diff --git a/sbi/utils/torchutils.py b/sbi/utils/torchutils.py index 4405bed65..ce657cf34 100644 --- a/sbi/utils/torchutils.py +++ b/sbi/utils/torchutils.py @@ -677,13 +677,14 @@ def _base_recursor( _active: Object identities on the active recursion path. Used internally to avoid infinite recursion in cyclic object graphs. """ - if _active is None: - _active = set() + _active_holder: Optional[set[int]] = _active + if _active_holder is None: + _active_holder = set() obj_id = id(obj) - if obj_id in _active: + if obj_id in _active_holder: return - _active.add(obj_id) + _active_holder.add(obj_id) if isinstance(obj, Module) and check(obj): action(obj) @@ -698,7 +699,7 @@ def _base_recursor( key=k, check=check, action=action, - _active=_active, + _active=_active_holder, ) elif isinstance(obj, type): # Skip class/type objects to avoid modifying immutable C extension types @@ -715,7 +716,7 @@ def _base_recursor( key=k, check=check, action=action, - _active=_active, + _active=_active_holder, ) elif isinstance(obj, (List, Tuple, Generator)): new_obj = [] @@ -723,12 +724,12 @@ def _base_recursor( if check(o): new_obj.append(action(o)) else: - _base_recursor(o, check=check, action=action, _active=_active) + _base_recursor(o, check=check, action=action, _active=_active_holder) new_obj.append(o) if parent is not None and key is not None: setattr(parent, key, type(obj)(new_obj)) # type: ignore - _active.remove(obj_id) + _active_holder.remove(obj_id) def move_all_tensor_to_device(obj: object, device: Union[str, torch.device]) -> None: diff --git a/sbi/utils/user_input_checks_utils.py b/sbi/utils/user_input_checks_utils.py index dbd8a1d5e..cbb0f66e8 100644 --- a/sbi/utils/user_input_checks_utils.py +++ b/sbi/utils/user_input_checks_utils.py @@ -2,6 +2,7 @@ # under the Apache License Version 2.0, see import warnings +from copy import deepcopy from typing import Dict, Optional, Sequence, Union import torch @@ -181,18 +182,15 @@ def __init__( event_shape=torch.Size(), validate_args=None, ): + self.prior = prior + self.device = None + self.return_type = return_type super().__init__( batch_shape=batch_shape, event_shape=event_shape, - validate_args=( - prior._validate_args if validate_args is None else validate_args - ), + validate_args=False, ) - self.prior = prior - self.device = None - self.return_type = return_type - def log_prob(self, value) -> Tensor: return torch.as_tensor( self.prior.log_prob(value), @@ -236,6 +234,15 @@ def to(self, device: Union[str, torch.device]) -> None: self.prior = move_distribution_to_device(self.prior, device) self.device = device + def __deepcopy__(self, memo): + """Ensure prior attribute is preserved during deepcopy.""" + cls = self.__class__ + result = cls.__new__(cls) + memo[id(self)] = result + for k, v in self.__dict__.items(): + setattr(result, k, deepcopy(v, memo)) + return result + class MultipleIndependent(Distribution): """Wrap a sequence of PyTorch distributions into a joint PyTorch distribution.""" @@ -434,6 +441,15 @@ def to(self, device: Union[str, torch.device]) -> None: self.dists[i] = move_distribution_to_device(self.dists[i], device) self.device = device + def __deepcopy__(self, memo): + """Ensure dists attribute is preserved during deepcopy.""" + cls = self.__class__ + result = cls.__new__(cls) + memo[id(self)] = result + for k, v in self.__dict__.items(): + setattr(result, k, deepcopy(v, memo)) + return result + def build_support( lower_bound: Optional[Tensor] = None, upper_bound: Optional[Tensor] = None @@ -523,15 +539,13 @@ class OneDimPriorWrapper(Distribution): """ def __init__(self, prior: Distribution, validate_args=None) -> None: + self.prior = prior + self.device = None super().__init__( batch_shape=prior.batch_shape, event_shape=prior.event_shape, - validate_args=( - prior._validate_args if validate_args is None else validate_args - ), + validate_args=False, ) - self.prior = prior - self.device = None def to(self, device: Union[str, torch.device]) -> None: """ @@ -546,6 +560,15 @@ def to(self, device: Union[str, torch.device]) -> None: self.prior = move_distribution_to_device(self.prior, device) self.device = device + def __deepcopy__(self, memo): + """Ensure prior attribute is preserved during deepcopy.""" + cls = self.__class__ + result = cls.__new__(cls) + memo[id(self)] = result + for k, v in self.__dict__.items(): + setattr(result, k, deepcopy(v, memo)) + return result + def sample(self, *args, **kwargs) -> Tensor: return self.prior.sample(*args, **kwargs) diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index 395f09a71..a56f323fa 100644 --- a/tests/inference_on_device_test.py +++ b/tests/inference_on_device_test.py @@ -396,6 +396,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + potential_fn = FakePotential( prior=MultivariateNormal( zeros(num_dim, device=device), eye(num_dim, device=device) @@ -450,6 +456,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + potential_fn = FakePotential(prior=prior, device=device) # Generate mock simulation data on device diff --git a/tests/vi_test.py b/tests/vi_test.py index e1327e56a..f3f187ce9 100644 --- a/tests/vi_test.py +++ b/tests/vi_test.py @@ -57,6 +57,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + def make_tractable_potential(target_distribution, prior): """Create a potential function from a known target distribution.""" @@ -70,6 +76,14 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + def bind( + self, x_o: torch.Tensor, x_is_iid: bool = True + ) -> "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) + return bound + return TractablePotential(prior=prior) From 3e52f56f0e8331f1debda65f197fbc9c29e5147a Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 07:26:34 +0200 Subject: [PATCH 02/14] Add .bind() method to potentials and fix _base_recursor scoping - Add bind() method to BasePotential (abstract) and CustomPotentialWrapper - Add bind() to LikelihoodBasedPotential, PosteriorBasedPotential, RatioBasedPotential - Add bind() to VectorFieldBasedPotential with full parameter support - Add bind() to EnsemblePotential in ensemble_posterior.py - Add bind() to ConditionedPotential in conditional_density_utils.py - Fix _base_recursor to use local variable _active_holder to avoid Python scoping issues with the _active parameter - Update posteriors to use bind() instead of set_x() for immutability - Fix PytorchReturnTypeWrapper, MultipleIndependent, OneDimPriorWrapper init order and add __deepcopy__ methods for proper deepcopy support - Update error messages to reference bind() instead of set_x() - Add bind() to test FakePotential and TractablePotential classes - Replace remaining set_x calls in VectorFieldPosterior with bind --- sbi/inference/posteriors/base_posterior.py | 4 +- .../posteriors/ensemble_posterior.py | 26 +++++++++- .../posteriors/importance_posterior.py | 4 +- sbi/inference/posteriors/mcmc_posterior.py | 29 ++++++------ .../posteriors/rejection_posterior.py | 4 +- .../posteriors/vector_field_posterior.py | 12 +++-- sbi/inference/posteriors/vi_posterior.py | 2 +- sbi/inference/potentials/base_potential.py | 27 ++++++++++- .../potentials/likelihood_based_potential.py | 11 +++++ .../potentials/posterior_based_potential.py | 11 +++++ .../potentials/ratio_based_potential.py | 11 +++++ .../potentials/vector_field_potential.py | 32 +++++++++++++ sbi/utils/conditional_density_utils.py | 12 ++++- sbi/utils/torchutils.py | 17 +++---- sbi/utils/user_input_checks_utils.py | 47 ++++++++++++++----- tests/inference_on_device_test.py | 12 +++++ tests/vi_test.py | 14 ++++++ 17 files changed, 223 insertions(+), 52 deletions(-) diff --git a/sbi/inference/posteriors/base_posterior.py b/sbi/inference/posteriors/base_posterior.py index 599622904..f5111cecb 100644 --- a/sbi/inference/posteriors/base_posterior.py +++ b/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, diff --git a/sbi/inference/posteriors/ensemble_posterior.py b/sbi/inference/posteriors/ensemble_posterior.py index 755f1cebe..fa4015f92 100644 --- a/sbi/inference/posteriors/ensemble_posterior.py +++ b/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. diff --git a/sbi/inference/posteriors/importance_posterior.py b/sbi/inference/posteriors/importance_posterior.py index cfbc3d2f8..fdc98a9a1 100644 --- a/sbi/inference/posteriors/importance_posterior.py +++ b/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( diff --git a/sbi/inference/posteriors/mcmc_posterior.py b/sbi/inference/posteriors/mcmc_posterior.py index f90db67d6..fc1d254cf 100644 --- a/sbi/inference/posteriors/mcmc_posterior.py +++ b/sbi/inference/posteriors/mcmc_posterior.py @@ -2,7 +2,6 @@ # under the Apache License Version 2.0, see 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() diff --git a/sbi/inference/posteriors/rejection_posterior.py b/sbi/inference/posteriors/rejection_posterior.py index 7b05160a2..92e173129 100644 --- a/sbi/inference/posteriors/rejection_posterior.py +++ b/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) diff --git a/sbi/inference/posteriors/vector_field_posterior.py b/sbi/inference/posteriors/vector_field_posterior.py index 20c1990f0..ea27025b5 100644 --- a/sbi/inference/posteriors/vector_field_posterior.py +++ b/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 ) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index e070d904f..d10162722 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/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) diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index d24b05b94..f75f3e050 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/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,7 +66,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 @@ -74,6 +74,20 @@ def x_o(self, x_o: Optional[Tensor]) -> None: """Check the shape of the observed data and, if valid, set it.""" self.set_x(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. @@ -154,3 +168,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, + ) diff --git a/sbi/inference/potentials/likelihood_based_potential.py b/sbi/inference/potentials/likelihood_based_potential.py index d2907430c..991254489 100644 --- a/sbi/inference/potentials/likelihood_based_potential.py +++ b/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))$. diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py index eda384626..91a4af0e3 100644 --- a/sbi/inference/potentials/posterior_based_potential.py +++ b/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. diff --git a/sbi/inference/potentials/ratio_based_potential.py b/sbi/inference/potentials/ratio_based_potential.py index 844405e75..917fede47 100644 --- a/sbi/inference/potentials/ratio_based_potential.py +++ b/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. diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index b7ac4ba45..fbac610fd 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/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,37 @@ 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.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, diff --git a/sbi/utils/conditional_density_utils.py b/sbi/utils/conditional_density_utils.py index 0d48efedb..a42e6b4ed 100644 --- a/sbi/utils/conditional_density_utils.py +++ b/sbi/utils/conditional_density_utils.py @@ -412,7 +412,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): @@ -420,7 +420,15 @@ def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): 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: diff --git a/sbi/utils/torchutils.py b/sbi/utils/torchutils.py index 4405bed65..ce657cf34 100644 --- a/sbi/utils/torchutils.py +++ b/sbi/utils/torchutils.py @@ -677,13 +677,14 @@ def _base_recursor( _active: Object identities on the active recursion path. Used internally to avoid infinite recursion in cyclic object graphs. """ - if _active is None: - _active = set() + _active_holder: Optional[set[int]] = _active + if _active_holder is None: + _active_holder = set() obj_id = id(obj) - if obj_id in _active: + if obj_id in _active_holder: return - _active.add(obj_id) + _active_holder.add(obj_id) if isinstance(obj, Module) and check(obj): action(obj) @@ -698,7 +699,7 @@ def _base_recursor( key=k, check=check, action=action, - _active=_active, + _active=_active_holder, ) elif isinstance(obj, type): # Skip class/type objects to avoid modifying immutable C extension types @@ -715,7 +716,7 @@ def _base_recursor( key=k, check=check, action=action, - _active=_active, + _active=_active_holder, ) elif isinstance(obj, (List, Tuple, Generator)): new_obj = [] @@ -723,12 +724,12 @@ def _base_recursor( if check(o): new_obj.append(action(o)) else: - _base_recursor(o, check=check, action=action, _active=_active) + _base_recursor(o, check=check, action=action, _active=_active_holder) new_obj.append(o) if parent is not None and key is not None: setattr(parent, key, type(obj)(new_obj)) # type: ignore - _active.remove(obj_id) + _active_holder.remove(obj_id) def move_all_tensor_to_device(obj: object, device: Union[str, torch.device]) -> None: diff --git a/sbi/utils/user_input_checks_utils.py b/sbi/utils/user_input_checks_utils.py index dbd8a1d5e..cbb0f66e8 100644 --- a/sbi/utils/user_input_checks_utils.py +++ b/sbi/utils/user_input_checks_utils.py @@ -2,6 +2,7 @@ # under the Apache License Version 2.0, see import warnings +from copy import deepcopy from typing import Dict, Optional, Sequence, Union import torch @@ -181,18 +182,15 @@ def __init__( event_shape=torch.Size(), validate_args=None, ): + self.prior = prior + self.device = None + self.return_type = return_type super().__init__( batch_shape=batch_shape, event_shape=event_shape, - validate_args=( - prior._validate_args if validate_args is None else validate_args - ), + validate_args=False, ) - self.prior = prior - self.device = None - self.return_type = return_type - def log_prob(self, value) -> Tensor: return torch.as_tensor( self.prior.log_prob(value), @@ -236,6 +234,15 @@ def to(self, device: Union[str, torch.device]) -> None: self.prior = move_distribution_to_device(self.prior, device) self.device = device + def __deepcopy__(self, memo): + """Ensure prior attribute is preserved during deepcopy.""" + cls = self.__class__ + result = cls.__new__(cls) + memo[id(self)] = result + for k, v in self.__dict__.items(): + setattr(result, k, deepcopy(v, memo)) + return result + class MultipleIndependent(Distribution): """Wrap a sequence of PyTorch distributions into a joint PyTorch distribution.""" @@ -434,6 +441,15 @@ def to(self, device: Union[str, torch.device]) -> None: self.dists[i] = move_distribution_to_device(self.dists[i], device) self.device = device + def __deepcopy__(self, memo): + """Ensure dists attribute is preserved during deepcopy.""" + cls = self.__class__ + result = cls.__new__(cls) + memo[id(self)] = result + for k, v in self.__dict__.items(): + setattr(result, k, deepcopy(v, memo)) + return result + def build_support( lower_bound: Optional[Tensor] = None, upper_bound: Optional[Tensor] = None @@ -523,15 +539,13 @@ class OneDimPriorWrapper(Distribution): """ def __init__(self, prior: Distribution, validate_args=None) -> None: + self.prior = prior + self.device = None super().__init__( batch_shape=prior.batch_shape, event_shape=prior.event_shape, - validate_args=( - prior._validate_args if validate_args is None else validate_args - ), + validate_args=False, ) - self.prior = prior - self.device = None def to(self, device: Union[str, torch.device]) -> None: """ @@ -546,6 +560,15 @@ def to(self, device: Union[str, torch.device]) -> None: self.prior = move_distribution_to_device(self.prior, device) self.device = device + def __deepcopy__(self, memo): + """Ensure prior attribute is preserved during deepcopy.""" + cls = self.__class__ + result = cls.__new__(cls) + memo[id(self)] = result + for k, v in self.__dict__.items(): + setattr(result, k, deepcopy(v, memo)) + return result + def sample(self, *args, **kwargs) -> Tensor: return self.prior.sample(*args, **kwargs) diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index 395f09a71..a56f323fa 100644 --- a/tests/inference_on_device_test.py +++ b/tests/inference_on_device_test.py @@ -396,6 +396,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + potential_fn = FakePotential( prior=MultivariateNormal( zeros(num_dim, device=device), eye(num_dim, device=device) @@ -450,6 +456,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + potential_fn = FakePotential(prior=prior, device=device) # Generate mock simulation data on device diff --git a/tests/vi_test.py b/tests/vi_test.py index e1327e56a..f3f187ce9 100644 --- a/tests/vi_test.py +++ b/tests/vi_test.py @@ -57,6 +57,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + def make_tractable_potential(target_distribution, prior): """Create a potential function from a known target distribution.""" @@ -70,6 +76,14 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + def bind( + self, x_o: torch.Tensor, x_is_iid: bool = True + ) -> "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) + return bound + return TractablePotential(prior=prior) From 8fdc10cf3c26b0f42a2e38d646050e06ad10dccc Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 08:26:23 +0200 Subject: [PATCH 03/14] fix attempt with bound.neural_ode.params.update --- sbi/inference/potentials/vector_field_potential.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index fbac610fd..5c1783beb 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/sbi/inference/potentials/vector_field_potential.py @@ -149,6 +149,7 @@ def bind( iid_params=self.iid_params, neural_ode_backend=self.neural_ode_backend, ) + bound.neural_ode.params.update(self.neural_ode.params) bound.set_x( x_o, x_is_iid=x_is_iid, From 968b58bf1660446faba628e345a8e0a1d7636078 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 08:54:05 +0200 Subject: [PATCH 04/14] checkpoint after full bind replacement of set_x From ee2adfa4327a592ca71e98d6c2b17271f1ddc98e Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 09:06:20 +0200 Subject: [PATCH 05/14] fixup:slightly wrong last merge --- sbi/inference/potentials/base_potential.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index f75f3e050..f353d535f 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/sbi/inference/potentials/base_potential.py @@ -69,11 +69,6 @@ def x_o(self) -> Tensor: "No observed data is available. Use `potential_fn.bind(x_o)`." ) - @x_o.setter - def x_o(self, x_o: Optional[Tensor]) -> None: - """Check the shape of the observed data and, if valid, set it.""" - self.set_x(x_o) - def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "BasePotential": """Create new potential with x bound, without mutable state. From 0fe9359ed1b19dc0788a025a975484c0a13308d2 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 09:50:11 +0200 Subject: [PATCH 06/14] 3 things: remove set_x from BasePotential.__init__; replace overlooked set_x inside vi_divergence; remove set_x usage inside bind() --- sbi/inference/potentials/base_potential.py | 5 ++- .../potentials/likelihood_based_potential.py | 6 +++- .../potentials/posterior_based_potential.py | 13 ++++++-- .../potentials/ratio_based_potential.py | 6 +++- .../potentials/vector_field_potential.py | 31 ++++++++++++------- sbi/samplers/vi/vi_divergence_optimizers.py | 4 +-- 6 files changed, 46 insertions(+), 19 deletions(-) diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index f353d535f..7b9676812 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/sbi/inference/potentials/base_potential.py @@ -30,7 +30,10 @@ def __init__( """ self.device = device self.prior = prior - self.set_x(x_o) + if x_o is not None: + x_o = process_x(x_o).to(self.device) + self._x_o = x_o + self._x_is_iid = True @abstractmethod def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: diff --git a/sbi/inference/potentials/likelihood_based_potential.py b/sbi/inference/potentials/likelihood_based_potential.py index 991254489..8513b5570 100644 --- a/sbi/inference/potentials/likelihood_based_potential.py +++ b/sbi/inference/potentials/likelihood_based_potential.py @@ -96,13 +96,17 @@ def to(self, device: Union[str, torch.device]) -> "LikelihoodBasedPotential": def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "LikelihoodBasedPotential": """Create new potential with x bound, without mutable state.""" + from sbi.utils.user_input_checks import process_x + 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) + x_o = process_x(x_o).to(self.device) + bound._x_o = x_o + bound._x_is_iid = x_is_iid return bound def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py index 91a4af0e3..f5a31cd1e 100644 --- a/sbi/inference/potentials/posterior_based_potential.py +++ b/sbi/inference/potentials/posterior_based_potential.py @@ -102,20 +102,29 @@ def to(self, device: Union[str, torch.device]) -> "PosteriorBasedPotential": def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential": """Create new potential with x bound, without mutable state.""" + from sbi.utils.user_input_checks import process_x + 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) + x_o = process_x(x_o).to(self.device) + bound._x_o = x_o + bound._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. """ - super().set_x(x_o, x_is_iid=x_is_iid) + if x_o is not None: + from sbi.utils.user_input_checks import process_x + + x_o = process_x(x_o).to(self.device) + self._x_o = x_o + self._x_is_iid = x_is_iid def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential for posterior-based methods. diff --git a/sbi/inference/potentials/ratio_based_potential.py b/sbi/inference/potentials/ratio_based_potential.py index 917fede47..f18de63ee 100644 --- a/sbi/inference/potentials/ratio_based_potential.py +++ b/sbi/inference/potentials/ratio_based_potential.py @@ -84,13 +84,17 @@ def to(self, device: Union[str, torch.device]) -> "RatioBasedPotential": def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "RatioBasedPotential": """Create new potential with x bound, without mutable state.""" + from sbi.utils.user_input_checks import process_x + 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) + x_o = process_x(x_o).to(self.device) + bound._x_o = x_o + bound._x_is_iid = x_is_iid return bound def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index 5c1783beb..1592b9493 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/sbi/inference/potentials/vector_field_potential.py @@ -119,7 +119,12 @@ def set_x( `IIDScoreFunction`. ode_kwargs: Additional keyword arguments for the neural ODE. """ - super().set_x(x_o, x_is_iid) + if x_o is not None: + from sbi.utils.user_input_checks import process_x + + 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 @@ -140,25 +145,27 @@ def bind( **ode_kwargs, ) -> "VectorFieldBasedPotential": """Create new potential with x bound, without mutable state.""" + from sbi.utils.user_input_checks import process_x + 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, + iid_method=iid_method or self.iid_method, + iid_params=iid_params, neural_ode_backend=self.neural_ode_backend, ) bound.neural_ode.params.update(self.neural_ode.params) - 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, - ) + x_o = process_x(x_o).to(self.device) + bound._x_o = x_o + bound._x_is_iid = x_is_iid + bound.guidance_method = guidance_method + bound.guidance_params = guidance_params + if not x_is_iid and (bound._x_o is not None): + bound.flow = bound.rebuild_flow(**ode_kwargs) + elif bound._x_o is not None: + bound.flows = bound.rebuild_flows_for_batch(**ode_kwargs) return bound def __call__( diff --git a/sbi/samplers/vi/vi_divergence_optimizers.py b/sbi/samplers/vi/vi_divergence_optimizers.py index 24ab3d9f1..e57ab621a 100644 --- a/sbi/samplers/vi/vi_divergence_optimizers.py +++ b/sbi/samplers/vi/vi_divergence_optimizers.py @@ -479,7 +479,7 @@ def generate_elbo_particles( samples = self.q.rsample(torch.Size((num_samples,))) log_q = self.q.log_prob(samples) - self.potential_fn.set_x(x_o) + self.potential_fn = self.potential_fn.bind(x_o) log_potential = self.potential_fn(samples) elbo = log_potential - log_q return elbo @@ -638,7 +638,7 @@ def _loss_q_proposal(self, x_o: Tensor) -> Tuple[Tensor, Tensor]: if hasattr(self.q, "clear_cache"): self.q.clear_cache() logq = self.q.log_prob(samples) - self.potential_fn.set_x(x_o) + self.potential_fn = self.potential_fn.bind(x_o) logp = self.potential_fn(samples) with torch.no_grad(): logweights = logp - logq From e8444c5f8ed99d231ac442c8d982d278f005b852 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 10:44:23 +0200 Subject: [PATCH 07/14] fix attempt: VI training broken due to bind --- sbi/inference/posteriors/vi_posterior.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index d10162722..80e0cf228 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -880,6 +880,10 @@ def train( break # Training finished: self._trained_on = x + + # Bind potential_fn to the trained x so it can be used for sampling/evaluation + self.potential_fn = self.potential_fn.bind(x) + if self._mode == "amortized": warnings.warn( "Switching from amortized to single-x mode. " From 81b070945204e196e62adfd1de423b68824d5f1e Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 10:56:50 +0200 Subject: [PATCH 08/14] fixup: AttributeError: 'VectorFieldBasedPotential' object has no attribute 'guidance_method' --- sbi/inference/potentials/vector_field_potential.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index 1592b9493..208071e62 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/sbi/inference/potentials/vector_field_potential.py @@ -61,6 +61,8 @@ def __init__( self.vector_field_estimator.eval() self.iid_method = iid_method self.iid_params = iid_params + self.guidance_method: Optional[str] = None + self.guidance_params: Optional[Dict[str, Any]] = None self.neural_ode_backend = neural_ode_backend neural_ode_kwargs = neural_ode_kwargs or {} @@ -153,15 +155,19 @@ def bind( x_o=None, device=self.device, iid_method=iid_method or self.iid_method, - iid_params=iid_params, + iid_params=iid_params if iid_params is not None else self.iid_params, neural_ode_backend=self.neural_ode_backend, ) bound.neural_ode.params.update(self.neural_ode.params) x_o = process_x(x_o).to(self.device) bound._x_o = x_o bound._x_is_iid = x_is_iid - bound.guidance_method = guidance_method - bound.guidance_params = guidance_params + bound.guidance_method = ( + guidance_method if guidance_method is not None else self.guidance_method + ) + bound.guidance_params = ( + guidance_params if guidance_params is not None else self.guidance_params + ) if not x_is_iid and (bound._x_o is not None): bound.flow = bound.rebuild_flow(**ode_kwargs) elif bound._x_o is not None: From 3292830de33580ba29f63fc90359e4e50ecbf6b8 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 13:21:09 +0200 Subject: [PATCH 09/14] set_x-bind-swap: set_x now internally uses bind to implement a statefull version of it --- .../posteriors/ensemble_posterior.py | 22 +++++++---- sbi/inference/potentials/base_potential.py | 19 +++++++--- .../potentials/posterior_based_potential.py | 16 +++++--- .../potentials/vector_field_potential.py | 38 ++++++++++++------- sbi/utils/conditional_density_utils.py | 20 +++++++--- 5 files changed, 78 insertions(+), 37 deletions(-) diff --git a/sbi/inference/posteriors/ensemble_posterior.py b/sbi/inference/posteriors/ensemble_posterior.py index fa4015f92..cc7c276e0 100644 --- a/sbi/inference/posteriors/ensemble_posterior.py +++ b/sbi/inference/posteriors/ensemble_posterior.py @@ -479,14 +479,20 @@ 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. + """ + import warnings + + warnings.warn( + "set_x() is deprecated, 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, diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index 7b9676812..946e9b134 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/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: diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py index f5a31cd1e..50679b2e9 100644 --- a/sbi/inference/potentials/posterior_based_potential.py +++ b/sbi/inference/potentials/posterior_based_potential.py @@ -118,13 +118,19 @@ def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential" 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. """ - if x_o is not None: - from sbi.utils.user_input_checks import process_x + import warnings - x_o = process_x(x_o).to(self.device) - self._x_o = x_o - self._x_is_iid = x_is_iid + 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 def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential for posterior-based methods. diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index 208071e62..d61161be4 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/sbi/inference/potentials/vector_field_potential.py @@ -112,6 +112,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. + Args: x_o: The observed data. x_is_iid: Whether the observed data is IID (if batch_dim>1). @@ -121,20 +123,30 @@ def set_x( `IIDScoreFunction`. ode_kwargs: Additional keyword arguments for the neural ODE. """ - if x_o is not None: - from sbi.utils.user_input_checks import process_x + import warnings - 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, 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 + self.flow = bound.flow + self.flows = bound.flows def bind( self, diff --git a/sbi/utils/conditional_density_utils.py b/sbi/utils/conditional_density_utils.py index 67115b850..73e35c5ef 100644 --- a/sbi/utils/conditional_density_utils.py +++ b/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): @@ -416,11 +415,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.""" From 6e9020ef925ca931f56dc5474c699d0483432407 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 13:56:03 +0200 Subject: [PATCH 10/14] fixup: flows and flow --- sbi/inference/potentials/vector_field_potential.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index d61161be4..2872ee305 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/sbi/inference/potentials/vector_field_potential.py @@ -145,8 +145,10 @@ def set_x( self.iid_params = bound.iid_params self.guidance_method = bound.guidance_method self.guidance_params = bound.guidance_params - self.flow = bound.flow - self.flows = bound.flows + 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, From 43e46b8bdbb2ba3ee64300cb747476e5288c7728 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 14:08:05 +0200 Subject: [PATCH 11/14] fix: some places still used set_x inside bind -> recursion loop... --- tests/inference_on_device_test.py | 12 ++++++++++-- tests/vi_test.py | 12 ++++++++++-- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index a56f323fa..d6fc9e845 100644 --- a/tests/inference_on_device_test.py +++ b/tests/inference_on_device_test.py @@ -398,8 +398,12 @@ 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.""" + from sbi.utils.user_input_checks import process_x + 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( @@ -458,8 +462,12 @@ 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.""" + from sbi.utils.user_input_checks import process_x + 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) diff --git a/tests/vi_test.py b/tests/vi_test.py index f3f187ce9..fa5414ee1 100644 --- a/tests/vi_test.py +++ b/tests/vi_test.py @@ -59,8 +59,12 @@ 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.""" + from sbi.utils.user_input_checks import process_x + 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 @@ -80,8 +84,12 @@ def bind( self, x_o: torch.Tensor, x_is_iid: bool = True ) -> "TractablePotential": """Create new potential with x bound, without mutable state.""" + from sbi.utils.user_input_checks import process_x + 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) From 12944af3e54f479c409ed3c342354d682084643b Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Thu, 23 Jul 2026 14:27:49 +0200 Subject: [PATCH 12/14] reverted user_input_checks_utils changes; moved imports up; fixed iid_method for VectorFieldBasedPotential; --- .../potentials/likelihood_based_potential.py | 3 +- .../potentials/posterior_based_potential.py | 19 +++------ .../potentials/ratio_based_potential.py | 3 +- .../potentials/vector_field_potential.py | 10 ++--- sbi/utils/user_input_checks_utils.py | 40 ++++--------------- 5 files changed, 20 insertions(+), 55 deletions(-) diff --git a/sbi/inference/potentials/likelihood_based_potential.py b/sbi/inference/potentials/likelihood_based_potential.py index 8513b5570..ce15d204e 100644 --- a/sbi/inference/potentials/likelihood_based_potential.py +++ b/sbi/inference/potentials/likelihood_based_potential.py @@ -19,6 +19,7 @@ ) from sbi.sbi_types import TorchTransform from sbi.utils.sbiutils import mcmc_transform +from sbi.utils.user_input_checks import process_x def likelihood_estimator_based_potential( @@ -96,8 +97,6 @@ def to(self, device: Union[str, torch.device]) -> "LikelihoodBasedPotential": def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "LikelihoodBasedPotential": """Create new potential with x bound, without mutable state.""" - from sbi.utils.user_input_checks import process_x - bound = LikelihoodBasedPotential( likelihood_estimator=self.likelihood_estimator, prior=self.prior, diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py index 50679b2e9..791bc8c30 100644 --- a/sbi/inference/potentials/posterior_based_potential.py +++ b/sbi/inference/potentials/posterior_based_potential.py @@ -21,6 +21,7 @@ within_support, ) from sbi.utils.torchutils import ensure_theta_batched, infer_module_device +from sbi.utils.user_input_checks import process_x def posterior_estimator_based_potential( @@ -102,8 +103,6 @@ def to(self, device: Union[str, torch.device]) -> "PosteriorBasedPotential": def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential": """Create new potential with x bound, without mutable state.""" - from sbi.utils.user_input_checks import process_x - bound = PosteriorBasedPotential( posterior_estimator=self.posterior_estimator, prior=self.prior, @@ -118,19 +117,11 @@ def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential" 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. """ - 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 + 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 def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential for posterior-based methods. diff --git a/sbi/inference/potentials/ratio_based_potential.py b/sbi/inference/potentials/ratio_based_potential.py index f18de63ee..6b328c933 100644 --- a/sbi/inference/potentials/ratio_based_potential.py +++ b/sbi/inference/potentials/ratio_based_potential.py @@ -11,6 +11,7 @@ from sbi.sbi_types import TorchTransform from sbi.utils.sbiutils import match_theta_and_x_batch_shapes, mcmc_transform from sbi.utils.torchutils import atleast_2d +from sbi.utils.user_input_checks import process_x def ratio_estimator_based_potential( @@ -84,8 +85,6 @@ def to(self, device: Union[str, torch.device]) -> "RatioBasedPotential": def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "RatioBasedPotential": """Create new potential with x bound, without mutable state.""" - from sbi.utils.user_input_checks import process_x - bound = RatioBasedPotential( ratio_estimator=self.ratio_estimator, prior=self.prior, diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index 2872ee305..17108d2c3 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/sbi/inference/potentials/vector_field_potential.py @@ -22,6 +22,7 @@ from sbi.sbi_types import TorchTransform from sbi.utils.sbiutils import mcmc_transform, within_support from sbi.utils.torchutils import ensure_theta_batched +from sbi.utils.user_input_checks import process_x class VectorFieldBasedPotential(BasePotential): @@ -113,7 +114,7 @@ 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). @@ -126,7 +127,8 @@ def set_x( import warnings warnings.warn( - "set_x() is deprecated, use bind() instead", + "set_x() is deprecated and will be removed in a future release. " + "Use bind() instead.", FutureWarning, stacklevel=2, ) @@ -161,14 +163,12 @@ def bind( **ode_kwargs, ) -> "VectorFieldBasedPotential": """Create new potential with x bound, without mutable state.""" - from sbi.utils.user_input_checks import process_x - bound = VectorFieldBasedPotential( vector_field_estimator=self.vector_field_estimator, prior=self.prior, x_o=None, device=self.device, - iid_method=iid_method or self.iid_method, + iid_method=iid_method if iid_method is not None else self.iid_method, iid_params=iid_params if iid_params is not None else self.iid_params, neural_ode_backend=self.neural_ode_backend, ) diff --git a/sbi/utils/user_input_checks_utils.py b/sbi/utils/user_input_checks_utils.py index cbb0f66e8..4a6ff8d0e 100644 --- a/sbi/utils/user_input_checks_utils.py +++ b/sbi/utils/user_input_checks_utils.py @@ -2,7 +2,6 @@ # under the Apache License Version 2.0, see import warnings -from copy import deepcopy from typing import Dict, Optional, Sequence, Union import torch @@ -188,7 +187,9 @@ def __init__( super().__init__( batch_shape=batch_shape, event_shape=event_shape, - validate_args=False, + validate_args=( + prior._validate_args if validate_args is None else validate_args + ), ) def log_prob(self, value) -> Tensor: @@ -234,15 +235,6 @@ def to(self, device: Union[str, torch.device]) -> None: self.prior = move_distribution_to_device(self.prior, device) self.device = device - def __deepcopy__(self, memo): - """Ensure prior attribute is preserved during deepcopy.""" - cls = self.__class__ - result = cls.__new__(cls) - memo[id(self)] = result - for k, v in self.__dict__.items(): - setattr(result, k, deepcopy(v, memo)) - return result - class MultipleIndependent(Distribution): """Wrap a sequence of PyTorch distributions into a joint PyTorch distribution.""" @@ -441,15 +433,6 @@ def to(self, device: Union[str, torch.device]) -> None: self.dists[i] = move_distribution_to_device(self.dists[i], device) self.device = device - def __deepcopy__(self, memo): - """Ensure dists attribute is preserved during deepcopy.""" - cls = self.__class__ - result = cls.__new__(cls) - memo[id(self)] = result - for k, v in self.__dict__.items(): - setattr(result, k, deepcopy(v, memo)) - return result - def build_support( lower_bound: Optional[Tensor] = None, upper_bound: Optional[Tensor] = None @@ -539,13 +522,15 @@ class OneDimPriorWrapper(Distribution): """ def __init__(self, prior: Distribution, validate_args=None) -> None: - self.prior = prior - self.device = None super().__init__( batch_shape=prior.batch_shape, event_shape=prior.event_shape, - validate_args=False, + validate_args=( + prior._validate_args if validate_args is None else validate_args + ), ) + self.prior = prior + self.device = None def to(self, device: Union[str, torch.device]) -> None: """ @@ -560,15 +545,6 @@ def to(self, device: Union[str, torch.device]) -> None: self.prior = move_distribution_to_device(self.prior, device) self.device = device - def __deepcopy__(self, memo): - """Ensure prior attribute is preserved during deepcopy.""" - cls = self.__class__ - result = cls.__new__(cls) - memo[id(self)] = result - for k, v in self.__dict__.items(): - setattr(result, k, deepcopy(v, memo)) - return result - def sample(self, *args, **kwargs) -> Tensor: return self.prior.sample(*args, **kwargs) From 0a15ef23eead2ccafee443a409acb49b52f75b2f Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Thu, 23 Jul 2026 14:29:35 +0200 Subject: [PATCH 13/14] moved imports up; changed warnings to include a note on future deletion --- .../posteriors/ensemble_posterior.py | 4 +++- .../potentials/posterior_based_potential.py | 21 +++++++++++++------ tests/inference_on_device_test.py | 6 +----- tests/vi_test.py | 5 +---- 4 files changed, 20 insertions(+), 16 deletions(-) diff --git a/sbi/inference/posteriors/ensemble_posterior.py b/sbi/inference/posteriors/ensemble_posterior.py index cc7c276e0..e87eeab03 100644 --- a/sbi/inference/posteriors/ensemble_posterior.py +++ b/sbi/inference/posteriors/ensemble_posterior.py @@ -482,11 +482,13 @@ def set_x(self, x_o: Optional[Tensor]): """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. """ import warnings warnings.warn( - "set_x() is deprecated, use bind() instead", + "set_x() is deprecated and will be removed in a future release. " + "Use bind() instead.", FutureWarning, stacklevel=2, ) diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py index 791bc8c30..4c82d0cc8 100644 --- a/sbi/inference/potentials/posterior_based_potential.py +++ b/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. diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index d6fc9e845..53d927f8f 100644 --- a/tests/inference_on_device_test.py +++ b/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." @@ -398,8 +398,6 @@ 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.""" - from sbi.utils.user_input_checks import process_x - bound = FakePotential(prior=self.prior, device=self.device) x_o = process_x(x_o).to(self.device) bound._x_o = x_o @@ -462,8 +460,6 @@ 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.""" - from sbi.utils.user_input_checks import process_x - bound = FakePotential(prior=self.prior, device=self.device) x_o = process_x(x_o).to(self.device) bound._x_o = x_o diff --git a/tests/vi_test.py b/tests/vi_test.py index fa5414ee1..9aac865cf 100644 --- a/tests/vi_test.py +++ b/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"] @@ -59,8 +60,6 @@ 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.""" - from sbi.utils.user_input_checks import process_x - bound = FakePotential(prior=self.prior, device=self.device) x_o = process_x(x_o).to(self.device) bound._x_o = x_o @@ -84,8 +83,6 @@ def bind( self, x_o: torch.Tensor, x_is_iid: bool = True ) -> "TractablePotential": """Create new potential with x bound, without mutable state.""" - from sbi.utils.user_input_checks import process_x - bound = TractablePotential(prior=self.prior, device=self.device) x_o = process_x(x_o).to(self.device) bound._x_o = x_o From 52476838af54d9c56a125367f4a16fe11bcef9ac Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Thu, 23 Jul 2026 23:52:40 +0200 Subject: [PATCH 14/14] moved warnings import up --- sbi/inference/posteriors/ensemble_posterior.py | 2 +- sbi/inference/potentials/vector_field_potential.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sbi/inference/posteriors/ensemble_posterior.py b/sbi/inference/posteriors/ensemble_posterior.py index e87eeab03..042c19265 100644 --- a/sbi/inference/posteriors/ensemble_posterior.py +++ b/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 +import warnings from typing import List, Optional, Tuple, Union import torch @@ -484,7 +485,6 @@ def set_x(self, x_o: Optional[Tensor]): DEPRECATED: Use bind() instead. This method delegates to bind() internally. It will be removed in a future release. """ - import warnings warnings.warn( "set_x() is deprecated and will be removed in a future release. " diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index 17108d2c3..7ce061e49 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/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 +import warnings from typing import Any, Dict, List, Literal, Optional, Tuple, Union import torch @@ -124,7 +125,6 @@ def set_x( `IIDScoreFunction`. ode_kwargs: Additional keyword arguments for the neural ODE. """ - import warnings warnings.warn( "set_x() is deprecated and will be removed in a future release. "