Add .bind() method to potentials by Jocho-Smith · Pull Request #1943 · sbi-dev/sbi

Jocho-Smith · 2026-07-22T05:38:18Z

This PR attempts to do the same as #1940 but without all the rebase/merge in between.

- 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

Jocho-Smith · 2026-07-22T05:51:08Z

Comming from #1940 this is what this PR does

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

Finally we got rid of all the errors with this commit :D

Now only Ensemble and VectorField remains :)

Jocho-Smith · 2026-07-22T06:03:14Z

EnsemblePotential.set_x() falls back to the set_x of the individual potentials within it. I'll not convert this to bind, b/c EnsemblePotential.set_x() is never used internally, but might need deprecation for users that potentially use this. I don't think a specific deprecation warning inside the ensemble set_x is required, since this will fall back to the ones of the potentials which will then throw the error.

- 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

…ind-method-clean

Jocho-Smith · 2026-07-22T06:21:32Z

vector field continues to fail with this:

FAILED tests/linearGaussian_vector_field_test.py::test_c2st_vector_field_on_linearGaussian[vp-3-gaussian-sample_with2] - AssertionError: D-KL=tensor([0.2521]) is more than 2 stds above the average performance.
assert tensor([0.2521]) < 0.25

I've already spend quite some time yesterday debugging this in #1940 maybe I'll leave this to someone more familiar with the actual method and continue with my deprecation plan?

Jocho-Smith · 2026-07-22T06:34:54Z

OMG it worked!

The fix was adding a single line to copy the neural ODE params when binding:

bound.neural_ode.params.update(self.neural_ode.params)

This ensures runtime parameter modifications (like exact, atol, rtol) persist across bind() calls, which was breaking the test's log probability evaluation.

dgedon

The name of the PR with the fix _base_recursor scoping is a bit misleading since it's mainly about the bind(), right. Is the base recursor scoping orthogonal?

dgedon · 2026-07-22T07:32:37Z

+            init_fn = self._build_mcmc_init_fn(
+                self.proposal,
+                potential_fn=potential_bound,
+                transform=self.theta_transform,
+                init_strategy=init_strategy,  # type: ignore
+                **kwargs,
+            )


Do we understand the implications of moving that inside the loop? What is happening there and is it fine to move it inside the loop?

Here is how I understand it:

_build_mcmc_init_fn takes a potential_fn which is used for weighting samples in sir and resample init strategies

Each xi in the batch has a different bound potential (potential_.bind(xi)). The initialization needs to evaluate the likelihood at the correct xi value to weight/resample properly

The old code passed self.potential_fn (unbound) to _build_mcmc_init_fn

Now each xi gets its own correctly-bound init function, otherwise this fails with bind()

There's a small overhead of rebuilding init_fn for each xi, but this is necessary for correctness.

