feat(mdn): implement transform_to_unconstrained z-scoring for MDN

sbi/inference/posteriors/npe_a_posterior.py

@@ -171,7 +171,7 @@ def _corrected_log_prob(self, theta: Tensor, condition: Tensor) -> Tensor:
 
         # Add log det jacobian for z-score transform
         log_probs = log_probs + self.posterior_estimator._log_det_jacobian_forward(
-            theta
+            theta, theta_transformed
         )
 
         return log_probs

sbi/inference/trainers/npe/npe_a.py

@@ -354,6 +354,12 @@ def _compute_z_scored_prior_mog(
         prior_cov = self._prior.covariance_matrix
 
         # Apply z-score transform if enabled
+        if getattr(density_estimator, "_prior_transform", None) is not None:
+            raise NotImplementedError(
+                "NPE-A's analytic leakage correction does not support "
+                "z_score_theta='transform_to_unconstrained' (it assumes an affine "
+                "z-score). Use 'independent'/'structured', or a different method."
+            )
         if density_estimator.has_input_transform:
             shift = density_estimator._transform_shift
             scale = density_estimator._transform_scale

sbi/inference/trainers/npe/npe_c.py

@@ -239,6 +239,12 @@ def _set_state_for_mog_proposal(self) -> None:
         """
         # Check if z-scoring is enabled on the MixtureDensityEstimator
         assert isinstance(self._neural_net, MixtureDensityEstimator)
+        if getattr(self._neural_net, "_prior_transform", None) is not None:
+            raise NotImplementedError(
+                "NPE-C's analytic proposal correction does not support "
+                "z_score_theta='transform_to_unconstrained' (it assumes an affine "
+                "z-score). Use 'independent'/'structured', or single-round NPE."
+            )
         self.z_score_theta = self._neural_net.has_input_transform
 
         self._set_maybe_z_scored_prior()

sbi/neural_nets/estimators/mixture_density_estimator.py

@@ -21,6 +21,7 @@
 
 from sbi.neural_nets.estimators.base import ConditionalDensityEstimator
 from sbi.neural_nets.estimators.mog import MoG
+from sbi.sbi_types import TorchTransform
 
 
 class MultivariateGaussianMDN(nn.Module):
@@ -337,6 +338,7 @@ def __init__(
         condition_shape: torch.Size,
         embedding_net: Optional[nn.Module] = None,
         transform_input: Optional[Tensor] = None,
+        prior_transform: Optional[TorchTransform] = None,
     ) -> None:
         """Initialize the mixture density estimator.
 
@@ -354,12 +356,21 @@ def __init__(
                 evaluation, and samples are inverse transformed as:
                 x = z * scale + shift.
                 This is used for z-scoring inputs to improve numerical stability.
+            prior_transform: Optional Transform mapping constrained -> unconstrained
+                space. Mutually exclusive with transform_input. When set, log_prob
+                applies this transform (with log-det correction) and sample applies
+                its inverse.
         """
         super().__init__(net, input_shape, condition_shape)
         self._embedding_net = (
             embedding_net if embedding_net is not None else nn.Identity()
         )
 
+        if transform_input is not None and prior_transform is not None:
+            raise ValueError(
+                "Only one of transform_input and prior_transform can be provided."
+            )
+
         # Validate that embedding_net output matches MDN input if embedding is provided
         if embedding_net is not None:
             with torch.no_grad():
@@ -372,6 +383,8 @@ def __init__(
                     f"MDN context_features ({net.context_features})"
                 )
 
+        self._prior_transform = prior_transform
+
         # Store z-score transform parameters as buffers (not trained, moved with model)
         if transform_input is not None:
             if transform_input.shape[0] != 2:
@@ -404,17 +417,23 @@ def has_input_transform(self) -> bool:
 
     def _transform_input(self, input: Tensor) -> Tensor:
         """Apply z-score transform to input: z = (x - shift) / scale."""
-        if not self.has_input_transform:
+        if self._prior_transform is not None:
+            return self._prior_transform(input)
+        if self._transform_shift is None:
             return input
         return (input - self._transform_shift) / self._transform_scale
 
     def _inverse_transform_input(self, z: Tensor) -> Tensor:
         """Apply inverse z-score transform: x = z * scale + shift."""
-        if not self.has_input_transform:
+        if self._prior_transform is not None:
+            return self._prior_transform.inv(z)
+        if self._transform_shift is None:
             return z
         return z * self._transform_scale + self._transform_shift
 
-    def _log_det_jacobian_forward(self, input: Tensor) -> Tensor:
+    def _log_det_jacobian_forward(
+        self, input: Tensor, transformed_input: Tensor
+    ) -> Tensor:
         """Compute log determinant of Jacobian for the forward z-score transform.
 
         For the forward affine transform z = (x - shift) / scale:
@@ -437,11 +456,15 @@ def _log_det_jacobian_forward(self, input: Tensor) -> Tensor:
 
         Args:
             input: Input tensor, used to determine device and dtype.
+            transformed_input: Transformed input tensor (for prior_transform path).
 
         Returns:
-            Log determinant of forward Jacobian (scalar), on the same device as input.
+            Log determinant of forward Jacobian (scalar for affine, same shape as
+            input for prior_transform).
         """
-        if not self.has_input_transform:
+        if self._prior_transform is not None:
+            return self._prior_transform.log_abs_det_jacobian(input, transformed_input)
+        if self._transform_shift is None:
             return torch.zeros(1, device=input.device, dtype=input.dtype).squeeze()
         return -torch.log(self._transform_scale).sum()
 
@@ -478,7 +501,7 @@ def log_prob(self, input: Tensor, condition: Tensor, **kwargs) -> Tensor:
         # Change of variables: log p(x) = log p(z) + log|det(dz/dx)|
         # where z = (x - shift) / scale and log|det(dz/dx)| = -sum(log(scale))
         log_probs = mog.log_prob(transformed_input)  # (sample_dim, batch_dim)
-        log_probs = log_probs + self._log_det_jacobian_forward(input)
+        log_probs = log_probs + self._log_det_jacobian_forward(input, transformed_input)
 
         if not has_sample_dim:
             log_probs = log_probs.squeeze(0)

sbi/neural_nets/net_builders/mdn.py

@@ -5,14 +5,15 @@
 
 import torch
 from torch import Tensor, nn
+from torch.distributions import Distribution
 
 from sbi.neural_nets.estimators.mixture_density_estimator import (
     MixtureDensityEstimator,
     MultivariateGaussianMDN,
 )
 from sbi.utils.nn_utils import get_numel
 from sbi.utils.sbiutils import (
-    assert_transform_to_unconstrained_supported,
+    mcmc_transform,
     standardizing_net,
     z_score_parser,
     z_standardization,
@@ -28,6 +29,7 @@ def build_mdn(
     hidden_features: int = 50,
     num_components: int = 10,
     embedding_net: nn.Module = nn.Identity(),
+    x_dist: Optional[Distribution] = None,
     **kwargs,
 ) -> MixtureDensityEstimator:
     """Builds MDN p(x|y).
@@ -41,31 +43,37 @@ def build_mdn(
             - `structured`: treat dimensions as related, therefore compute mean and std
             over the entire batch, instead of per-dimension. Should be used when each
             sample is, for example, a time series or an image.
+            - `transform_to_unconstrained`: transform inputs to unconstrained space
+            using the prior's support. Requires `x_dist` to be provided.
         z_score_y: Whether to z-score ys passing into the network, same options as
             z_score_x.
         hidden_features: Number of hidden features.
         num_components: Number of components.
         embedding_net: Optional embedding network for y.
+        x_dist: Optional prior distribution. Required when
+            ``z_score_x="transform_to_unconstrained"``.
         kwargs: Additional arguments that are passed by the build function but are not
             relevant for MDNs and are therefore ignored.
 
     Returns:
         MixtureDensityEstimator for conditional density estimation.
     """
     check_data_device(batch_x, batch_y)
-    assert_transform_to_unconstrained_supported(
-        z_score_x,
-        "build_mdn",
-        "Use a `zuko_*` model (e.g. `zuko_maf`, `zuko_nsf`), which supports it, "
-        "or one of 'none', 'independent', 'structured'.",
-    )
+
     x_numel = get_numel(batch_x, embedding_net=None)
     y_numel = get_numel(batch_y, embedding_net=embedding_net)
 
     # Handle z-scoring for x (input)
     transform_input = None
+    prior_transform = None
     z_score_x_bool, structured_x = z_score_parser(z_score_x)
-    if z_score_x_bool:
+    if z_score_x == "transform_to_unconstrained":
+        if x_dist is None:
+            raise ValueError(
+                "x_dist must be provided when z_score_x='transform_to_unconstrained'."
+            )
+        prior_transform = mcmc_transform(x_dist, device=batch_x.device)
+    elif z_score_x_bool:
         x_mean, x_std = z_standardization(batch_x, structured_x)
         # Store as [shift, scale] tensor for the estimator
         transform_input = torch.stack([x_mean, x_std], dim=0)
@@ -93,6 +101,7 @@ def build_mdn(
         condition_shape=batch_y[0].shape,
         embedding_net=embedding_net,
         transform_input=transform_input,
+        prior_transform=prior_transform,
     )
 
     return estimator

sbi/utils/conditional_density_utils.py

@@ -193,6 +193,12 @@ def extract_and_transform_mog(
     assert precfs is not None  # For type checker
 
     # Transform means and precision factors to original space if z-scoring is enabled
+    if getattr(estimator, "_prior_transform", None) is not None:
+        raise NotImplementedError(
+            "Conditional density extraction is not supported for "
+            "z_score='transform_to_unconstrained' (it assumes an affine "
+            "z-score)."
+        )
     if estimator.has_input_transform:
         # The z-score transform is: z = (theta - shift) / scale
         # To transform means: theta = z * scale + shift

tests/sbiutils_test.py

@@ -497,7 +497,11 @@ def test_z_scoring_structured(z_x, z_theta, build_fn):
 
         # Unsupported combination: the modeled variable requests the unconstrained
         # transform on a non-Zuko-conditional builder -> expect a clear ValueError.
-        if modeled_z == "transform_to_unconstrained" and not model.startswith("zuko"):
+        if (
+            modeled_z == "transform_to_unconstrained"
+            and not model.startswith("zuko")
+            and model != "mdn"
+        ):
             with pytest.raises(ValueError, match="transform_to_unconstrained"):
                 build_fun = build_fn(**kwargs)
                 build_fun(theta, model_x)
@@ -549,30 +553,27 @@ def test_z_scoring_structured(z_x, z_theta, build_fn):
         "build_maf",
         "build_maf_rqs",
         "build_nsf",
-        "build_mdn",
         "build_linear_classifier",
         "build_mlp_classifier",
         "build_resnet_classifier",
     ],
 )
 def test_transform_to_unconstrained_raises_for_unsupported_builders(builder_name):
-    """nflows, MDN, and ratio-classifier builders raise a clear error for
+    """nflows and ratio-classifier builders raise a clear error for
     `transform_to_unconstrained` instead of silently building a model without the
-    reparametrization."""
+    reparametrization. MDN now supports it and is excluded."""
     from sbi.neural_nets.net_builders import flow as flow_builders
     from sbi.neural_nets.net_builders.classifier import (
         build_linear_classifier,
         build_mlp_classifier,
         build_resnet_classifier,
     )
-    from sbi.neural_nets.net_builders.mdn import build_mdn
 
     builders = {
         "build_made": flow_builders.build_made,
         "build_maf": flow_builders.build_maf,
         "build_maf_rqs": flow_builders.build_maf_rqs,
         "build_nsf": flow_builders.build_nsf,
-        "build_mdn": build_mdn,
         "build_linear_classifier": build_linear_classifier,
         "build_mlp_classifier": build_mlp_classifier,
         "build_resnet_classifier": build_resnet_classifier,
@@ -690,3 +691,21 @@ def test_higher_dimensional_tensor(self):
         x_3d_outlier[0, 2, 3] = 10000.0  # Extreme outlier at position [2,3]
         with pytest.warns(UserWarning, match="extreme outliers"):
             warn_if_invalid_for_zscoring(x_3d_outlier)
+
+
+def test_mdn_transform_to_unconstrained():
+    """MDN with transform_to_unconstrained produces valid log_probs and samples."""
+    from sbi.neural_nets.net_builders.mdn import build_mdn
+
+    prior = BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
+    bx, by = prior.sample((512,)), torch.randn(512, 3)
+    est = build_mdn(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
+    theta, cond = prior.sample((5,)), torch.randn(1, 3)
+    lp = est.log_prob(theta.unsqueeze(1), cond)
+    assert lp.shape == (5, 1)
+    z = est._prior_transform(theta.unsqueeze(1))
+    mog = est.get_uncorrected_mog(cond)
+    ldj = est._prior_transform.log_abs_det_jacobian(theta.unsqueeze(1), z)
+    assert torch.allclose(lp, mog.log_prob(z) + ldj, atol=1e-5)
+    s = est.sample((10,), cond)
+    assert s.shape[0] == 10 and torch.isfinite(s).all()