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:
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
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..43cb43830
--- /dev/null
+++ b/docs/notebooks/sbi_old_vs_new_api.ipynb
@@ -0,0 +1,264 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Old API vs New API Demo"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 6,
+ "metadata": {},
+ "outputs": [
+ {
+ "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": 7,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "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.9753, 0.0945])\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": 8,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "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.9765, 0.0895])\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": 9,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "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.9757, 0.0981])\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": 10,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "OLD API: tensor([0.9753, 0.0945])\n",
+ "NEW API: tensor([0.9765, 0.0895])\n",
+ "Custom: tensor([0.9757, 0.0981])\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 c88d6301c..a2fb09991 100644
--- a/sbi/inference/potentials/__init__.py
+++ b/sbi/inference/potentials/__init__.py
@@ -1,13 +1,23 @@
# 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,
+)
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,
)
+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",
+ "BoundPotential",
"likelihood_estimator_based_potential",
+ "likelihood_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/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
new file mode 100644
index 000000000..3e92ddb03
--- /dev/null
+++ b/sbi/inference/potentials/binding.py
@@ -0,0 +1,127 @@
+# 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 Union
+
+import torch
+from torch import Tensor
+
+from sbi.inference.potentials.base_potential import BasePotential
+from sbi.inference.potentials.protocol import PotentialFunction
+
+
+def bind_observation(
+ potential: PotentialFunction,
+ x_o: Tensor,
+ sum_iid: bool = True,
+) -> "BoundPotential":
+ """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 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
+ >>> 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)
+
+ return BoundPotential(potential, x_o, sum_iid)
+
+
+class BoundPotential(BasePotential):
+ """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
+ self._x_is_iid = True
+
+ @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)
+
+ 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.
+
+ 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)
+ self._x_o = self._x_o.to(device)
+ return self
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/sbi/inference/potentials/protocol.py b/sbi/inference/potentials/protocol.py
new file mode 100644
index 000000000..414ceffe2
--- /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 (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..6c125998b
--- /dev/null
+++ b/tests/potential_protocol_test.py
@@ -0,0 +1,357 @@
+# 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 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."""
+
+ 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