From 5c2b72aac5022ce60c081acc2fcff224afb8c7d3 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Wed, 22 Jul 2026 16:56:38 +0200 Subject: [PATCH 1/4] migrated tests from set_x to bind --- tests/inference_on_device_test.py | 16 ++++++++++++---- tests/linearGaussian_snpe_test.py | 2 +- tests/linearGaussian_vector_field_test.py | 2 +- tests/posterior_nn_test.py | 2 +- tests/score_samplers_test.py | 10 ++++++---- tests/vi_test.py | 12 ++++++++++-- 6 files changed, 31 insertions(+), 13 deletions(-) diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index a56f323fa..ca3648530 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) @@ -770,8 +778,8 @@ def test_to_method_on_npe_posteriors(trained_npe_for_device_test, posterior_para assert sample_device.device.type == device.split(":")[0], ( f"sample was not correctly moved to {device}." ) - posterior.potential_fn.set_x(x_o) - potential_values = posterior.potential_fn(sample_device) + bound_potential = posterior.potential_fn.bind(x_o) + potential_values = bound_potential(sample_device) assert potential_values.device.type == device.split(":")[0], ( f"potential was not correctly evaluated on {device}." ) diff --git a/tests/linearGaussian_snpe_test.py b/tests/linearGaussian_snpe_test.py index 1603b3294..1a77e30a2 100644 --- a/tests/linearGaussian_snpe_test.py +++ b/tests/linearGaussian_snpe_test.py @@ -685,7 +685,7 @@ def simulator(theta): condition=samples[0], dims_to_sample=[dim_to_sample_1, dim_to_sample_2], ) - conditioned_potential_fn.set_x(x_o, x_is_iid=False) + conditioned_potential_fn = conditioned_potential_fn.bind(x_o, x_is_iid=False) mcmc_posterior = MCMCPosterior( potential_fn=conditioned_potential_fn, theta_transform=restricted_tf, diff --git a/tests/linearGaussian_vector_field_test.py b/tests/linearGaussian_vector_field_test.py index d7eb4fb40..8883e1d6e 100644 --- a/tests/linearGaussian_vector_field_test.py +++ b/tests/linearGaussian_vector_field_test.py @@ -583,7 +583,7 @@ def simulator(theta): condition=samples[0], dims_to_sample=[dim_to_sample_1, dim_to_sample_2], ) - conditioned_potential_fn.set_x(x_o, x_is_iid=False) + conditioned_potential_fn = conditioned_potential_fn.bind(x_o, x_is_iid=False) mcmc_posterior = MCMCPosterior( potential_fn=conditioned_potential_fn, theta_transform=restricted_tf, diff --git a/tests/posterior_nn_test.py b/tests/posterior_nn_test.py index c956fbcbd..2aefc965b 100644 --- a/tests/posterior_nn_test.py +++ b/tests/posterior_nn_test.py @@ -516,7 +516,7 @@ def test_posterior_based_potential_iid_log_prob(iid_batch_size: int): potential_fn, _ = posterior_estimator_based_potential( posterior_estimator, prior, x_o=None ) - potential_fn.set_x(x_o, x_is_iid=True) + potential_fn = potential_fn.bind(x_o, x_is_iid=True) posterior_samples = true_posterior.sample((num_posterior_samples,)) true_prob = true_posterior.log_prob(posterior_samples) diff --git a/tests/score_samplers_test.py b/tests/score_samplers_test.py index f480a3bd6..3c8c57776 100644 --- a/tests/score_samplers_test.py +++ b/tests/score_samplers_test.py @@ -67,10 +67,10 @@ def test_score_fn_iid_on_different_priors(sde_type, iid_method, num_dim): priors = build_random_priors(num_dim) x_o_iid = torch.ones((5, 1)) - score_fn.set_x(x_o_iid, x_is_iid=True, iid_method=iid_method) + score_fn = score_fn.bind(x_o_iid, x_is_iid=True, iid_method=iid_method) inputs = torch.ones((1, 1, num_dim)) + time = torch.ones(1) for prior in priors: - time = torch.ones(1) score_fn.prior = prior output = score_fn.gradient(inputs, time=time) @@ -116,7 +116,9 @@ def test_score_fn_guidance_general(sde_type, guidance_method, num_dim): score_fn = _build_gaussian_score_estimator(sde_type, (num_dim,), mean0, std0) x_o = torch.ones((1, 1)) guidance_name, guidance_params = guidance_method - score_fn.set_x(x_o, guidance_method=guidance_name, guidance_params=guidance_params) + score_fn = score_fn.bind( + x_o, guidance_method=guidance_name, guidance_params=guidance_params + ) inputs = torch.ones((1, 1, num_dim)) time = torch.ones(1) @@ -149,7 +151,7 @@ def test_score_fn_combined_guidance_and_iid(sde_type, iid_method, guidance_metho x_o_iid = torch.ones((5, 1)) guidance_name, guidance_params = guidance_method - score_fn.set_x( + score_fn = score_fn.bind( x_o_iid, x_is_iid=True, iid_method=iid_method, 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 b93fd7d4e2a8707f57e59842139bb60e35575eec Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Thu, 23 Jul 2026 15:09:21 +0200 Subject: [PATCH 2/4] replaced set_x with bind in SimplePotential inside a test --- tests/rejection_sampling_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/rejection_sampling_test.py b/tests/rejection_sampling_test.py index 123733227..cff0c50c2 100644 --- a/tests/rejection_sampling_test.py +++ b/tests/rejection_sampling_test.py @@ -87,7 +87,7 @@ class SimplePotential: def __call__(self, theta, x_o=None): return torch.zeros(theta.shape[0]) - def set_x(self, x): + def bind(self, x_o, x_is_iid=True): pass def to(self, device): From 0b0f7ba7bf31cf60e379819d598164de649ea524 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Fri, 24 Jul 2026 08:21:16 +0200 Subject: [PATCH 3/4] moved imports up --- tests/inference_on_device_test.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index ca3648530..7c17649c7 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,7 +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) @@ -462,7 +461,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) From eed5805c023d3e168b7853b1093c46b5fb8bec68 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Fri, 24 Jul 2026 08:30:31 +0200 Subject: [PATCH 4/4] moved another import up --- tests/vi_test.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/vi_test.py b/tests/vi_test.py index fa5414ee1..551689e07 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,7 +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) @@ -84,7 +84,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)