dgedon · 2026-07-22T07:33:35Z

            Tensor: initial parameters, one for each chain
        """

-        potential_ = deepcopy(self.potential_fn)


Do we understand why a deepcopy was necessary here and why we can replace it it with an assignment to self?

Jan told me this on discord:

And batched MCMC does deepcopy: look at _get_initial_params_batched in mcmc_posterior.py, first line is potential = deepcopy(self.potentialfn), and then potential.set_x(xi) inside the loop over observations.
The deepcopy is there purely so the per-xi set_x calls don't hold on to the shared self.potential_fn, it's a workaround for the statefulness.

And here is the nice part: you don't need to figure out how to deepcopy inside the sampling method. in the new design the deepcopy just disappears. Instead of copy-then-mutate you create one cheap bind_observation(raw_potential, xi) per observation inside the loop.

dgedon · 2026-07-22T08:47:04Z

+    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


There is code duplication nwith line 237-244

yes, I removed it.

dgedon · 2026-07-22T08:47:17Z

        self.prior = move_distribution_to_device(self.prior, device)
        self.device = device

+    def __deepcopy__(self, memo):


unexplained why this is needed. Add a comment

also it is unsued, no?

you are correct here. I removed it.

dgedon · 2026-07-22T08:48:07Z

-            validate_args=(
-                prior._validate_args if validate_args is None else validate_args
-            ),
+            validate_args=False,


same as above

dgedon · 2026-07-22T08:55:12Z

Here bind() builds a fresh ConditionedPotential. But ConditionedPotential does not call super().__init__(), so _x_is_iid only ever gets set inside set_x(). Then, bind() builds a new ConditionedPotential straight from the constructor, so _x_is_iid is never set on it. Calling .x_is_iid on a potential returned by .bind() will raise an error

you are right. I fixed it.

Jocho-Smith · 2026-07-22T09:37:19Z

Thanks for the review! You are right to be confused about the deepcopy and torchutils.py changes. I think they sneaked in after I started this new PR here. Let me go over everything again and remove everything that is not needed.

Jocho-Smith · 2026-07-22T09:46:09Z

simplifying torchutils.py and deepcopy worked successfully. All tests pass.

…) is necessary

dgedon · 2026-07-22T12:37:17Z

            condition=self.condition,
            dims_to_sample=self.dims_to_sample,
        )
+        bound._x_is_iid = x_is_iid


This will probably be changed downstream again, once we cahnge the Potentials, right? Then it is fine to keep for now.

What would you like to change about this exactly? Isn't this resolving your concern you had?

Yes this is fine. I was flagging it because so far we used after binding the potnetial a .set_x() here instead. And we plan to remove those set_x() calls later on. So I was wondering if we also remove that x_is_iid setting later on.

codecov · 2026-07-22T13:41:19Z

Codecov Report

❌ Patch coverage is 67.27273% with 18 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.87%. Comparing base (587f922) to head (a0cebbf).
⚠️ Report is 40 commits behind head on gsoc-2026.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
sbi/utils/conditional_density_utils.py	12.50%	7 Missing ⚠️
sbi/inference/posteriors/ensemble_posterior.py	14.28%	6 Missing ⚠️
sbi/inference/posteriors/mcmc_posterior.py	71.42%	2 Missing ⚠️
sbi/inference/posteriors/rejection_posterior.py	50.00%	1 Missing ⚠️
sbi/inference/posteriors/vector_field_posterior.py	75.00%	1 Missing ⚠️
sbi/inference/potentials/base_potential.py	75.00%	1 Missing ⚠️

Additional details and impacted files

@@              Coverage Diff              @@
##           gsoc-2026    #1943      +/-   ##
=============================================
- Coverage      87.99%   87.87%   -0.13%     
=============================================
  Files            144      144              
  Lines          13499    13637     +138     
=============================================
+ Hits           11879    11984     +105     
- Misses          1620     1653      +33

Flag	Coverage Δ
fast	`81.56% <67.27%> (?)`

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
sbi/inference/posteriors/base_posterior.py	`85.18% <100.00%> (ø)`
sbi/inference/posteriors/importance_posterior.py	`76.62% <100.00%> (ø)`
sbi/inference/posteriors/vi_posterior.py	`78.55% <100.00%> (ø)`
...inference/potentials/likelihood_based_potential.py	`82.82% <100.00%> (+0.72%)`	⬆️
.../inference/potentials/posterior_based_potential.py	`96.61% <100.00%> (+0.24%)`	⬆️
sbi/inference/potentials/ratio_based_potential.py	`100.00% <100.00%> (ø)`
sbi/inference/potentials/vector_field_potential.py	`91.40% <100.00%> (+0.42%)`	⬆️
sbi/inference/posteriors/rejection_posterior.py	`76.36% <50.00%> (-1.82%)`	⬇️
sbi/inference/posteriors/vector_field_posterior.py	`77.55% <75.00%> (-0.69%)`	⬇️
sbi/inference/potentials/base_potential.py	`91.22% <75.00%> (-1.63%)`	⬇️
... and 3 more

... and 3 files with indirect coverage changes

Jocho-Smith mentioned this pull request Jul 22, 2026

bind implementation still using set_x as fallback yet #1940

Closed

Jocho-Smith added 2 commits July 22, 2026 08:04

Merge branch 'bind-method-clean' of github.com:jocho-smith/sbi into b…

30459d9

…ind-method-clean

fix attempt with bound.neural_ode.params.update

8fdc10c

dgedon self-requested a review July 22, 2026 06:49

Jocho-Smith and others added 3 commits July 22, 2026 08:54

checkpoint after full bind replacement of set_x

968b58b

Merge branch 'gsoc-2026' into bind-method-clean

84c4731

fixup:slightly wrong last merge

ee2adfa

Jocho-Smith mentioned this pull request Jul 22, 2026

Separate iid from set x #1945

Merged

1 task

dgedon reviewed Jul 22, 2026

View reviewed changes

Jocho-Smith changed the title ~~Add .bind() method to potentials and fix _base_recursor scoping~~ Add .bind() method to potentials Jul 22, 2026

revert unnecessary torchutils changes

b98bd84

Jocho-Smith added 3 commits July 22, 2026 11:47

remove deepcopy utils

a6284bc

add comment why bound.neural_ode.params.update(self.neural_ode.params…

d345ffb

…) is necessary

fixup: restored unnecessary changes

a0cebbf

Jocho-Smith requested a review from dgedon July 22, 2026 10:32

Jocho-Smith mentioned this pull request Jul 22, 2026

Test set x deprecation #1946

Merged

dgedon approved these changes Jul 22, 2026

View reviewed changes

dgedon merged commit b6eff5e into sbi-dev:gsoc-2026 Jul 22, 2026
15 checks passed

Jocho-Smith mentioned this pull request Jul 23, 2026

GSoC2026: Project 3 dashboard #1950

Open

8 tasks

Uh oh!

Conversation

Jocho-Smith commented Jul 22, 2026

Uh oh!

Jocho-Smith commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Jocho-Smith commented Jul 22, 2026

Uh oh!

Jocho-Smith commented Jul 22, 2026

Uh oh!

Jocho-Smith commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dgedon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jocho-Smith commented Jul 22, 2026

Uh oh!

Jocho-Smith commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Jocho-Smith Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

codecov Bot commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Jocho-Smith commented Jul 22, 2026 •

edited

Loading

Jocho-Smith commented Jul 22, 2026 •

edited

Loading

Jocho-Smith commented Jul 22, 2026 •

edited

Loading

Jocho-Smith Jul 22, 2026 •

edited

Loading

codecov Bot commented Jul 22, 2026 •

edited

Loading