PR #1-2, Weeks 3-5: implement core protocol + bind_observation + minimal notebook example

[Jocho-Smith]

Implements the first GSoC 2026 deliverable for the PotentialFunction API redesign.

Layer 1 - Core Protocol:

• Add PotentialFunction protocol (stateless callable: theta, x -> log_prob)
• Add PotentialFunctionWithDevice protocol with device property
  • maybe over-engineered that these are two separate protocols?
• Add validate_potential() helper for runtime protocol validation
  • (somewhat optional). Is something like this of use?
  • "import inspect" inside function. What do you think about this?

Layer 2 - Binding Utilities:

• Add bind_observation() utility for creating sampler-ready functions
• Add BoundPotential class as alternative class-based binding
• Implement IID handling via sum_iid parameter
• Implement device validation with error messages

Testing:

• Add comprehensive unit tests (18 tests) covering:
  • Protocol validation and compliance
  • bind_observation functionality
  • IID handling (sum_iid)
  • Device validation and mismatch errors
  • Statelessness (property-based: repeated calls yield identical outputs)
  • Gradient tracking (pure pytorch)
  • Backward compatibility with existing LikelihoodBasedPotential

Open questions:

• should I (and maybe Satwik too) merge this to another 'gsoc specific' branch instead of sbi/main? Or use the github functionality "Create draft pull request"?
  • I think it might be useful to stage the 'half ready' PRs instead of merging them directly. That way core maintainers don't have to worry about unfinished code in the repo and 'we developers' can work in iterated PRs without having to worry about 'breaking things' too much too early. Does that make sense?

Note:

• I tried opencode for this PR, especially for the tests and doc strings. I didn't use any agent.md/skill.md files for now.
• This PR#1 does not enable the new potential routine yet, only after PR#2 we can start checking numerical tests, deprecation scenarios and first migrations.

[codecov]

Codecov Report

❌ Patch coverage is 81.05263% with 18 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.98%. Comparing base (9cf0d9e) to head (5001697).
⚠️ Report is 11 commits behind head on gsoc-2026.

Files with missing lines	Patch %	Lines
...inference/potentials/likelihood_based_potential.py	64.00%	9 Missing ⚠️
sbi/inference/potentials/binding.py	83.33%	7 Missing ⚠️
sbi/inference/potentials/base_potential.py	87.50%	1 Missing ⚠️
sbi/inference/potentials/protocol.py	94.44%	1 Missing ⚠️

Additional details and impacted files

@@              Coverage Diff              @@
##           gsoc-2026    #1876      +/-   ##
=============================================
+ Coverage      87.89%   87.98%   +0.09%     
=============================================
  Files            143      146       +3     
  Lines          13353    13684     +331     
=============================================
+ Hits           11736    12040     +304     
- Misses          1617     1644      +27     

Flag	Coverage Δ
fast	81.52% <81.05%> (?)

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

Files with missing lines	Coverage Δ
sbi/inference/potentials/__init__.py	100.00% <100.00%> (ø)
sbi/inference/potentials/base_potential.py	92.06% <87.50%> (-0.80%)	⬇️
sbi/inference/potentials/protocol.py	94.44% <94.44%> (ø)
sbi/inference/potentials/binding.py	83.33% <83.33%> (ø)
...inference/potentials/likelihood_based_potential.py	78.33% <64.00%> (-3.78%)	⬇️

... and 5 files with indirect coverage changes

[satwiksps]

should I (and maybe Satwik too) merge this to another 'gsoc specific' branch instead of sbi/main? Or use the github functionality "Create draft pull request"?

I was thinking the same thing. For NN Builder API Refactor also there would be 8 PRs. Also my upcoming PR 2 requires my PR 1 to be merged, so I would be creating PR 2 as draft PR. Maybe a GSoC specific branch would be a good thing. Am okay with whatever maintainers prefer.

[janfb]

Thanks, @Jocho-Smith and @satwiksps! just briefly here to unblock you. Yes, it's actually a good idea, if your PRs are introducing breaking changes. We are planning to do a release soon, and we don't want to introduce breaking changes in there.

That said, the PRs that go into that branch should be treated as if they go into main, e.g., fully reviewed by us and functional. We want to avoid the situation where half ready PRs accumulate in a branch with a massive diff to main, that we then have to review again.

