fix: map amortized VI flow to the prior support

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,
@@ -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
@@ -208,7 +209,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
@@ -278,9 +279,29 @@ 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)
+            # 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
 
+    @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, 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 = link_transform
+        self._amortized_q._support = self._prior_support(prior)
+
     def _build_unconditional_flow(
         self,
         flow_type: str,
@@ -411,7 +432,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 +448,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,14 +463,19 @@ 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,
+        )
+        return TransformedConditionalZukoFlow(
+            flow, self.link_transform, self._prior_support(self._prior)
         ).to(self._device)
 
     @property
@@ -561,6 +587,17 @@ 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)
+                # 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(
+                    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
             self._x = q._x

sbi/samplers/vi/vi_utils.py

@@ -20,9 +20,10 @@
     Normal,
     TransformedDistribution,
 )
+from torch.distributions.constraints import Constraint
 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 +137,71 @@ 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,
+        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:
+        # 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)
+        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
+    ) -> 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.
 

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,25 @@ 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()
+    outside = torch.tensor([[4.0, 0.0]])
+    assert posterior.log_prob(outside, x=x[:1]).item() == -torch.inf
+
+
 @pytest.mark.slow
 def test_amortized_vi_log_prob(linear_gaussian_setup):
     """Test log_prob evaluation in amortized mode."""