migrated tests from set_x to bind

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}."
     )

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,

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,

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)

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):

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,

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)