So, yes, let's merge your PRs, into an intermediate gsoc-2026 branch, treat it as your main branch, i.e., you both use the same one.

Hope this helps!

[Jocho-Smith]

In this comment I do two things:

• Report on the progress on the execution of the proposal and
• suggest a pivot for this project (containing 2 changes).

Both points are related in a somewhat tricky way, which I try to explain here.

Contribution of this PR

So with some delay, here the minimal working example of the proposed protocol change: docs/notebooks/sbi_old_vs_new_api.ipynb, there one can see the new API in action for NLE.

Some notes and quetsions

1: I have been as explicit about the 'build posterior' step as possible while still being within suggestions of the how_to_guides. For the old API I demonstrate this by calling

potential_fn, parameter_transform = likelihood_estimator_based_potential(
    likelihood_estimator, prior, x_o
)
posterior_old = MCMCPosterior(
    potential_fn, proposal=prior, theta_transform=parameter_transform, warmup_steps=10
)

which is actually equivalent to

posterior = trainer.build_posterior(prior=prior, sample_with='mcmc')

(except for setting x_o before the sampling, more on that see below)

This way of calling the potential function before the sampler in two separate expressions is motivated here: https://github.com/sbi-dev/sbi/blob/main/docs/how_to_guide/09_sampler_interface.ipynb

However, I don't see any situation where users would have an incentive to manually call these potential functions. Even the most finely tuned MCMCPosterior setup I found in the how_to_guides just call .build_posterior, with many parameters, see the last code block of https://github.com/sbi-dev/sbi/blob/main/docs/how_to_guide/24_abstraction_levels.ipynb

• @maintainers is the an appreciated refactoring step for the potential function to restrict the user to .build_posterior? One way to ensure they don't lose options is to rewrite all example from 09_sampler_interface.ipynb to only use .build_posterior. If everything works the same, I think this would simplify the user experience.

2: The bind_observation function asks for an argument called potential with type PotentialFunction, this is a Protocol. I'm not quite sure how to combine this with the partial() function to allow for other samples outside of sbi yet though.

• @maintainers Is this the Protocol based design you envisioned in the beginning?

3: I discovered an unexpected dependence on the .device() property of the potential_fn here https://github.com/sbi-dev/sbi/blob/main/sbi/inference/posteriors/mcmc_posterior.py#L108

This motivated the introduction of PotentialFunctionWithDevice in this PR, however I'm not using it explicitly in this test notebook, only implicitly in the "Lower-Level: Using the Protocol Directly" section for the custom potential I would there.

• @maintainers Is this the right way to deal with this device fallback in this new setup? Or should I merge PotentialFunctionWithDevice and PotentialFunction?

Pivot idea 1: A simple alternative solution

When playing with some examples I made the following observation:

Currently the BasePotential handles x_o in this statefull manner (simplified code block):

class BasePotential:
    def __init__(self, prior, x_o=None, device="cpu"):
        self.device = device
        self.prior = prior
        self.set_x(x_o)

    def set_x(self, x_o, x_is_iid=True):
        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

More precisely:

• set_x() returns nothing, but (over)writes
• self._x_o is not defined in __init__(), instead set_x() is called directly, after 'user-input check' inside process_x()

Usage looks like this then:

potential = BasePotential(prior)

potential.set_x(x_o)

# same object was modified

