From 7601c08f5feb8c7b93513f06d44f7969604420e2 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Thu, 4 Jun 2026 20:34:46 +0200 Subject: [PATCH 1/9] PR #1, Weeks 3-4: implement core protocol + bind_observation 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 - Add validate_potential() helper for runtime protocol validation 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 clear 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 - Backward compatibility with existing LikelihoodBasedPotential This enables: - Stateless potential functions without hidden state - Easy integration with PyMC, and custom PyTorch samplers - Clear device handling with explicit error messages - Backward compatibility: existing potentials continue to work --- sbi/inference/potentials/__init__.py | 16 ++ sbi/inference/potentials/binding.py | 144 +++++++++++++ sbi/inference/potentials/protocol.py | 106 ++++++++++ tests/potential_protocol_test.py | 303 +++++++++++++++++++++++++++ 4 files changed, 569 insertions(+) create mode 100644 sbi/inference/potentials/binding.py create mode 100644 sbi/inference/potentials/protocol.py create mode 100644 tests/potential_protocol_test.py diff --git a/sbi/inference/potentials/__init__.py b/sbi/inference/potentials/__init__.py index c88d6301c..a03bf4f83 100644 --- a/sbi/inference/potentials/__init__.py +++ b/sbi/inference/potentials/__init__.py @@ -1,6 +1,11 @@ # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed # under the Apache License Version 2.0, see +from sbi.inference.potentials.binding import ( + BoundPotential, + bind_observation, + bind_observation_class, +) from sbi.inference.potentials.likelihood_based_potential import ( likelihood_estimator_based_potential, mixed_likelihood_estimator_based_potential, @@ -8,6 +13,11 @@ from sbi.inference.potentials.posterior_based_potential import ( posterior_estimator_based_potential, ) +from sbi.inference.potentials.protocol import ( + PotentialFunction, + PotentialFunctionWithDevice, + validate_potential, +) from sbi.inference.potentials.ratio_based_potential import ( ratio_estimator_based_potential, ) @@ -16,9 +26,15 @@ ) __all__ = [ + "bind_observation", + "bind_observation_class", + "BoundPotential", "likelihood_estimator_based_potential", "mixed_likelihood_estimator_based_potential", "posterior_estimator_based_potential", + "PotentialFunction", + "PotentialFunctionWithDevice", "ratio_estimator_based_potential", + "validate_potential", "vector_field_estimator_based_potential", ] diff --git a/sbi/inference/potentials/binding.py b/sbi/inference/potentials/binding.py new file mode 100644 index 000000000..6d49e5675 --- /dev/null +++ b/sbi/inference/potentials/binding.py @@ -0,0 +1,144 @@ +# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed +# under the Apache License Version 2.0, see + +"""Binding utilities for creating sampler-ready potential functions. + +This module provides Layer 2 utilities for binding observations to stateless +potential functions, making them ready for MCMC/VI samplers. +""" + +from typing import Callable, Union + +import torch +from torch import Tensor + +from sbi.inference.potentials.protocol import PotentialFunction + + +def bind_observation( + potential: PotentialFunction, + x_o: Tensor, + sum_iid: bool = True, +) -> Callable[[Tensor], Tensor]: + """Create a sampler-compatible theta -> log_prob function. + + This utility binds observed data to a stateless potential function, producing + a callable ready for use with MCMC, VI, or rejection sampling algorithms. + + The function infers the device from x_o and validates compatibility with + the potential's device requirements (if the potential has a device attribute). + + Args: + potential: Stateless (theta, x) -> log_prob function satisfying the + PotentialFunction protocol. + x_o: Observed data tensor. Shape should be (obs_batch, *obs_shape) for + IID observations, or (*obs_shape) for single observation. + sum_iid: If True, sum log probabilities over the observation batch dimension. + This is the typical behavior for IID observations where we want to + compute log p(x_o | theta) = sum_i log p(x_i | theta). + + Returns: + A callable that takes theta (parameters) and returns the log-probability + log p(theta | x_o) + log p(theta). The returned function is stateless + from the sampler's perspective. + + Example: + >>> # Create stateless potential + >>> def likelihood_potential(theta: Tensor, x: Tensor) -> Tensor: + ... return estimator.log_prob(x, context=theta) + prior.log_prob(theta) + >>> + >>> # Bind observation for sampling + >>> bound_fn = bind_observation(likelihood_potential, x_o, sum_iid=True) + >>> log_probs = bound_fn(theta_samples) # Ready for MCMC/VI + """ + device = x_o.device + + if hasattr(potential, "device"): + potential_device = potential.device + if isinstance(potential_device, str): + potential_device = torch.device(potential_device) + if potential_device != device: + raise ValueError( + f"Device mismatch: x_o is on {device}, but potential expects " + f"{potential_device}. Use x_o.to('{potential_device}') before " + f"binding, or ensure the potential was created on the same device " + f"as x_o." + ) + + x_o = x_o.to(device) + + def bound_potential(theta: Tensor) -> Tensor: + theta = theta.to(device) + + log_prob = potential(theta, x_o) + + if sum_iid and x_o.ndim > 1: + log_prob = log_prob.sum(dim=0) + elif sum_iid and x_o.ndim == 0: + pass + + return log_prob + + return bound_potential + + +class BoundPotential: + """Wrapper class for binding observation to a potential function. + + This provides an alternative to the functional bind_observation, maintaining + compatibility with existing code patterns while supporting the new stateless + protocol internally. + """ + + def __init__( + self, + potential: PotentialFunction, + x_o: Tensor, + sum_iid: bool = True, + ): + self._potential = potential + self._x_o = x_o.to(x_o.device) + self._sum_iid = sum_iid + self._device = x_o.device + + @property + def device(self) -> torch.device: + return self._device + + def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: + theta = theta.to(self._device) + + with torch.set_grad_enabled(track_gradients): + log_prob = self._potential(theta, self._x_o) + + if self._sum_iid and self._x_o.ndim > 1: + log_prob = log_prob.sum(dim=0) + + return log_prob + + def to(self, device: Union[str, torch.device]) -> "BoundPotential": + self._device = torch.device(device) + self._x_o = self._x_o.to(device) + return self + + +def bind_observation_class( + potential: PotentialFunction, + x_o: Tensor, + sum_iid: bool = True, +) -> BoundPotential: + """Create a BoundPotential instance for sampler-compatible usage. + + This is an alternative to bind_observation that returns a class instance + instead of a simple function. Useful when you need method chaining or + want to maintain compatibility with code that expects an object. + + Args: + potential: Stateless (theta, x) -> log_prob function. + x_o: Observed data tensor. + sum_iid: If True, sum log probabilities over IID batch dimension. + + Returns: + BoundPotential instance that can be called like a function. + """ + return BoundPotential(potential, x_o, sum_iid) diff --git a/sbi/inference/potentials/protocol.py b/sbi/inference/potentials/protocol.py new file mode 100644 index 000000000..39d129d0f --- /dev/null +++ b/sbi/inference/potentials/protocol.py @@ -0,0 +1,106 @@ +# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed +# under the Apache License Version 2.0, see + +"""Protocol definitions for stateless potential functions. + +This module provides the core protocol (Layer 1) for potential functions, +enabling stateless, composable potential evaluation. +""" + +from typing import Protocol + +import torch +from torch import Tensor + + +class PotentialFunction(Protocol): + """Stateless callable protocol for potential functions. + + A PotentialFunction takes parameters theta and observed data x, and returns + the unnormalized log-probability log p(theta, x). This is the core interface + that all potential functions must satisfy. + + The protocol is stateless - all inputs (theta, x) are passed as arguments, + with no hidden state dependencies. This enables: + - Easy testing (same inputs -> same outputs) + - Simple serialization + - Composable inference workflows + - Integration with external samplers (NumPyro, PyMC, custom PyTorch) + + Example: + >>> def my_potential(theta: Tensor, x: Tensor) -> Tensor: + ... ll = density_estimator.log_prob(x, context=theta) + ... return ll + prior.log_prob(theta) + >>> # Satisfies PotentialFunction protocol + """ + + def __call__(self, theta: Tensor, x: Tensor) -> Tensor: + """Evaluate the potential function at theta given observed data x. + + Args: + theta: Parameter samples with shape (batch_size, *param_shape). + x: Observed data with shape (obs_batch, *obs_shape). + + Returns: + Log-probability with shape (batch_size,) or (obs_batch, batch_size) + depending on the potential implementation. + """ + ... + + +class PotentialFunctionWithDevice(Protocol): + """Extended protocol that reports device requirements. + + This protocol adds a device property to enable automatic device validation + and inference in bind_observation. + """ + + def __call__(self, theta: Tensor, x: Tensor) -> Tensor: + """Evaluate the potential function at theta given observed data x.""" + ... + + @property + def device(self) -> torch.device: + """Return the expected device for inputs. + + This property allows bind_observation to validate device compatibility + and provide clear error messages for device mismatches. + """ + ... + + +def validate_potential(potential: PotentialFunction) -> bool: + """Validate that a potential function satisfies the protocol. + + This helper checks that the potential is callable and has the correct + signature. It performs runtime validation since Protocol typing + doesn't catch errors at definition time. + + Args: + potential: The potential function to validate. + + Returns: + True if the potential satisfies the protocol. + + Raises: + TypeError: If potential doesn't satisfy the protocol. + ValueError: If potential is callable but has wrong signature. + """ + if not callable(potential): + raise TypeError( + "Potential must be callable. " + "Expected a function or object with __call__ method." + ) + + import inspect + + sig = inspect.signature(potential.__call__) + params = list(sig.parameters.keys()) + + if len(params) < 2: + raise ValueError( + f"Potential __call__ must have at least 2 parameters (theta, x), " + f"got signature: {sig}" + ) + + return True diff --git a/tests/potential_protocol_test.py b/tests/potential_protocol_test.py new file mode 100644 index 000000000..aa4f68e6f --- /dev/null +++ b/tests/potential_protocol_test.py @@ -0,0 +1,303 @@ +# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed +# under the Apache License Version 2.0, see + +"""Unit tests for the stateless PotentialFunction protocol and bind_observation. + +These tests verify: +1. Protocol compliance for potential functions +2. bind_observation correctly binds x_o +3. IID handling (sum_iid parameter) +4. Device validation +5. Statelessness (repeated calls yield identical outputs) +""" + +import pytest +import torch +from torch import Tensor + +from sbi.inference.potentials.binding import BoundPotential, bind_observation +from sbi.inference.potentials.protocol import ( + validate_potential, +) + + +class SimplePotential: + """Simple potential function for testing.""" + + def __init__(self): + self.device = torch.device("cpu") + + def __call__(self, theta: Tensor, x: Tensor) -> Tensor: + return -0.5 * ((theta - x.mean()).pow(2)).sum(dim=-1) + + +class DeviceAwarePotential: + """Potential with explicit device attribute.""" + + def __init__(self, device: torch.device): + self.device = device + + def __call__(self, theta: Tensor, x: Tensor) -> Tensor: + return -0.5 * ((theta - x.mean()).pow(2)).sum(dim=-1) + + +def stateless_potential(theta: Tensor, x: Tensor) -> Tensor: + """Simple stateless potential function.""" + return -0.5 * ((theta - x.mean()).pow(2)).sum(dim=-1) + + +class TestPotentialFunctionProtocol: + """Tests for PotentialFunction protocol.""" + + def test_validate_callable_function(self): + """Test that validate_potential accepts valid functions.""" + assert validate_potential(stateless_potential) is True + + def test_validate_callable_class(self): + """Test that validate_potential accepts callable classes.""" + assert validate_potential(SimplePotential()) is True + + def test_validate_non_callable_raises(self): + """Test that validate_potential rejects non-callables.""" + with pytest.raises(TypeError, match="must be callable"): + validate_potential("not a potential") + + def test_validate_missing_call_raises(self): + """Test that validate_potential catches missing __call__.""" + + class NoCall: + pass + + with pytest.raises(TypeError, match="must be callable"): + validate_potential(NoCall()) + + +class TestBindObservation: + """Tests for bind_observation utility.""" + + def test_bind_simple_potential(self): + """Test binding a simple stateless potential.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta = torch.tensor([[0.0], [1.0], [2.0]]) + result = bound_fn(theta) + + assert result.shape == (3,) + assert torch.isfinite(result).all() + + def test_bind_preserves_device(self): + """Test that bound function works on correct device.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta = torch.tensor([[0.0], [1.0]]) + result = bound_fn(theta) + + assert result.device == x_o.device + + def test_iid_summing(self): + """Test that sum_iid=True correctly sums over IID observations.""" + x_o = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + + def potential(theta: Tensor, x: Tensor) -> Tensor: + return -((theta.unsqueeze(0) - x.unsqueeze(1)).pow(2)).sum(dim=-1) + + bound_fn_sum = bind_observation(potential, x_o, sum_iid=True) + bound_fn_no_sum = bind_observation(potential, x_o, sum_iid=False) + + theta = torch.tensor([[0.0, 0.0], [1.0, 1.0]]) + + result_sum = bound_fn_sum(theta) + result_no_sum = bound_fn_no_sum(theta) + + assert result_sum.shape == (2,) + assert result_no_sum.shape == (2, 2) + + def test_single_observation(self): + """Test binding with single observation (no IID dimension).""" + x_o = torch.tensor([1.0, 2.0]) + + def potential(theta: Tensor, x: Tensor) -> Tensor: + return -((theta - x).pow(2)).sum(dim=-1) + + bound_fn = bind_observation(potential, x_o, sum_iid=True) + theta = torch.tensor([[0.0, 0.0]]) + result = bound_fn(theta) + + assert result.shape == (1,) + assert torch.isfinite(result).all() + + def test_device_mismatch_raises(self): + """Test that device mismatch raises clear error.""" + potential = DeviceAwarePotential(torch.device("cpu")) + x_o = torch.tensor([1.0, 2.0, 3.0]) + + if torch.cuda.is_available(): + x_o_cuda = x_o.to("cuda") + with pytest.raises(ValueError, match="Device mismatch"): + bind_observation(potential, x_o_cuda, sum_iid=True) + + +class TestBoundPotential: + """Tests for BoundPotential class.""" + + def test_bound_potential_callable(self): + """Test that BoundPotential is callable.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound = BoundPotential(stateless_potential, x_o, sum_iid=True) + + theta = torch.tensor([[0.0], [1.0], [2.0]]) + result = bound(theta) + + assert result.shape == (3,) + assert torch.isfinite(result).all() + + def test_bound_potential_device_property(self): + """Test device property returns correct device.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound = BoundPotential(stateless_potential, x_o, sum_iid=True) + + assert bound.device == x_o.device + + def test_bound_potential_to_method(self): + """Test .to() method changes device.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound = BoundPotential(stateless_potential, x_o, sum_iid=True) + + if torch.cuda.is_available(): + bound_cuda = bound.to("cuda") + assert bound_cuda.device.type == "cuda" + + +class TestStatelessness: + """Tests verifying stateless behavior - key property of the new design.""" + + def test_repeated_calls_same_output(self): + """Test that repeated calls with same inputs yield identical outputs.""" + x_o = torch.tensor([1.0, 2.0, 3.0, 4.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta = torch.tensor([[0.0], [1.0], [2.0]]) + + result1 = bound_fn(theta.clone()) + result2 = bound_fn(theta.clone()) + result3 = bound_fn(theta.clone()) + + assert torch.allclose(result1, result2) + assert torch.allclose(result2, result3) + + def test_different_theta_different_output(self): + """Test that different theta values yield different outputs.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta1 = torch.tensor([[0.0]]) + theta2 = torch.tensor([[1.0]]) + theta3 = torch.tensor([[2.0]]) + + result1 = bound_fn(theta1) + result2 = bound_fn(theta2) + result3 = bound_fn(theta3) + + assert not torch.allclose(result1, result2) + assert not torch.allclose(result2, result3) + + +class TestIntegration: + """Integration tests with realistic potential functions.""" + + def test_with_prior_and_likelihood(self): + """Test binding a potential combining prior and likelihood.""" + prior = torch.distributions.Normal(0.0, 1.0) + + class MockEstimator: + def log_prob(self, x: Tensor, context: Tensor) -> Tensor: + return -((x - context).pow(2)).sum(dim=-1) + + estimator = MockEstimator() + + def potential(theta: Tensor, x: Tensor) -> Tensor: + likelihood = estimator.log_prob(x, context=theta) + prior_log_prob = prior.log_prob(theta).sum(dim=-1) + return likelihood + prior_log_prob + + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(potential, x_o, sum_iid=True) + + theta = torch.randn(10, 1) + result = bound_fn(theta) + + assert result.shape == (10,) + assert torch.isfinite(result).all() + + def test_with_gradient_tracking(self): + """Test that gradients can be tracked through bound potential.""" + x_o = torch.tensor([1.0, 2.0, 3.0]) + bound_fn = bind_observation(stateless_potential, x_o, sum_iid=True) + + theta = torch.randn(5, 1, requires_grad=True) + result = bound_fn(theta) + + assert result.requires_grad + result.sum().backward() + assert theta.grad is not None + + +class TestBackwardCompatibility: + """Tests verifying backward compatibility with existing sbi workflows.""" + + def test_existing_likelihood_potential_still_works(self): + """Test that existing likelihood_estimator_based_potential works.""" + from sbi.inference import NLE + from sbi.inference.potentials.likelihood_based_potential import ( + likelihood_estimator_based_potential, + ) + from sbi.utils import BoxUniform + + torch.manual_seed(42) + + prior = BoxUniform(low=torch.zeros(3), high=torch.ones(3)) + theta = prior.sample((50,)) + x = theta + torch.randn_like(theta) * 0.1 + x_o = torch.randn(1, 3) + + trainer = NLE() + likelihood_estimator = trainer.append_simulations(theta, x).train( + training_batch_size=16, max_num_epochs=2 + ) + + potential_fn, parameter_transform = likelihood_estimator_based_potential( + likelihood_estimator, prior, x_o + ) + + assert hasattr(potential_fn, "x_o") + assert hasattr(potential_fn, "set_x") + + theta_test = torch.randn(5, 3) + result = potential_fn(theta_test) + assert result.shape == (5,) + + def test_validate_potential_with_real_potential(self): + """Test that validate_potential works with existing potentials.""" + from sbi.inference import NLE + from sbi.inference.potentials.likelihood_based_potential import ( + LikelihoodBasedPotential, + ) + from sbi.utils import BoxUniform + + torch.manual_seed(42) + + prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2)) + theta = prior.sample((30,)) + x = theta + torch.randn_like(theta) * 0.1 + + trainer = NLE() + likelihood_estimator = trainer.append_simulations(theta, x).train( + training_batch_size=16, max_num_epochs=2 + ) + + potential = LikelihoodBasedPotential(likelihood_estimator, prior) + potential.set_x(torch.randn(1, 2)) + + assert validate_potential(potential) is True From 2ece7d9589c8086799bfb62010f6c72a29519fab Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 15 Jun 2026 11:09:22 +0200 Subject: [PATCH 2/9] build(deps): bump codecov/codecov-action from 6.0.1 to 7.0.0 (#1880) Bumps [codecov/codecov-action](https://github.com/codecov/codecov-action) from 6.0.1 to 7.0.0. - [Release notes](https://github.com/codecov/codecov-action/releases) - [Changelog](https://github.com/codecov/codecov-action/blob/main/CHANGELOG.md) - [Commits](https://github.com/codecov/codecov-action/compare/e79a6962e0d4c0c17b229090214935d2e33f8354...fb8b3582c8e4def4969c97caa2f19720cb33a72f) --- updated-dependencies: - dependency-name: codecov/codecov-action dependency-version: 7.0.0 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/cd.yml | 4 ++-- .github/workflows/ci.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 8d181012b..ede6ee394 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -51,7 +51,7 @@ jobs: run: uv run pytest -v -x -n auto -m "not gpu" --cov=sbi --cov-report=xml --junitxml=junit.xml -o junit_family=legacy tests/ - name: Upload coverage to Codecov - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: env_vars: OS,PYTHON files: ./coverage.xml @@ -62,7 +62,7 @@ jobs: - name: Upload test results to Codecov if: ${{ !cancelled() }} - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: report_type: test_results env: diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c937915ba..43054217c 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -64,7 +64,7 @@ jobs: run: uv run pytest -v -n auto -m "not slow and not gpu" --cov=sbi --cov-report=xml --junitxml=junit.xml -o junit_family=legacy --splitting-algorithm least_duration --splits ${{ matrix.split_size }} --group ${{ matrix.group_number }} tests/ - name: Upload coverage to Codecov - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: env_vars: OS,PYTHON file: ./coverage.xml @@ -75,7 +75,7 @@ jobs: - name: Upload test results to Codecov if: ${{ !cancelled() }} - uses: codecov/codecov-action@e79a6962e0d4c0c17b229090214935d2e33f8354 # v6.0.1 + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 with: report_type: test_results env: From 6660ea0dfa2787bcb8f686693725575454e83843 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 15 Jun 2026 11:09:50 +0200 Subject: [PATCH 3/9] build(deps): bump sigstore/gh-action-sigstore-python from 3.3.0 to 3.4.0 (#1881) Bumps [sigstore/gh-action-sigstore-python](https://github.com/sigstore/gh-action-sigstore-python) from 3.3.0 to 3.4.0. - [Release notes](https://github.com/sigstore/gh-action-sigstore-python/releases) - [Changelog](https://github.com/sigstore/gh-action-sigstore-python/blob/main/CHANGELOG.md) - [Commits](https://github.com/sigstore/gh-action-sigstore-python/compare/04cffa1d795717b140764e8b640de88853c92acc...5b79a39c381910c090341a2c9b0bf022c8b387e1) --- updated-dependencies: - dependency-name: sigstore/gh-action-sigstore-python dependency-version: 3.4.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> --- .github/workflows/publish.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index fd27f72ba..9ae4fe341 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -77,7 +77,7 @@ jobs: name: python-package-distributions path: dist/ - name: Sign the dists with Sigstore - uses: sigstore/gh-action-sigstore-python@04cffa1d795717b140764e8b640de88853c92acc # v3.3.0 + uses: sigstore/gh-action-sigstore-python@5b79a39c381910c090341a2c9b0bf022c8b387e1 # v3.4.0 with: inputs: >- ./dist/*.tar.gz From 89b9e76718b62e75ebbd54aa3bcffc95d7ee5192 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Tue, 16 Jun 2026 08:37:23 +0200 Subject: [PATCH 4/9] prepare minimal notebook test --- sbi/inference/potentials/__init__.py | 6 ++ .../potentials/likelihood_based_potential.py | 67 +++++++++++++++++++ tests/potential_protocol_test.py | 54 +++++++++++++++ 3 files changed, 127 insertions(+) diff --git a/sbi/inference/potentials/__init__.py b/sbi/inference/potentials/__init__.py index a03bf4f83..8ed377efd 100644 --- a/sbi/inference/potentials/__init__.py +++ b/sbi/inference/potentials/__init__.py @@ -8,10 +8,12 @@ ) from sbi.inference.potentials.likelihood_based_potential import ( likelihood_estimator_based_potential, + likelihood_potential, mixed_likelihood_estimator_based_potential, ) from sbi.inference.potentials.posterior_based_potential import ( posterior_estimator_based_potential, + posterior_potential, ) from sbi.inference.potentials.protocol import ( PotentialFunction, @@ -20,6 +22,7 @@ ) from sbi.inference.potentials.ratio_based_potential import ( ratio_estimator_based_potential, + ratio_potential, ) from sbi.inference.potentials.vector_field_potential import ( vector_field_estimator_based_potential, @@ -30,11 +33,14 @@ "bind_observation_class", "BoundPotential", "likelihood_estimator_based_potential", + "likelihood_potential", "mixed_likelihood_estimator_based_potential", "posterior_estimator_based_potential", + "posterior_potential", "PotentialFunction", "PotentialFunctionWithDevice", "ratio_estimator_based_potential", + "ratio_potential", "validate_potential", "vector_field_estimator_based_potential", ] diff --git a/sbi/inference/potentials/likelihood_based_potential.py b/sbi/inference/potentials/likelihood_based_potential.py index d2907430c..411742e61 100644 --- a/sbi/inference/potentials/likelihood_based_potential.py +++ b/sbi/inference/potentials/likelihood_based_potential.py @@ -21,6 +21,73 @@ from sbi.utils.sbiutils import mcmc_transform +def likelihood_potential( + likelihood_estimator: ConditionalDensityEstimator, + prior: Distribution, +) -> Callable[[Tensor, Tensor], Tensor]: + r"""Create a stateless likelihood potential function. + + This returns a function that satisfies the PotentialFunction protocol: + (theta, x) -> log p(x|theta) + log p(theta) + + This is the new recommended API. For use with samplers, combine with + bind_observation(). + + Args: + likelihood_estimator: The density estimator modelling the likelihood. + prior: The prior distribution. + + Returns: + A callable (theta, x) -> log_prob that computes the likelihood potential. + The returned function has a `device` attribute for compatibility with + bind_observation. + """ + device = str(next(likelihood_estimator.parameters()).device) + + def _potential(theta: Tensor, x: Tensor) -> Tensor: + theta = theta.to(device) + x = x.to(device) + + x_batch_size = x.shape[0] + theta_batch_size = theta.shape[0] + + with torch.set_grad_enabled(True): + if x_batch_size == 1: + x_for_eval = x + log_likelihood = likelihood_estimator.log_prob( + x_for_eval, condition=theta + ) + log_likelihood = log_likelihood.squeeze(1) + elif x_batch_size == theta_batch_size: + log_likelihood = likelihood_estimator.log_prob(x, condition=theta) + log_likelihood = log_likelihood.reshape(-1) + else: + x_expanded = reshape_to_sample_batch_event( + x, event_shape=x.shape[1:], leading_is_sample=True + ) + trailing_minus_ones = [-1 for _ in range(x_expanded.dim() - 2)] + x_expanded = x_expanded.expand( + -1, theta_batch_size, *trailing_minus_ones + ) + + theta_batch = reshape_to_batch_event(theta, event_shape=theta.shape[1:]) + + log_likelihood_batch = likelihood_estimator.log_prob( + x_expanded, condition=theta_batch + ) + log_likelihood = log_likelihood_batch.reshape( + x_batch_size, theta_batch_size + ) + + prior_log_prob = prior.log_prob(theta) + + return log_likelihood + prior_log_prob + + _potential.device = torch.device(device) + + return _potential + + def likelihood_estimator_based_potential( likelihood_estimator: ConditionalDensityEstimator, prior: Distribution, # type: ignore diff --git a/tests/potential_protocol_test.py b/tests/potential_protocol_test.py index aa4f68e6f..6c125998b 100644 --- a/tests/potential_protocol_test.py +++ b/tests/potential_protocol_test.py @@ -244,6 +244,60 @@ def test_with_gradient_tracking(self): assert theta.grad is not None +class TestNewStatelessAPI: + """Tests for the new stateless potential functions (likelihood_potential, etc.).""" + + def test_likelihood_potential_creates_valid_function(self): + """Test that likelihood_potential returns a valid PotentialFunction.""" + from sbi.inference import NLE + from sbi.inference.potentials import bind_observation, likelihood_potential + from sbi.utils import BoxUniform + + torch.manual_seed(42) + + prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2)) + theta = prior.sample((50,)) + x = theta + torch.randn_like(theta) * 0.1 + x_o = torch.randn(1, 2) + + trainer = NLE() + estimator = trainer.append_simulations(theta, x).train( + training_batch_size=16, max_num_epochs=2 + ) + + stateless_fn = likelihood_potential(estimator, prior) + + assert hasattr(stateless_fn, "device") + assert callable(stateless_fn) + + bound_fn = bind_observation(stateless_fn, x_o) + theta_test = torch.rand(5, 2) + result = bound_fn(theta_test) + + assert result.shape == (5,) + assert torch.isfinite(result).all() + + def test_new_stateless_functions_satisfy_protocol(self): + """Test that new stateless functions satisfy PotentialFunction protocol.""" + from sbi.inference import NLE + from sbi.inference.potentials import likelihood_potential, validate_potential + from sbi.utils import BoxUniform + + torch.manual_seed(42) + + prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2)) + theta = prior.sample((30,)) + x = theta + torch.randn_like(theta) * 0.1 + + trainer = NLE() + estimator = trainer.append_simulations(theta, x).train( + training_batch_size=16, max_num_epochs=2 + ) + + stateless_fn = likelihood_potential(estimator, prior) + assert validate_potential(stateless_fn) is True + + class TestBackwardCompatibility: """Tests verifying backward compatibility with existing sbi workflows.""" From b67f0596194f1fdf1ff2f125aa8fec099d34dd4d Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Tue, 16 Jun 2026 08:55:15 +0200 Subject: [PATCH 5/9] added notebook and fixed old __init__ --- docs/notebooks/sbi_old_vs_new_api.ipynb | 274 ++++++++++++++++++++++++ sbi/inference/potentials/__init__.py | 4 - 2 files changed, 274 insertions(+), 4 deletions(-) create mode 100644 docs/notebooks/sbi_old_vs_new_api.ipynb diff --git a/docs/notebooks/sbi_old_vs_new_api.ipynb b/docs/notebooks/sbi_old_vs_new_api.ipynb new file mode 100644 index 000000000..5a8e147cc --- /dev/null +++ b/docs/notebooks/sbi_old_vs_new_api.ipynb @@ -0,0 +1,274 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Old API vs New API Demo" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/schmidt/gsoc2026/sbi/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " Neural network successfully converged after 77 epochs." + ] + } + ], + "source": [ + "import warnings\n", + "\n", + "import torch\n", + "\n", + "from sbi.inference import NLE, MCMCPosterior\n", + "from sbi.utils import BoxUniform\n", + "\n", + "warnings.filterwarnings('default')\n", + "\n", + "\n", + "# Setup\n", + "prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2))\n", + "def simulator(θ):\n", + " return θ + torch.randn_like(θ) * 0.1\n", + "\n", + "torch.manual_seed(42)\n", + "theta_train = prior.sample((1000,))\n", + "x_train = simulator(theta_train)\n", + "x_o = torch.randn(1, 2)\n", + "\n", + "# Train\n", + "trainer = NLE()\n", + "likelihood_estimator = trainer.append_simulations(theta_train, x_train).train()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## OLD API (Deprecated)\n", + "\n", + "The old API uses stateful classes that bind `x_o` at construction time." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/schmidt/gsoc2026/sbi/sbi/inference/potentials/likelihood_based_potential.py:118: DeprecationWarning: LikelihoodBasedPotential is deprecated and will be removed in a future release. Use the stateless `likelihood_potential()` function instead, combined with `bind_observation()` for sampler compatibility.\n", + " potential_fn = LikelihoodBasedPotential(\n", + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 46.64it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:04<00:00, 533.34it/s] " + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "OLD API - samples mean: tensor([0.9758, 0.0908])\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "# OLD API: stateful, emits DeprecationWarning\n", + "from sbi.inference.potentials import likelihood_estimator_based_potential\n", + "\n", + "potential_fn, parameter_transform = likelihood_estimator_based_potential(\n", + " likelihood_estimator, prior, x_o\n", + ")\n", + "posterior_old = MCMCPosterior(\n", + " potential_fn, proposal=prior, theta_transform=parameter_transform, warmup_steps=10\n", + ")\n", + "samples_old = posterior_old.sample((1000,))\n", + "print(f\"OLD API - samples mean: {samples_old.mean(dim=0)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## NEW API (Recommended)\n", + "\n", + "The new API uses stateless functions that follow the `PotentialFunction` protocol:\n", + "- `(theta, x) -> log_prob`\n", + "- No hidden state, all inputs explicit\n", + "- Combine with `bind_observation()`" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 61.52it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:05<00:00, 374.81it/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "NEW API - samples mean: tensor([0.9756, 0.0942])\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "# NEW API: stateless + bind_observation\n", + "from sbi.inference.potentials import bind_observation, likelihood_potential\n", + "from sbi.utils import mcmc_transform\n", + "\n", + "# Step 1: Create stateless potential (satisfies PotentialFunction protocol)\n", + "potential_fn = likelihood_potential(likelihood_estimator, prior)\n", + "\n", + "# Step 2: Get transform from prior (for MCMC in unconstrained space)\n", + "theta_transform = mcmc_transform(prior, enable_transform=True)\n", + "\n", + "# Step 3: Bind observation to create sampler-ready function\n", + "bound_fn = bind_observation(potential_fn, x_o, sum_iid=True)\n", + "\n", + "# Step 4: Use with MCMC sampler (pass theta_transform!)\n", + "posterior_new = MCMCPosterior(\n", + " bound_fn, proposal=prior, theta_transform=theta_transform, warmup_steps=10\n", + ")\n", + "samples_new = posterior_new.sample((1000,))\n", + "print(f\"NEW API - samples mean: {samples_new.mean(dim=0)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Lower-Level: Using the Protocol Directly\n", + "\n", + "You can also implement your own potential function that satisfies the protocol." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 60.21it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:02<00:00, 984.18it/s] " + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Custom potential - samples mean: tensor([0.9775, 0.0960])\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "# Custom potential function satisfying PotentialFunction protocol\n", + "def my_custom_potential(theta: torch.Tensor, x: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"Custom potential: log p(x|theta) + log p(theta)\"\"\"\n", + " log_likelihood = likelihood_estimator.log_prob(x, condition=theta)\n", + " log_prior = prior.log_prob(theta)\n", + " return log_likelihood + log_prior\n", + "\n", + "# Add device attribute for bind_observation compatibility\n", + "my_custom_potential.device = next(likelihood_estimator.parameters()).device\n", + "\n", + "# Bind observation\n", + "bound_custom = bind_observation(my_custom_potential, x_o, sum_iid=True)\n", + "\n", + "# Use with sampler\n", + "posterior_custom = MCMCPosterior(\n", + " bound_custom, proposal=prior, warmup_steps=10\n", + ")\n", + "samples_custom = posterior_custom.sample((1000,))\n", + "print(f\"Custom potential - samples mean: {samples_custom.mean(dim=0)}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "OLD API: tensor([0.9758, 0.0908])\n", + "NEW API: tensor([0.9756, 0.0942])\n", + "Custom: tensor([0.9775, 0.0960])\n" + ] + } + ], + "source": [ + "print(f\"OLD API: {samples_old.mean(dim=0)}\")\n", + "print(f\"NEW API: {samples_new.mean(dim=0)}\")\n", + "print(f\"Custom: {samples_custom.mean(dim=0)}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "sbi (3.12.12)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/sbi/inference/potentials/__init__.py b/sbi/inference/potentials/__init__.py index 8ed377efd..a5049d009 100644 --- a/sbi/inference/potentials/__init__.py +++ b/sbi/inference/potentials/__init__.py @@ -13,7 +13,6 @@ ) from sbi.inference.potentials.posterior_based_potential import ( posterior_estimator_based_potential, - posterior_potential, ) from sbi.inference.potentials.protocol import ( PotentialFunction, @@ -22,7 +21,6 @@ ) from sbi.inference.potentials.ratio_based_potential import ( ratio_estimator_based_potential, - ratio_potential, ) from sbi.inference.potentials.vector_field_potential import ( vector_field_estimator_based_potential, @@ -36,11 +34,9 @@ "likelihood_potential", "mixed_likelihood_estimator_based_potential", "posterior_estimator_based_potential", - "posterior_potential", "PotentialFunction", "PotentialFunctionWithDevice", "ratio_estimator_based_potential", - "ratio_potential", "validate_potential", "vector_field_estimator_based_potential", ] From b1137318432f4935badf53b867aadf1ed6c107ec Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Tue, 16 Jun 2026 10:03:52 +0200 Subject: [PATCH 6/9] fix: set_x() for actually making the notebook run --- docs/notebooks/sbi_old_vs_new_api.ipynb | 44 +++++++++------------- sbi/inference/potentials/base_potential.py | 11 +++++- sbi/inference/potentials/binding.py | 42 +++++++++++---------- 3 files changed, 49 insertions(+), 48 deletions(-) diff --git a/docs/notebooks/sbi_old_vs_new_api.ipynb b/docs/notebooks/sbi_old_vs_new_api.ipynb index 5a8e147cc..43cb43830 100644 --- a/docs/notebooks/sbi_old_vs_new_api.ipynb +++ b/docs/notebooks/sbi_old_vs_new_api.ipynb @@ -9,17 +9,9 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 6, "metadata": {}, "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/home/schmidt/gsoc2026/sbi/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n" - ] - }, { "name": "stdout", "output_type": "stream", @@ -65,24 +57,22 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "/home/schmidt/gsoc2026/sbi/sbi/inference/potentials/likelihood_based_potential.py:118: DeprecationWarning: LikelihoodBasedPotential is deprecated and will be removed in a future release. Use the stateless `likelihood_potential()` function instead, combined with `bind_observation()` for sampler compatibility.\n", - " potential_fn = LikelihoodBasedPotential(\n", - "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 46.64it/s]\n", - "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:04<00:00, 533.34it/s] " + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:02<00:00, 7.25it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:03<00:00, 594.07it/s] " ] }, { "name": "stdout", "output_type": "stream", "text": [ - "OLD API - samples mean: tensor([0.9758, 0.0908])\n" + "OLD API - samples mean: tensor([0.9753, 0.0945])\n" ] }, { @@ -121,22 +111,22 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 61.52it/s]\n", - "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:05<00:00, 374.81it/s]" + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 46.07it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:05<00:00, 425.35it/s]" ] }, { "name": "stdout", "output_type": "stream", "text": [ - "NEW API - samples mean: tensor([0.9756, 0.0942])\n" + "NEW API - samples mean: tensor([0.9765, 0.0895])\n" ] }, { @@ -180,22 +170,22 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 9, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ - "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 60.21it/s]\n", - "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:02<00:00, 984.18it/s] " + "Generating 20 MCMC inits via resample strategy: 100%|██████████| 20/20 [00:00<00:00, 45.43it/s]\n", + "Running vectorized MCMC with 20 chains: 100%|██████████| 2200/2200 [00:02<00:00, 998.14it/s] " ] }, { "name": "stdout", "output_type": "stream", "text": [ - "Custom potential - samples mean: tensor([0.9775, 0.0960])\n" + "Custom potential - samples mean: tensor([0.9757, 0.0981])\n" ] }, { @@ -230,16 +220,16 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 10, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ - "OLD API: tensor([0.9758, 0.0908])\n", - "NEW API: tensor([0.9756, 0.0942])\n", - "Custom: tensor([0.9775, 0.0960])\n" + "OLD API: tensor([0.9753, 0.0945])\n", + "NEW API: tensor([0.9765, 0.0895])\n", + "Custom: tensor([0.9757, 0.0981])\n" ] } ], diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index d24b05b94..f640d6e80 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/sbi/inference/potentials/base_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 inspect from abc import ABCMeta, abstractmethod from typing import Optional, Protocol, Union @@ -147,10 +148,18 @@ def to(self, device: Union[str, torch.device]) -> "CustomPotentialWrapper": super().to(device) return self + def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): + super().set_x(x_o, x_is_iid) + def __call__(self, theta, track_gradients: bool = True): """Calls the custom potential function on given theta. Note, x_o is re-used from the initialization of the potential function. """ + sig = inspect.signature(self.potential_fn) + num_params = len(sig.parameters) with torch.set_grad_enabled(track_gradients): - return self.potential_fn(theta, self.x_o) + if num_params == 1: + return self.potential_fn(theta) + else: + return self.potential_fn(theta, self.x_o) diff --git a/sbi/inference/potentials/binding.py b/sbi/inference/potentials/binding.py index 6d49e5675..2ca6d7aed 100644 --- a/sbi/inference/potentials/binding.py +++ b/sbi/inference/potentials/binding.py @@ -7,11 +7,12 @@ potential functions, making them ready for MCMC/VI samplers. """ -from typing import Callable, Union +from typing import Union import torch from torch import Tensor +from sbi.inference.potentials.base_potential import BasePotential from sbi.inference.potentials.protocol import PotentialFunction @@ -19,7 +20,7 @@ def bind_observation( potential: PotentialFunction, x_o: Tensor, sum_iid: bool = True, -) -> Callable[[Tensor], Tensor]: +) -> "BoundPotential": """Create a sampler-compatible theta -> log_prob function. This utility binds observed data to a stateless potential function, producing @@ -38,9 +39,9 @@ def bind_observation( compute log p(x_o | theta) = sum_i log p(x_i | theta). Returns: - A callable that takes theta (parameters) and returns the log-probability - log p(theta | x_o) + log p(theta). The returned function is stateless - from the sampler's perspective. + A BoundPotential instance that takes theta (parameters) and returns the + log-probability log p(theta | x_o) + log p(theta). The returned object + is compatible with samplers that expect set_x() method. Example: >>> # Create stateless potential @@ -67,22 +68,10 @@ def bind_observation( x_o = x_o.to(device) - def bound_potential(theta: Tensor) -> Tensor: - theta = theta.to(device) - - log_prob = potential(theta, x_o) - - if sum_iid and x_o.ndim > 1: - log_prob = log_prob.sum(dim=0) - elif sum_iid and x_o.ndim == 0: - pass - - return log_prob - - return bound_potential + return BoundPotential(potential, x_o, sum_iid) -class BoundPotential: +class BoundPotential(BasePotential): """Wrapper class for binding observation to a potential function. This provides an alternative to the functional bind_observation, maintaining @@ -100,6 +89,7 @@ def __init__( self._x_o = x_o.to(x_o.device) self._sum_iid = sum_iid self._device = x_o.device + self._x_is_iid = True @property def device(self) -> torch.device: @@ -114,7 +104,19 @@ def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: if self._sum_iid and self._x_o.ndim > 1: log_prob = log_prob.sum(dim=0) - return log_prob + return log_prob.reshape(-1) if log_prob.ndim > 0 else log_prob.unsqueeze(0) + + def set_x(self, x_o: Tensor, x_is_iid: bool = True) -> None: + """Set observation for the potential function. + + For BoundPotential, x_o is already bound in __init__, so this is a no-op + for backward compatibility with samplers. + """ + pass + + def return_x_o(self) -> Tensor: + """Return the bound observation.""" + return self._x_o def to(self, device: Union[str, torch.device]) -> "BoundPotential": self._device = torch.device(device) From 24315adc34515fdf00607d63970b89f6b84ba8bc Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Fri, 19 Jun 2026 13:05:15 +0200 Subject: [PATCH 7/9] fix: iid handling was broken, I got tests/potential_protocol_test.py:115: AssertionError --- sbi/inference/potentials/binding.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/sbi/inference/potentials/binding.py b/sbi/inference/potentials/binding.py index 2ca6d7aed..324d6f23d 100644 --- a/sbi/inference/potentials/binding.py +++ b/sbi/inference/potentials/binding.py @@ -104,7 +104,10 @@ def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: if self._sum_iid and self._x_o.ndim > 1: log_prob = log_prob.sum(dim=0) - return log_prob.reshape(-1) if log_prob.ndim > 0 else log_prob.unsqueeze(0) + if self._sum_iid: + return log_prob.reshape(-1) if log_prob.ndim > 0 else log_prob.unsqueeze(0) + else: + return log_prob def set_x(self, x_o: Tensor, x_is_iid: bool = True) -> None: """Set observation for the potential function. From a4fe9eb8b37539aa2a4d0881f444fe1c7484e036 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Sat, 20 Jun 2026 13:00:27 +0200 Subject: [PATCH 8/9] removed unused components --- sbi/inference/potentials/binding.py | 22 ---------------------- sbi/inference/potentials/protocol.py | 2 +- 2 files changed, 1 insertion(+), 23 deletions(-) diff --git a/sbi/inference/potentials/binding.py b/sbi/inference/potentials/binding.py index 324d6f23d..3e92ddb03 100644 --- a/sbi/inference/potentials/binding.py +++ b/sbi/inference/potentials/binding.py @@ -125,25 +125,3 @@ def to(self, device: Union[str, torch.device]) -> "BoundPotential": self._device = torch.device(device) self._x_o = self._x_o.to(device) return self - - -def bind_observation_class( - potential: PotentialFunction, - x_o: Tensor, - sum_iid: bool = True, -) -> BoundPotential: - """Create a BoundPotential instance for sampler-compatible usage. - - This is an alternative to bind_observation that returns a class instance - instead of a simple function. Useful when you need method chaining or - want to maintain compatibility with code that expects an object. - - Args: - potential: Stateless (theta, x) -> log_prob function. - x_o: Observed data tensor. - sum_iid: If True, sum log probabilities over IID batch dimension. - - Returns: - BoundPotential instance that can be called like a function. - """ - return BoundPotential(potential, x_o, sum_iid) diff --git a/sbi/inference/potentials/protocol.py b/sbi/inference/potentials/protocol.py index 39d129d0f..414ceffe2 100644 --- a/sbi/inference/potentials/protocol.py +++ b/sbi/inference/potentials/protocol.py @@ -25,7 +25,7 @@ class PotentialFunction(Protocol): - Easy testing (same inputs -> same outputs) - Simple serialization - Composable inference workflows - - Integration with external samplers (NumPyro, PyMC, custom PyTorch) + - Integration with external samplers (PyMC, custom PyTorch) Example: >>> def my_potential(theta: Tensor, x: Tensor) -> Tensor: From 50016974477a49e1ee1746185fa0bd3147134176 Mon Sep 17 00:00:00 2001 From: Jocho-Smith Date: Sat, 20 Jun 2026 13:03:31 +0200 Subject: [PATCH 9/9] fixup:remove links inside __init__ as well --- sbi/inference/potentials/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/sbi/inference/potentials/__init__.py b/sbi/inference/potentials/__init__.py index a5049d009..a2fb09991 100644 --- a/sbi/inference/potentials/__init__.py +++ b/sbi/inference/potentials/__init__.py @@ -4,7 +4,6 @@ from sbi.inference.potentials.binding import ( BoundPotential, bind_observation, - bind_observation_class, ) from sbi.inference.potentials.likelihood_based_potential import ( likelihood_estimator_based_potential, @@ -28,7 +27,6 @@ __all__ = [ "bind_observation", - "bind_observation_class", "BoundPotential", "likelihood_estimator_based_potential", "likelihood_potential",