From 0183b4ba9bd6faceb9fb96905997ee0f857d68ed Mon Sep 17 00:00:00 2001 From: andre israel Date: Sun, 1 Nov 2026 10:52:09 -0300 Subject: [PATCH 1/7] fix: map amortized VI flow to the prior support (Fixes #2137) --- sbi/inference/posteriors/vi_posterior.py | 19 +++++---- sbi/samplers/vi/vi_utils.py | 53 +++++++++++++++++++++++- tests/vi_test.py | 19 ++++++++- 3 files changed, 82 insertions(+), 9 deletions(-) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index 07c22a27d..c834cd58a 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -20,7 +20,6 @@ if TYPE_CHECKING: from sbi.inference.posteriors.posterior_parameters import VIPosteriorParameters from sbi.inference.potentials.base_potential import BasePotential, CustomPotential -from sbi.neural_nets.estimators.base import ConditionalDensityEstimator from sbi.neural_nets.estimators.zuko_flow import ZukoUnconditionalFlow from sbi.neural_nets.factory import ZukoFlowType from sbi.neural_nets.net_builders.flow import ( @@ -32,6 +31,7 @@ from sbi.samplers.vi.vi_utils import ( AdaptedVariationalDistribution, LearnableGaussian, + TransformedConditionalZukoFlow, TransformedZukoFlow, check_variational_distribution, detach_and_deepcopy, @@ -208,7 +208,7 @@ def __init__( self._mode: Optional[Literal["single_x", "amortized"]] = None # Amortized mode: conditional flow q(θ|x) - self._amortized_q: Optional[ConditionalDensityEstimator] = None + self._amortized_q: Optional[TransformedConditionalZukoFlow] = None self._num_transforms: int = num_transforms self._hidden_features: int = hidden_features @@ -411,7 +411,7 @@ def _build_conditional_flow( hidden_features: int = 32, z_score_theta: Literal["none", "independent", "structured"] = "independent", z_score_x: Literal["none", "independent", "structured"] = "independent", - ) -> ConditionalDensityEstimator: + ) -> TransformedConditionalZukoFlow: """Build a conditional Zuko flow for amortized variational inference. Args: @@ -427,7 +427,7 @@ def _build_conditional_flow( structured data like images. Returns: - ConditionalDensityEstimator: The constructed conditional flow q(θ|x). + TransformedConditionalZukoFlow: The conditional flow q(θ|x) in prior space. Raises: ValueError: If flow_type is not supported. @@ -442,15 +442,20 @@ def _build_conditional_flow( f"Supported types: {[t.name for t in ZukoFlowType]}." ) from e - return build_zuko_flow( + # The flow models θ in unconstrained space; the wrapper maps it to the prior + # support so that bounded priors are not evaluated outside their support. + flow = build_zuko_flow( flow_type.value.upper(), - batch_x=theta, # θ is what we model + batch_x=self.link_transform.inv(theta), # θ is what we model batch_y=x, # x is the condition z_score_x=z_score_theta, # z-score for θ (naming mismatch) z_score_y=z_score_x, # z-score for x condition num_transforms=num_transforms, hidden_features=hidden_features, - ).to(self._device) + ) + return TransformedConditionalZukoFlow(flow, self.link_transform).to( + self._device + ) @property def q( diff --git a/sbi/samplers/vi/vi_utils.py b/sbi/samplers/vi/vi_utils.py index d41a6a621..39d0a5ec8 100644 --- a/sbi/samplers/vi/vi_utils.py +++ b/sbi/samplers/vi/vi_utils.py @@ -22,7 +22,7 @@ ) from torch.nn import Module -from sbi.neural_nets.estimators.zuko_flow import ZukoUnconditionalFlow +from sbi.neural_nets.estimators.zuko_flow import ZukoFlow, ZukoUnconditionalFlow from sbi.sbi_types import Shape, TorchTransform, VariationalDistribution from sbi.utils.torchutils import _base_recursor @@ -136,6 +136,57 @@ def sample_and_log_prob(self, sample_shape: Shape) -> tuple[Tensor, Tensor]: return theta, log_prob_theta +class TransformedConditionalZukoFlow(nn.Module): + """Conditional counterpart of `TransformedZukoFlow` for amortized VI. + + The wrapped flow q(z|x) lives in unconstrained space. Samples and log-probs of + this wrapper are in the prior's (constrained) space, so a bounded prior is never + evaluated outside its support. + """ + + def __init__(self, flow: ZukoFlow, link_transform: TorchTransform): + """Initialize the wrapper. + + Args: + flow: Conditional Zuko flow operating in unconstrained space. + link_transform: Transform from unconstrained to constrained space. + """ + super().__init__() + self._flow = flow + self._link_transform = link_transform + + @staticmethod + def _sum_log_det(log_det: Tensor, like: Tensor) -> Tensor: + # Elementwise transforms return per-dimension Jacobians; sum them if needed. + return log_det.sum(dim=-1) if log_det.dim() > like.dim() else log_det + + def sample(self, sample_shape: Shape, condition: Tensor) -> Tensor: + """Sample θ ~ q(θ|x) in constrained space.""" + theta = self._link_transform(self._flow.sample(sample_shape, condition)) + assert isinstance(theta, Tensor) + return theta + + def log_prob(self, theta: Tensor, condition: Tensor) -> Tensor: + """Log density of constrained θ, including the Jacobian of the link.""" + z = self._link_transform.inv(theta) + assert isinstance(z, Tensor) + log_prob_z = self._flow.log_prob(z, condition) + log_det = self._link_transform.inv.log_abs_det_jacobian(theta, z) + return log_prob_z + self._sum_log_det(log_det, log_prob_z) + + def sample_and_log_prob( + self, sample_shape: Shape, condition: Tensor + ) -> tuple[Tensor, Tensor]: + """Reparameterized constrained samples and their log density.""" + z, log_prob_z = self._flow.sample_and_log_prob( + torch.Size(sample_shape), condition + ) + theta = self._link_transform(z) + assert isinstance(theta, Tensor) + log_det = self._link_transform.log_abs_det_jacobian(z, theta) + return theta, log_prob_z - self._sum_log_det(log_det, log_prob_z) + + class LearnableGaussian(nn.Module): """Learnable Gaussian distribution for variational inference. diff --git a/tests/vi_test.py b/tests/vi_test.py index 18b23cc03..f441346f8 100644 --- a/tests/vi_test.py +++ b/tests/vi_test.py @@ -37,7 +37,7 @@ linear_gaussian, true_posterior_linear_gaussian_mvn_prior, ) -from sbi.utils import MultipleIndependent +from sbi.utils import BoxUniform, MultipleIndependent from sbi.utils.metrics import c2st, check_c2st from sbi.utils.user_input_checks import process_x @@ -878,6 +878,23 @@ def test_amortized_vi_batched_sampling(linear_gaussian_setup): assert samples.shape == (num_samples, num_obs, setup["num_dim"]) +def test_amortized_vi_with_bounded_prior(): + """Amortized VI trains and samples inside the support of a `BoxUniform` prior.""" + prior = BoxUniform(-3 * ones(2), 3 * ones(2)) + theta = prior.sample((300,)) + x = theta + 0.3 * torch.randn_like(theta) + nle = NLE(prior, show_progress_bars=False) + nle.append_simulations(theta, x).train(max_num_epochs=3) + posterior = nle.build_posterior(sample_with="vi") + + posterior.train_amortized(theta, x, max_num_iters=5, show_progress_bar=False) + + samples = posterior.sample((50,), x=x[:1], show_progress_bars=False) + assert torch.isfinite(samples).all() + assert (samples.abs() <= 3).all() + assert torch.isfinite(posterior.log_prob(samples, x=x[:1])).all() + + @pytest.mark.slow def test_amortized_vi_log_prob(linear_gaussian_setup): """Test log_prob evaluation in amortized mode.""" From 5ee94917c98252d332091d2b49f0c9be083d5ef7 Mon Sep 17 00:00:00 2001 From: andre israel Date: Sun, 1 Nov 2026 11:29:15 -0300 Subject: [PATCH 2/7] fix: keep amortized link in sync on .to() and return -inf outside the prior support --- sbi/inference/posteriors/vi_posterior.py | 8 +++++--- sbi/samplers/vi/vi_utils.py | 19 +++++++++++++++++-- tests/vi_test.py | 2 ++ 3 files changed, 24 insertions(+), 5 deletions(-) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index c834cd58a..8776783e6 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -278,6 +278,8 @@ def to(self, device: Union[str, torch.device]) -> "VIPosterior": self._q._link_transform = self.link_transform # type: ignore[union-attr] if self._amortized_q is not None: self._amortized_q.to(device) + self._amortized_q._link_transform = self.link_transform + self._amortized_q._support = getattr(self._prior, "support", None) return self @@ -453,9 +455,9 @@ def _build_conditional_flow( num_transforms=num_transforms, hidden_features=hidden_features, ) - return TransformedConditionalZukoFlow(flow, self.link_transform).to( - self._device - ) + return TransformedConditionalZukoFlow( + flow, self.link_transform, getattr(self._prior, "support", None) + ).to(self._device) @property def q( diff --git a/sbi/samplers/vi/vi_utils.py b/sbi/samplers/vi/vi_utils.py index 39d0a5ec8..f27fc53d8 100644 --- a/sbi/samplers/vi/vi_utils.py +++ b/sbi/samplers/vi/vi_utils.py @@ -20,6 +20,7 @@ Normal, TransformedDistribution, ) +from torch.distributions.constraints import Constraint from torch.nn import Module from sbi.neural_nets.estimators.zuko_flow import ZukoFlow, ZukoUnconditionalFlow @@ -144,16 +145,24 @@ class TransformedConditionalZukoFlow(nn.Module): evaluated outside its support. """ - def __init__(self, flow: ZukoFlow, link_transform: TorchTransform): + def __init__( + self, + flow: ZukoFlow, + link_transform: TorchTransform, + support: Optional[Constraint] = None, + ): """Initialize the wrapper. Args: flow: Conditional Zuko flow operating in unconstrained space. link_transform: Transform from unconstrained to constrained space. + support: Support of the prior. `log_prob` returns `-inf` outside of it + (the inverse link would otherwise clamp such points into it). """ super().__init__() self._flow = flow self._link_transform = link_transform + self._support = support @staticmethod def _sum_log_det(log_det: Tensor, like: Tensor) -> Tensor: @@ -172,7 +181,13 @@ def log_prob(self, theta: Tensor, condition: Tensor) -> Tensor: assert isinstance(z, Tensor) log_prob_z = self._flow.log_prob(z, condition) log_det = self._link_transform.inv.log_abs_det_jacobian(theta, z) - return log_prob_z + self._sum_log_det(log_det, log_prob_z) + log_prob = log_prob_z + self._sum_log_det(log_det, log_prob_z) + if self._support is None: + return log_prob + in_support = self._support.check(theta) + if in_support.dim() > log_prob.dim(): + in_support = in_support.all(dim=-1) + return torch.where(in_support, log_prob, torch.full_like(log_prob, -torch.inf)) def sample_and_log_prob( self, sample_shape: Shape, condition: Tensor diff --git a/tests/vi_test.py b/tests/vi_test.py index f441346f8..02c819366 100644 --- a/tests/vi_test.py +++ b/tests/vi_test.py @@ -893,6 +893,8 @@ def test_amortized_vi_with_bounded_prior(): assert torch.isfinite(samples).all() assert (samples.abs() <= 3).all() assert torch.isfinite(posterior.log_prob(samples, x=x[:1])).all() + outside = torch.tensor([[4.0, 0.0]]) + assert posterior.log_prob(outside, x=x[:1]).item() == -torch.inf @pytest.mark.slow From bc1d3a56417267a3a3d312d9cd4e925020ef9e83 Mon Sep 17 00:00:00 2001 From: andre israel Date: Sun, 1 Nov 2026 11:51:28 -0300 Subject: [PATCH 3/7] fix: tolerate priors without support and rebind amortized flow in set_q --- sbi/inference/posteriors/vi_posterior.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index 8776783e6..6c18c4c44 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -278,11 +278,24 @@ def to(self, device: Union[str, torch.device]) -> "VIPosterior": self._q._link_transform = self.link_transform # type: ignore[union-attr] if self._amortized_q is not None: self._amortized_q.to(device) - self._amortized_q._link_transform = self.link_transform - self._amortized_q._support = getattr(self._prior, "support", None) + self._rebind_amortized_q(self._prior) return self + @staticmethod + def _prior_support(prior: Distribution): + """Support of the prior, or None if it does not define one.""" + try: + return prior.support + except (AttributeError, NotImplementedError): + return None + + def _rebind_amortized_q(self, prior: Distribution) -> None: + """Point the amortized flow's link and support at tensors on this device.""" + assert self._amortized_q is not None + self._amortized_q._link_transform = self.link_transform + self._amortized_q._support = self._prior_support(prior) + def _build_unconditional_flow( self, flow_type: str, @@ -456,7 +469,7 @@ def _build_conditional_flow( hidden_features=hidden_features, ) return TransformedConditionalZukoFlow( - flow, self.link_transform, getattr(self._prior, "support", None) + flow, self.link_transform, self._prior_support(self._prior) ).to(self._device) @property @@ -568,6 +581,7 @@ def set_q( self._amortized_q = deepcopy(getattr(q, "_amortized_q", None)) if self._amortized_q is not None: self._amortized_q.to(self._device) + self._rebind_amortized_q(q._prior) self.vi_method = q.vi_method # type: ignore self._prior = q._prior self._x = q._x From 04aa84e558f3abde2b857dbecd83c1bb23b459a7 Mon Sep 17 00:00:00 2001 From: andre israel Date: Sun, 1 Nov 2026 12:00:03 -0300 Subject: [PATCH 4/7] fix: bind a device-local copy of the source prior's support in set_q --- sbi/inference/posteriors/vi_posterior.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index 6c18c4c44..a0551d104 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -581,7 +581,10 @@ def set_q( self._amortized_q = deepcopy(getattr(q, "_amortized_q", None)) if self._amortized_q is not None: self._amortized_q.to(self._device) - self._rebind_amortized_q(q._prior) + # Copy, so the source posterior's prior stays on its own device. + self._rebind_amortized_q( + move_distribution_to_device(deepcopy(q._prior), self._device) + ) self.vi_method = q.vi_method # type: ignore self._prior = q._prior self._x = q._x From 3e52c51fe55ed90369baed1a2e5fb2e3d522f492 Mon Sep 17 00:00:00 2001 From: andre israel Date: Sun, 1 Nov 2026 12:06:04 -0300 Subject: [PATCH 5/7] fix: rebind a copied amortized flow to the source prior's own link and support --- sbi/inference/posteriors/vi_posterior.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index a0551d104..de236bb5d 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -278,7 +278,7 @@ def to(self, device: Union[str, torch.device]) -> "VIPosterior": self._q._link_transform = self.link_transform # type: ignore[union-attr] if self._amortized_q is not None: self._amortized_q.to(device) - self._rebind_amortized_q(self._prior) + self._rebind_amortized_q(self._prior, self.link_transform) return self @@ -290,10 +290,12 @@ def _prior_support(prior: Distribution): except (AttributeError, NotImplementedError): return None - def _rebind_amortized_q(self, prior: Distribution) -> None: + def _rebind_amortized_q( + self, prior: Distribution, link_transform: TorchTransform + ) -> None: """Point the amortized flow's link and support at tensors on this device.""" assert self._amortized_q is not None - self._amortized_q._link_transform = self.link_transform + self._amortized_q._link_transform = link_transform self._amortized_q._support = self._prior_support(prior) def _build_unconditional_flow( @@ -581,10 +583,13 @@ def set_q( self._amortized_q = deepcopy(getattr(q, "_amortized_q", None)) if self._amortized_q is not None: self._amortized_q.to(self._device) - # Copy, so the source posterior's prior stays on its own device. - self._rebind_amortized_q( - move_distribution_to_device(deepcopy(q._prior), self._device) - ) + # Link and support come from the source posterior's prior (not this + # one's), copied so that the source stays on its own device. + prior = move_distribution_to_device(deepcopy(q._prior), self._device) + theta_transform = q.theta_transform + if theta_transform is None: + theta_transform = mcmc_transform(prior, device=self._device) + self._rebind_amortized_q(prior, theta_transform.inv) self.vi_method = q.vi_method # type: ignore self._prior = q._prior self._x = q._x From 8baea453c8f037f2c175f36b2f400286a771ba8c Mon Sep 17 00:00:00 2001 From: andre israel Date: Sun, 1 Nov 2026 14:34:43 -0300 Subject: [PATCH 6/7] fix: copy the source prior and custom transform safely in set_q --- sbi/inference/posteriors/vi_posterior.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index de236bb5d..35f4407d1 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -46,6 +46,7 @@ from sbi.utils.torchutils import ( atleast_2d_float32_tensor, ensure_theta_batched, + move_all_tensor_to_device, process_device, ) from sbi.utils.user_input_checks_utils import move_distribution_to_device @@ -585,10 +586,14 @@ def set_q( self._amortized_q.to(self._device) # Link and support come from the source posterior's prior (not this # one's), copied so that the source stays on its own device. - prior = move_distribution_to_device(deepcopy(q._prior), self._device) - theta_transform = q.theta_transform - if theta_transform is None: + prior = move_distribution_to_device( + detach_and_deepcopy(q._prior), self._device + ) + if q.theta_transform is None: theta_transform = mcmc_transform(prior, device=self._device) + else: + theta_transform = detach_and_deepcopy(q.theta_transform) + move_all_tensor_to_device(theta_transform, self._device) self._rebind_amortized_q(prior, theta_transform.inv) self.vi_method = q.vi_method # type: ignore self._prior = q._prior From b6a28da9b771dabed8a5da9303723d1008611650 Mon Sep 17 00:00:00 2001 From: andre israel Date: Sun, 1 Nov 2026 14:50:07 -0300 Subject: [PATCH 7/7] fix: keep the trained amortized link transform when moving devices --- sbi/inference/posteriors/vi_posterior.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index 35f4407d1..9fb05044b 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -279,7 +279,10 @@ def to(self, device: Union[str, torch.device]) -> "VIPosterior": self._q._link_transform = self.link_transform # type: ignore[union-attr] if self._amortized_q is not None: self._amortized_q.to(device) - self._rebind_amortized_q(self._prior, self.link_transform) + # Keep the trained link: rebuilding it can re-estimate prior moments. + link_transform = self._amortized_q._link_transform + move_all_tensor_to_device(link_transform, device) + self._rebind_amortized_q(self._prior, link_transform) return self