This set_x() method is actually used in some occasions throughout the code base (I'm working on a document that examines all occurrences of it, WIP)

Assuming that the goal is only to make the potentials stateless, this simple modification would do the trick

class BasePotential:
    def __init__(
        self,
        prior,
        x_o=None,
        device="cpu",
        x_is_iid=True,
    ):
        self.device = device
        self.prior = prior
        self._x_o = x_o
        self._x_is_iid = x_is_iid

    def set_x(self, x_o, x_is_iid=True):
        if x_o is not None:
            x_o = process_x(x_o).to(self.device)

        return self.__class__(
            prior=self.prior,
            x_o=x_o,
            device=self.device,
            x_is_iid=x_is_iid,
        )

So the usage becomes

potential1 = BasePotential(prior)

potential2 = potential1.set_x(x_o)

The obvious consequence is that we cannot 'reuse' the same potential_fn object anymore the way it was possible before, without copying it (which would increase memory consumption and operations needed). The question is, is that happening? From what I see, we have 3 cases here:

• The 'vanilla' sbi workflow: .build_posterior() and .sample() are called once - we just set x_o once. Copying the potential function will not cause problems here.

• Sequential Methods: There we train -> build_postarior -> sample -> train -> build_posterior -> sample -> ... In this case copying will also not slow things down, b/c a completely new posterior is build again anyway.

• Batched sampling: Here I'm unsettled b/c of this line

  sbi/sbi/inference/posteriors/mcmc_posterior.py

  Line 676 in a7f115c

  potential_.set_x(xi)

  • this would very likely slow down computations, if one had to create a new potential function every time (for large x).

• Does anyone see how we could deal with this batched sampling situation? This might be the only thing holding us back from simplifying the potential function API by a lot.

Pivot idea 2: remove set_x from BasePotential.__init__

Last (potentially stupid) question:

• Why do we allow the users to set x_o in 2 ways anyway?

I think there is no point of setting x_o before sampling, and it would be easy to get rid of it - here is why: Currently there is no way of specifying x_o when calling build_potential!

Higher up in this post I argue that build_potential should be the only way to build the potential, so running

posterior = trainer.build_posterior(prior=prior, sample_with='mcmc')
samples = posterior.sample((1000,), x=x_o)

instead of

potential_fn, parameter_transform = likelihood_estimator_based_potential(
    likelihood_estimator, prior, None
)
posterior_old = MCMCPosterior(
    potential_fn, proposal=prior, theta_transform=parameter_transform, warmup_steps=10
)
samples_old = posterior_old.sample((1000,), x=x_o)

Notice, that this second block could have been rewritten like this:

potential_fn, parameter_transform = likelihood_estimator_based_potential(
    likelihood_estimator, prior, x_o
)
posterior_old = MCMCPosterior(
    potential_fn, proposal=prior, theta_transform=parameter_transform, warmup_steps=10
)
samples_old = posterior_old.sample((1000,))

So we can set x_o in the setup step of the potential directly, even though we cannot do this through .build_posterior() - this way of calling likelihood_estimator_based_potential(likelihood_estimator, prior, x_o) is implemented and leads to the same results. However, it is the less 'natural' way of using the sbi code base already, b/c of the following observation:

• When posteriors require potential_fn and theta_transform we get into this case distinction https://github.com/sbi-dev/sbi/blob/main/sbi/inference/trainers/base.py#L932
• this always runs self._get_potential_function depending on the chosen method (NPE,NLE,NRE). In the case of NLE this is exactly where likelihood_estimator_based_potential(likelihood_estimator, prior, x_o) is called and (and this is the punchline) the value of x_o for all methods in all cases is set to None:
  • NLE: https://github.com/sbi-dev/sbi/blob/main/sbi/inference/trainers/nle/nle_base.py#L326
  • NPE: https://github.com/sbi-dev/sbi/blob/main/sbi/inference/trainers/npe/npe_base.py#L479
  • NRE: https://github.com/sbi-dev/sbi/blob/main/sbi/inference/trainers/nre/nre_base.py#L415

So one could add the possibility to provide x_o to .build_potential or remove this option entirely and make it mandatory to provide it to sample(), I advocate for the later.

I haven't fully analyzed it yet (for example I didn't look VectorFields now), but I plan to just implement and test this in a separate branch to see if this works without breaking any tests.

The biggest motivation for this refined approach is that one might achieve the desired stateless potential function design without introducing more complexity, but reducing it, which seems to be the more elegant solution then what we proposed initially.

Sorry if this post is a bit clumsy to read, this is already my third attempt to communicate my observations and ideas and I was unhappy with it every time and still am, but since this caused a one week delay of this PR#2 already, I decided that 'released/done is better then perfect' :)

[dgedon]

Overall, I think the PotentialFunction(theta, x) protocol and bind_observation are the right building blocks. Below are some design decisions we should agree on. I also found a currently undocumented case, and a few concrete issues to fix.

Main idea

Potentials should carry no x_o state at all. Not at construction and not via set_x(). The only available mechanism to hand over x_o is a function bind_observation(raw_potential, x_o), which produces an immutable BoundPotential. This call happens internally the moment sampling or evaluation starts. No potential constructor should accept x_o.

How x_o reaches the potential: three paths

There are three ways x_o ends up in the potential evaluation. All three should go through bind_observation, not through set_x().

Path 1: explicit at call time: posterior.sample(x=x_o) or posterior.log_prob(theta, x=x_o)

Path 2: pre-configured default: posterior.set_default_x(x_o). This one is important because it is the only option when the posterior is passed to external code that calls log_prob(theta) without knowing about x_o. One clear example is TSNPE:

posterior1 = DirectPosterior(...).set_default_x(x_o)
accept_reject_fn = get_density_thresholder(posterior1, quantile=1e-4)
proposal = RestrictedPrior(prior, accept_reject_fn, posterior=posterior1, sample_with="sir")

get_density_thresholder calls dist.log_prob(theta) with no x_o argument. RestrictedPrior does the same. There is no way to pass x_o through sample(x=...) here, so set_default_x is necessary and must be kept.

Path 3: what we implement for an expert user: bind_observation(raw, x_o) directly, without a posterior object, for users who configure their own samplers in all detail.

Undocumented case: Specific Samplers in TSNPE

When ImportanceSamplingPosterior (or others??) is used as the proposal inside RestrictedPrior with SIR sampling, the call chain goes (see here):

RestrictedPrior.sample(sir)
  -> ImportanceSamplingPosterior.log_prob(theta, x=None)
       x = self._x_else_default_x(None) 
       self.potential_fn.set_x(x)    
       self.potential_fn(theta)
       estimate_normalization_constant(x)
         -> importance_sample(self.potential_fn, ...)
         -> self.potential_fn(samples)

Under the new design this should be:

def log_prob(self, theta, x=None, ...):
    x = self._x_else_default_x(x)
    bound = bind_observation(self._raw_potential, x)
    potential_values = bound(theta)
    normalization = self.estimate_normalization_constant(x, bound)
    ...

The same refactor applies to MCMCPosterior, RejectionPosterior, VIPosterior, and VectorFieldPosterior, all of which currently call self.potential_fn.set_x(x) at the top of their sample() or log_prob() methods.

What we should keep based on my investigation and our discussion today

set_default_x on NeuralPosterior must stay for three reasons: SNPE multi-round training uses it to configure the proposal (npe_a/b/c), ImportanceSamplingPosterior uses it for normalization constant caching, and the TSNPE pattern described above depends on it entirely.

Issues in this PR (DG: IMPORTANTLY this is mostly claude generated, take it as an inspiration for now since we are discussing the design so far)

1. BoundPotential.set_x() is a silent no-op (blocking)

When ImportanceSamplingPosterior.log_prob calls self.potential_fn.set_x(self.default_x) on a BoundPotential, the call is silently discarded and the original construction-time x_o is used instead. This can produce wrong results without any error. It should raise NotImplementedError with a migration message pointing to bind_observation.

2. IID prior double-counting (correctness bug)

In the x_batch_size > 1 branch of likelihood_potential(), the prior log-probability is broadcast to shape (n_iid, n_theta) and returned as part of the result. BoundPotential then sums over dim=0, accumulating n_iid copies of the prior. The correct formula adds the prior once, after summing the likelihoods. The fix is to return only the likelihood from likelihood_potential in the multi-observation branch and add the prior after aggregation.

3. CustomPotentialWrapper runtime introspection is fragile

sig = inspect.signature(self.potential_fn)
num_params = len(sig.parameters)
if num_params == 1:
    return self.potential_fn(theta)
else:
    return self.potential_fn(theta, self.x_o)

This breaks for functools.partial, **kwargs signatures, and callables where self is counted. Callers who pass a single-argument callable should wrap it in BoundPotential explicitly before passing it in.

4. BoundPotential bypasses BasePotential.__init__()

BoundPotential inherits from BasePotential but skips super().__init__(), so process_x validation, the prior attribute, and to() device handling are all missing. isinstance(bound, BasePotential) returns True but the contract is not actually satisfied. Prefer composition over inheritance here, or call super().__init__(prior=None, x_o=x_o, device=x_o.device).

5. VectorFieldBasedPotential cannot be naively wrapped

VectorFieldBasedPotential.set_x() rebuilds the ODE flow on every call. Wrapping it in BoundPotential means the flow is never rebuilt, which silently breaks results. This needs an explicit decision: either exclude VectorFieldBasedPotential from this refactor for now (and document it), or trigger the flow rebuild at bind_observation time.