From ab02643f09174e0beb4665784738fb4d8efd6297 Mon Sep 17 00:00:00 2001
From: BHARATH0153 <chandrabharath463@gmail.com>
Date: Wed, 24 Jun 2026 00:14:32 +0530
Subject: [PATCH 1/3] feat(mdn): implement transform_to_unconstrained z-scoring
 for MDN

---
 .../estimators/mixture_density_estimator.py   | 67 +++++++++++++------
 sbi/neural_nets/net_builders/mdn.py           | 22 +++---
 tests/sbiutils_test.py                        | 14 ++--
 3 files changed, 67 insertions(+), 36 deletions(-)

diff --git a/sbi/neural_nets/estimators/mixture_density_estimator.py b/sbi/neural_nets/estimators/mixture_density_estimator.py
index dbbc608b7..310a35965 100644
--- a/sbi/neural_nets/estimators/mixture_density_estimator.py
+++ b/sbi/neural_nets/estimators/mixture_density_estimator.py
@@ -17,6 +17,7 @@
 import numpy as np
 import torch
 from torch import Tensor, nn
+from torch.distributions import Transform as TorchTransform
 from torch.nn import functional as F
 
 from sbi.neural_nets.estimators.base import ConditionalDensityEstimator
@@ -337,6 +338,7 @@ def __init__(
         condition_shape: torch.Size,
         embedding_net: Optional[nn.Module] = None,
         transform_input: Optional[Tensor] = None,
+        input_transform: Optional[TorchTransform] = None,
     ) -> None:
         """Initialize the mixture density estimator.
 
@@ -354,12 +356,22 @@ def __init__(
                 evaluation, and samples are inverse transformed as:
                 x = z * scale + shift.
                 This is used for z-scoring inputs to improve numerical stability.
+            input_transform: Optional PyTorch Transform for general bijective
+                transformation of inputs (e.g., for transform_to_unconstrained).
+                If provided, inputs are transformed as z = transform(x) and samples
+                are inverse transformed as x = transform.inv(z). Mutually exclusive
+                with transform_input.
         """
         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 input_transform is not None:
+            raise ValueError(
+                "Only one of transform_input and input_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 +384,9 @@ def __init__(
                     f"MDN context_features ({net.context_features})"
                 )
 
+        # Store general bijective input transform (e.g., for transform_to_unconstrained)
+        self._input_transform = input_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:
@@ -399,28 +414,29 @@ def embedding_net(self) -> nn.Module:
 
     @property
     def has_input_transform(self) -> bool:
-        """Whether input z-score transform is enabled."""
-        return self._transform_shift is not None
+        """Whether input transform is enabled (z-score or general transform)."""
+        return self._input_transform is not None or self._transform_shift is not None
 
     def _transform_input(self, input: Tensor) -> Tensor:
-        """Apply z-score transform to input: z = (x - shift) / scale."""
-        if not self.has_input_transform:
+        """Apply input transform: z = transform(x) or z = (x - shift) / scale."""
+        if self._input_transform is not None:
+            return self._input_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:
+        """Apply inverse input transform: x = transform.inv(z) or x = z * scale + shift."""
+        if self._input_transform is not None:
+            return self._input_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:
-        """Compute log determinant of Jacobian for the forward z-score transform.
-
-        For the forward affine transform z = (x - shift) / scale:
-            - The Jacobian matrix is: dz/dx = diag(1/scale)
-            - The determinant is: |det(dz/dx)| = prod(1/scale)
-            - The log determinant is: log|det(dz/dx)| = -sum(log(scale))
+    def _log_det_jacobian_forward(
+        self, input: Tensor, transformed_input: Tensor
+    ) -> Tensor:
+        """Compute log determinant of Jacobian for the forward input transform.
 
         Change of Variables Formula:
             When we have a density p_z(z) and want p_x(x), we use:
@@ -429,19 +445,25 @@ def _log_det_jacobian_forward(self, input: Tensor) -> Tensor:
             In log space:
             log p_x(x) = log p_z(z(x)) + log|det(dz/dx)|
 
-            Since log|det(dz/dx)| = -sum(log(scale)), we have:
-            log p_x(x) = log p_z(z) - sum(log(scale))
+        For the forward affine transform z = (x - shift) / scale:
+            log|det(dz/dx)| = -sum(log(scale))
 
-        This method returns log|det(dz/dx)| = -sum(log(scale)), which should
-        be ADDED to log p_z(z) to get log p_x(x).
+        For a general bijective transform:
+            log|det(dz/dx)| = transform.log_abs_det_jacobian(x, z).sum(dim=-1)
 
         Args:
-            input: Input tensor, used to determine device and dtype.
+            input: Input in original space.
+            transformed_input: Input after forward transform.
 
         Returns:
-            Log determinant of forward Jacobian (scalar), on the same device as input.
+            Log determinant of forward Jacobian. Shape matches batch dimensions.
         """
-        if not self.has_input_transform:
+        if self._input_transform is not None:
+            jac = self._input_transform.log_abs_det_jacobian(
+                input, transformed_input
+            )
+            return jac.sum(dim=-1)
+        if self._transform_shift is None:
             return torch.zeros(1, device=input.device, dtype=input.dtype).squeeze()
         return -torch.log(self._transform_scale).sum()
 
@@ -476,9 +498,10 @@ def log_prob(self, input: Tensor, condition: Tensor, **kwargs) -> Tensor:
 
         # MoG.log_prob handles (sample_dim, batch_dim, dim) input
         # 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)
diff --git a/sbi/neural_nets/net_builders/mdn.py b/sbi/neural_nets/net_builders/mdn.py
index 294fa7131..2dc1c88bb 100644
--- a/sbi/neural_nets/net_builders/mdn.py
+++ b/sbi/neural_nets/net_builders/mdn.py
@@ -12,7 +12,7 @@
 )
 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,
@@ -41,6 +41,8 @@ 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.
@@ -53,19 +55,22 @@ def build_mdn(
         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
+    input_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":
+        x_dist = kwargs.get("x_dist", None)
+        if x_dist is None:
+            raise ValueError(
+                "x_dist must be provided when z_score_x='transform_to_unconstrained'."
+            )
+        input_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 +98,7 @@ def build_mdn(
         condition_shape=batch_y[0].shape,
         embedding_net=embedding_net,
         transform_input=transform_input,
+        input_transform=input_transform,
     )
 
     return estimator
diff --git a/tests/sbiutils_test.py b/tests/sbiutils_test.py
index fcde9cab9..b3ad2fbd4 100644
--- a/tests/sbiutils_test.py
+++ b/tests/sbiutils_test.py
@@ -497,7 +497,12 @@ 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"):
+        # MDN supports it (via input_transform), so it's excluded from this check.
+        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 +554,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,

From 1df53a7336255d26d1b234a54143278375a5e816 Mon Sep 17 00:00:00 2001
From: BHARATH0153 <chandrabharath463@gmail.com>
Date: Wed, 24 Jun 2026 00:18:16 +0530
Subject: [PATCH 2/3] docs: simplify docstrings in MixtureDensityEstimator

---
 .../estimators/mixture_density_estimator.py   | 29 +++----------------
 1 file changed, 4 insertions(+), 25 deletions(-)

diff --git a/sbi/neural_nets/estimators/mixture_density_estimator.py b/sbi/neural_nets/estimators/mixture_density_estimator.py
index 310a35965..439bf73bd 100644
--- a/sbi/neural_nets/estimators/mixture_density_estimator.py
+++ b/sbi/neural_nets/estimators/mixture_density_estimator.py
@@ -357,10 +357,7 @@ def __init__(
                 x = z * scale + shift.
                 This is used for z-scoring inputs to improve numerical stability.
             input_transform: Optional PyTorch Transform for general bijective
-                transformation of inputs (e.g., for transform_to_unconstrained).
-                If provided, inputs are transformed as z = transform(x) and samples
-                are inverse transformed as x = transform.inv(z). Mutually exclusive
-                with transform_input.
+                transformation of inputs. Mutually exclusive with transform_input.
         """
         super().__init__(net, input_shape, condition_shape)
         self._embedding_net = (
@@ -384,7 +381,6 @@ def __init__(
                     f"MDN context_features ({net.context_features})"
                 )
 
-        # Store general bijective input transform (e.g., for transform_to_unconstrained)
         self._input_transform = input_transform
 
         # Store z-score transform parameters as buffers (not trained, moved with model)
@@ -436,27 +432,10 @@ def _inverse_transform_input(self, z: Tensor) -> Tensor:
     def _log_det_jacobian_forward(
         self, input: Tensor, transformed_input: Tensor
     ) -> Tensor:
-        """Compute log determinant of Jacobian for the forward input transform.
+        """Log determinant of the forward input transform Jacobian.
 
-        Change of Variables Formula:
-            When we have a density p_z(z) and want p_x(x), we use:
-            p_x(x) = p_z(z(x)) * |det(dz/dx)|
-
-            In log space:
-            log p_x(x) = log p_z(z(x)) + log|det(dz/dx)|
-
-        For the forward affine transform z = (x - shift) / scale:
-            log|det(dz/dx)| = -sum(log(scale))
-
-        For a general bijective transform:
-            log|det(dz/dx)| = transform.log_abs_det_jacobian(x, z).sum(dim=-1)
-
-        Args:
-            input: Input in original space.
-            transformed_input: Input after forward transform.
-
-        Returns:
-            Log determinant of forward Jacobian. Shape matches batch dimensions.
+        For the affine z-score transform: -sum(log(scale)).
+        For a general transform: transform.log_abs_det_jacobian(x, z).sum(dim=-1).
         """
         if self._input_transform is not None:
             jac = self._input_transform.log_abs_det_jacobian(

From 16d884e46dc7f2cfb3be43e493be61f60f9f8c08 Mon Sep 17 00:00:00 2001
From: BHARATH0153 <chandrabharath463@gmail.com>
Date: Wed, 24 Jun 2026 00:41:51 +0530
Subject: [PATCH 3/3] feat(vfpe): implement transform_to_unconstrained
 z-scoring for vector field estimators

Closes #1885 (vector field portion)

- Remove assert_transform_to_unconstrained_supported guard from
  build_vector_field_estimator
- Add x_dist parameter; build mcmc_transform when z_score_x is
  transform_to_unconstrained
- Add _input_transform attribute and input_transform property to
  ConditionalVectorFieldEstimator base class
- Update FlowMatchingEstimator.ode_fn() and .loss() to apply the
  transform before forward/loss computation
- Update score estimators (ConditionalScore, VP, SubVP, VE) same
  pattern in ode_fn() and loss()
- Update VectorFieldBasedPotential.__call__() to transform theta,
  add log|det J| correction
- Update VectorFieldPosterior.sample_via_ode() and
  _sample_via_diffusion() to inverse-transform samples
---
 .../posteriors/vector_field_posterior.py      | 14 +++++-
 .../potentials/vector_field_potential.py      | 25 +++++++++-
 sbi/neural_nets/estimators/base.py            | 11 ++++
 .../estimators/flowmatching_estimator.py      | 42 ++++++++++++++--
 sbi/neural_nets/estimators/score_estimator.py | 35 ++++++++++++-
 .../net_builders/vector_field_nets.py         | 50 +++++++++++++++----
 6 files changed, 157 insertions(+), 20 deletions(-)

diff --git a/sbi/inference/posteriors/vector_field_posterior.py b/sbi/inference/posteriors/vector_field_posterior.py
index 20c1990f0..ea2b6b9e9 100644
--- a/sbi/inference/posteriors/vector_field_posterior.py
+++ b/sbi/inference/posteriors/vector_field_posterior.py
@@ -395,6 +395,11 @@ def _sample_via_diffusion(
         # Concatenate all batches and ensure we return exactly the requested number
         samples = torch.cat(all_samples, dim=0)[:total_samples_needed]
 
+        # Inverse-transform samples from unconstrained back to original space
+        input_transform = self.vector_field_estimator.input_transform
+        if input_transform is not None:
+            samples = input_transform.inv(samples)
+
         if torch.isnan(samples).all():
             raise RuntimeError(
                 "All samples NaN after diffusion sampling. "
@@ -412,7 +417,9 @@ def sample_via_ode(
         Return samples from posterior distribution with probability flow ODE.
 
         This builds the probability flow ODE and then samples from the corresponding
-        flow.
+        flow. When ``input_transform`` is set on the estimator, the ODE integrates
+        in unconstrained space and samples are inverse-transformed back to the
+        original constrained space.
 
         Args:
             sample_shape: The shape of the samples to be returned.
@@ -429,6 +436,11 @@ def sample_via_ode(
             torch.Size((num_samples,))
         )
 
+        # Inverse-transform samples from unconstrained back to original space
+        input_transform = self.vector_field_estimator.input_transform
+        if input_transform is not None:
+            samples = input_transform.inv(samples)
+
         return samples
 
     def log_prob(
diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py
index b7ac4ba45..39c2c8579 100644
--- a/sbi/inference/potentials/vector_field_potential.py
+++ b/sbi/inference/potentials/vector_field_potential.py
@@ -136,6 +136,10 @@ def __call__(
         """
         Return the potential (posterior log prob) via probability flow ODE.
 
+        When ``input_transform`` is set on the estimator, theta is transformed
+        to unconstrained space before the flow evaluation, and the log-det-Jacobian
+        of the transform is added to account for the change of variables.
+
         Args:
             theta: The parameters at which to evaluate the potential.
             track_gradients: Whether to track gradients. Default is False.
@@ -155,6 +159,8 @@ def __call__(
         )
         self.vector_field_estimator.eval()
 
+        input_transform = self.vector_field_estimator.input_transform
+
         with torch.set_grad_enabled(track_gradients):
             if self.x_is_iid:
                 assert self.prior is not None, (
@@ -164,10 +170,14 @@ def __call__(
                     "Flows for each iid x are required for evaluating log_prob."
                 )
                 num_iid = self.x_o.shape[0]  # number of iid samples
+                if input_transform is not None:
+                    theta_z = input_transform(theta_density_estimator)
+                else:
+                    theta_z = theta_density_estimator
                 iid_posteriors_prob = torch.sum(
                     torch.stack(
                         [
-                            flow.log_prob(theta_density_estimator).squeeze(-1)
+                            flow.log_prob(theta_z).squeeze(-1)
                             for flow in self.flows
                         ],
                         dim=0,
@@ -180,7 +190,18 @@ def __call__(
                     theta_density_estimator
                 ).squeeze(-1)
             else:
-                log_probs = self.flow.log_prob(theta_density_estimator).squeeze(-1)
+                if input_transform is not None:
+                    theta_z = input_transform(theta_density_estimator)
+                else:
+                    theta_z = theta_density_estimator
+                log_probs = self.flow.log_prob(theta_z).squeeze(-1)
+
+            # Add log-det-Jacobian correction for change of variables
+            if input_transform is not None:
+                log_probs = log_probs + input_transform.log_abs_det_jacobian(
+                    theta_density_estimator, theta_z
+                ).sum(dim=-1)
+
             # Force probability to be zero outside prior support.
             in_prior_support = within_support(self.prior, theta)
 
diff --git a/sbi/neural_nets/estimators/base.py b/sbi/neural_nets/estimators/base.py
index 68eb7f996..a9abda0e2 100644
--- a/sbi/neural_nets/estimators/base.py
+++ b/sbi/neural_nets/estimators/base.py
@@ -6,6 +6,7 @@
 
 import torch
 from torch import Tensor, nn
+from torch.distributions import Transform as TorchTransform
 
 ConditionalEstimatorType = TypeVar(
     'ConditionalEstimatorType',
@@ -350,6 +351,7 @@ def __init__(
         embedding_net: Optional[nn.Module] = None,
         mean_base: Union[float, Tensor] = 0.0,
         std_base: Union[float, Tensor] = 1.0,
+        input_transform: Optional[TorchTransform] = None,
     ) -> None:
         r"""Base class for vector field estimators.
 
@@ -364,6 +366,9 @@ def __init__(
                 condition.
             mean_base: Mean of the base distribution.
             std_base: Standard deviation of the base distribution.
+            input_transform: Optional transform mapping constrained -> unconstrained
+                space. When set, ``forward()`` operates in unconstrained space and
+                ``ode_fn()`` applies this transform before calling ``forward()``.
         """
         super().__init__(input_shape, condition_shape)
         self.net = net
@@ -385,12 +390,18 @@ def __init__(
         self._embedding_net = (
             embedding_net if embedding_net is not None else nn.Identity()
         )
+        self._input_transform = input_transform
 
     @property
     def embedding_net(self) -> nn.Module:
         r"""Return the embedding network if it exists."""
         return self._embedding_net
 
+    @property
+    def input_transform(self) -> Optional[TorchTransform]:
+        r"""Return the input transform (constrained -> unconstrained) if set."""
+        return self._input_transform
+
     @abstractmethod
     def forward(self, input: Tensor, condition: Tensor, **kwargs) -> Tensor:
         r"""Forward pass of the score estimator.
diff --git a/sbi/neural_nets/estimators/flowmatching_estimator.py b/sbi/neural_nets/estimators/flowmatching_estimator.py
index ad76d1e80..3f6d614d0 100644
--- a/sbi/neural_nets/estimators/flowmatching_estimator.py
+++ b/sbi/neural_nets/estimators/flowmatching_estimator.py
@@ -7,6 +7,7 @@
 import torch
 import torch.nn as nn
 from torch import Tensor
+from torch.distributions import Transform as TorchTransform
 
 from sbi.neural_nets.estimators.base import ConditionalVectorFieldEstimator
 from sbi.utils.vector_field_utils import VectorFieldNet
@@ -65,6 +66,7 @@ def __init__(
         mean_0: float = 0.0,
         std_0: float = 1.0,
         gaussian_baseline: bool = False,
+        input_transform: Optional[TorchTransform] = None,
         **kwargs,
     ) -> None:
         r"""Creates a vector field estimator for Flow Matching.
@@ -81,8 +83,12 @@ def __init__(
             std_0: Std of the data distribution at t=0 (used for time-dependent
                 z-scoring).
             gaussian_baseline: If True, use analytical Gaussian baseline velocity
-                derived from Bayes' rule: v = factor * (x - μ_true) - mean.
-                The network then only learns the residual. Default: False.
+                derived from Bayes' rule. The network then only learns the residual.
+                Default: False.
+            input_transform: Optional transform mapping constrained -> unconstrained
+                space. When set, ``ode_fn()`` applies this transform before calling
+                ``forward()``, so the ODE integrates in unconstrained space, and
+                ``loss()`` transforms training data accordingly.
         """
 
         if "num_freqs" in kwargs:
@@ -99,6 +105,7 @@ def __init__(
             input_shape=input_shape,
             condition_shape=condition_shape,
             embedding_net=embedding_net,
+            input_transform=input_transform,
         )
         self.noise_scale = noise_scale
         self.gaussian_baseline = gaussian_baseline
@@ -196,7 +203,13 @@ def _compute_velocity_baseline(self, input: Tensor, time: Tensor) -> Tensor:
     def forward(self, input: Tensor, condition: Tensor, time: Tensor) -> Tensor:
         """Forward pass of the FlowMatchingEstimator.
 
-        Returns velocity in ORIGINAL SPACE for ODE integration.
+        Returns velocity in internal space (unconstrained z-space when
+        ``input_transform`` is set).
+
+        Note:
+            This method does **not** apply ``input_transform`` internally. The
+            caller (e.g. ``ode_fn()``, ``loss()``) is responsible for transforming
+            inputs if needed.
 
         Args:
             input: Inputs to evaluate the vector field on of shape
@@ -207,7 +220,7 @@ def forward(self, input: Tensor, condition: Tensor, time: Tensor) -> Tensor:
                 `(batch_dim_time, *event_shape_time)`.
 
         Returns:
-            The estimated vector field in original space.
+            The estimated vector field in internal space.
         """
         # Continue with standard processing (broadcast shapes etc.)
         batch_shape_input = input.shape[: -len(self.input_shape)]
@@ -274,6 +287,9 @@ def loss(
             \theta_t = t \cdot \theta_1 + (1 - t) \cdot \theta_0,
             \left[ \| v(\theta_t, t; x_o = x_o) - (\theta_1 - \theta_0) \|^2 \right]
 
+        When ``input_transform`` is set, the input is first transformed to
+        unconstrained space, and all computations happen in that space.
+
         Args:
             input: Parameters (:math:`\theta_0`).
             condition: Observed data (:math:`x_o`).
@@ -288,12 +304,16 @@ def loss(
             times = torch.rand(input.shape[:-1], device=input.device, dtype=input.dtype)
         times_ = times[..., None]
 
+        # Transform to unconstrained space if needed
+        if self._input_transform is not None:
+            input = self._input_transform(input)
+
         # Sample from probability path at time t
         # θ_t = (1-t) * θ_data + t * θ_noise, where θ_noise ~ N(0,1)
         theta_1 = torch.randn_like(input)
         theta_t = (1 - times_) * input + (times_ + self.noise_scale) * theta_1
 
-        # Target vector field in original space
+        # Target vector field in internal space
         vector_field = theta_1 - input
 
         # Embed condition
@@ -348,6 +368,10 @@ def ode_fn(self, input: Tensor, condition: Tensor, times: Tensor) -> Tensor:
         flow matching neural network (see Equation 1 in [1]_ with added
         conditioning on :math:`x_o`).
 
+        When ``input_transform`` is set, this method transforms the input to
+        unconstrained space before calling ``forward()``, so the ODE integrates
+        in the unconstrained space.
+
         Args:
             input: :math:`\theta_t`.
             condition: Conditioning variable :math:`x_o`.
@@ -357,6 +381,8 @@ def ode_fn(self, input: Tensor, condition: Tensor, times: Tensor) -> Tensor:
             Estimated vector field :math:`v(\theta_t, t; x_o)`.
             The shape is the same as the input.
         """
+        if self._input_transform is not None:
+            input = self._input_transform(input)
         return self.forward(input, condition, times)
 
     def score(self, input: Tensor, condition: Tensor, t: Tensor) -> Tensor:
@@ -372,6 +398,12 @@ def score(self, input: Tensor, condition: Tensor, t: Tensor) -> Tensor:
         :math:`\nabla_{\theta_t} \log p(\theta_t | x_o) =
             (- (1 - t) v(\theta_t, t; x_o) - \theta_0 ) / (t + \sigma_{min})`.
 
+        Note:
+            This method does **not** apply ``input_transform`` internally. When
+            ``input_transform`` is set, this method expects input in unconstrained
+            (z) space. Use ``ode_fn()`` for automatic transform handling during
+            ODE/SDE sampling.
+
         Args:
             input: variable whose distribution is estimated.
             condition: Conditioning variable.
diff --git a/sbi/neural_nets/estimators/score_estimator.py b/sbi/neural_nets/estimators/score_estimator.py
index 87bead1ee..e25e1ba29 100644
--- a/sbi/neural_nets/estimators/score_estimator.py
+++ b/sbi/neural_nets/estimators/score_estimator.py
@@ -7,6 +7,7 @@
 
 import torch
 from torch import Tensor, nn
+from torch.distributions import Transform as TorchTransform
 
 from sbi.neural_nets.estimators.base import ConditionalVectorFieldEstimator
 from sbi.utils.vector_field_utils import VectorFieldNet
@@ -73,6 +74,7 @@ def __init__(
         std_0: Union[Tensor, float] = 1.0,
         t_min: float = 1e-3,
         t_max: float = 1.0,
+        input_transform: Optional[TorchTransform] = None,
     ) -> None:
         r"""Score estimator class that estimates the
         conditional score function, i.e.,
@@ -94,6 +96,10 @@ def __init__(
             std_0: Starting standard deviation of the target distribution.
             t_min: Minimum time for diffusion (0 can be numerically unstable).
             t_max: Maximum time for diffusion.
+            input_transform: Optional transform mapping constrained -> unconstrained
+                space. When set, ``ode_fn()`` applies this transform before calling
+                ``forward()``, so the ODE integrates in unconstrained space, and
+                ``loss()`` transforms training data accordingly.
         """
 
         # Starting mean and std of the target distribution (otherwise assumes 0,1).
@@ -112,6 +118,7 @@ def __init__(
             embedding_net=embedding_net,
             t_min=t_min,
             t_max=t_max,
+            input_transform=input_transform,
         )
 
         # Min/max values for noise variance beta
@@ -141,6 +148,11 @@ def forward(self, input: Tensor, condition: Tensor, time: Tensor) -> Tensor:
         network to compute the conditional score
         at a given time.
 
+        Note:
+            This method does **not** apply ``input_transform`` internally. The
+            caller (e.g. ``ode_fn()``, ``loss()``) is responsible for transforming
+            inputs if needed.
+
         Args:
             input: Inputs to evaluate the score estimator on of shape
                     `(sample_dim_input, batch_dim_input, *event_shape_input)`.
@@ -232,6 +244,9 @@ def loss(
         which is equivalent to predicting the (scaled) Gaussian noise added to the
         input.
 
+        When ``input_transform`` is set, the input is first transformed to
+        unconstrained space, and all computations happen in that space.
+
         Args:
             input: Input variable i.e. theta.
             condition: Conditioning variable.
@@ -253,6 +268,10 @@ def loss(
             times = self.train_schedule(input.shape[0])
         times = times.to(input.device)
 
+        # Transform to unconstrained space if needed
+        if self._input_transform is not None:
+            input = self._input_transform(input)
+
         # Sample noise.
         eps = torch.randn_like(input)
 
@@ -501,6 +520,10 @@ def ode_fn(self, input: Tensor, condition: Tensor, times: Tensor) -> Tensor:
 
         For reference, see Equation 13 in [1]_.
 
+        When ``input_transform`` is set, this method transforms the input to
+        unconstrained space before calling ``forward()`` / ``score()``, so the
+        ODE integrates in the unconstrained space.
+
         Args:
             input: variable whose distribution is estimated.
             condition: Conditioning variable.
@@ -509,6 +532,8 @@ def ode_fn(self, input: Tensor, condition: Tensor, times: Tensor) -> Tensor:
         Returns:
             ODE flow function value at a given time.
         """
+        if self._input_transform is not None:
+            input = self._input_transform(input)
         score = self.score(input=input, condition=condition, t=times)
         f = self.drift_fn(input, times)
         g = self.diffusion_fn(input, times)
@@ -548,10 +573,11 @@ def __init__(
         weight_fn: Union[str, Callable] = "max_likelihood",
         beta_min: float = 0.01,
         beta_max: float = 10.0,
-        mean_0: Union[Tensor, float] = 0.0,
-        std_0: Union[Tensor, float] = 1.0,
+        mean_0: float = 0.0,
+        std_0: float = 1.0,
         t_min: float = 1e-3,
         t_max: float = 1.0,
+        input_transform: Optional[TorchTransform] = None,
     ) -> None:
         super().__init__(
             net,
@@ -565,6 +591,7 @@ def __init__(
             beta_max=beta_max,
             t_min=t_min,
             t_max=t_max,
+            input_transform=input_transform,
         )
 
     def mean_t_fn(self, times: Tensor) -> Tensor:
@@ -665,6 +692,7 @@ def __init__(
         std_0: float = 1.0,
         t_min: float = 1e-2,
         t_max: float = 1.0,
+        input_transform: Optional[TorchTransform] = None,
     ) -> None:
         super().__init__(
             net,
@@ -678,6 +706,7 @@ def __init__(
             std_0=std_0,
             t_min=t_min,
             t_max=t_max,
+            input_transform=input_transform,
         )
 
     def mean_t_fn(self, times: Tensor) -> Tensor:
@@ -811,6 +840,7 @@ def __init__(
         lognormal_mean: float = -1.2,
         lognormal_std: float = 1.2,
         power_law_exponent: float = 7.0,
+        input_transform: Optional[TorchTransform] = None,
     ) -> None:
         # Validate sigma bounds (required for VE SDE math and log computations).
         if sigma_min <= 0:
@@ -863,6 +893,7 @@ def __init__(
             std_0=std_0,
             t_min=t_min,
             t_max=t_max,
+            input_transform=input_transform,
         )
 
         self._warn_on_inappropriate_config()
diff --git a/sbi/neural_nets/net_builders/vector_field_nets.py b/sbi/neural_nets/net_builders/vector_field_nets.py
index 3585007c7..dfbbcbeb4 100644
--- a/sbi/neural_nets/net_builders/vector_field_nets.py
+++ b/sbi/neural_nets/net_builders/vector_field_nets.py
@@ -17,9 +17,12 @@
     VPScoreEstimator,
 )
 from sbi.neural_nets.net_builders.estimator_configs import _EstimatorConfigBase
+from torch.distributions import Distribution
+
+from sbi.sbi_types import TorchTransform
 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,
@@ -58,6 +61,7 @@ class _VectorFieldBaseConfig(_EstimatorConfigBase):
     # passed through **kwargs at the factory level
     net: Optional[Any] = None
     z_score_x: Optional[Any] = None
+    x_dist: Optional[Any] = None
     z_score_y: Optional[Any] = None
     hidden_features: Optional[Any] = None
     num_layers: Optional[int] = None
@@ -127,6 +131,7 @@ def build_vector_field_estimator(
         VectorFieldNet,
     ] = "mlp",
     gaussian_baseline: bool = False,
+    x_dist: Optional[Distribution] = None,
     **kwargs,
 ) -> Union[FlowMatchingEstimator, ConditionalScoreEstimator]:
     """Builds a vector field estimator (flow matching or score matching) with the given
@@ -136,8 +141,16 @@ def build_vector_field_estimator(
         batch_x: Batch of xs, used to infer dimensionality.
         batch_y: Batch of ys, used to infer dimensionality.
         estimator_type: Type of estimator to build, either "flow" or "score".
-        z_score_x: Whether to z-score xs passing into the network.
-        z_score_y: Whether to z-score ys passing into the network.
+        z_score_x: Whether to z-score xs passing into the network, can be one of:
+            - `none`, or None: do not z-score.
+            - `independent`: z-score each dimension independently.
+            - `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.
         embedding_net: Embedding network for batch_y.
         sde_type: SDE type for score estimator, one of "vp", "subvp", or "ve".
         hidden_features: Number of hidden features in each layer (for MLP) or dimension
@@ -152,6 +165,8 @@ def build_vector_field_estimator(
         gaussian_baseline: If True, use analytical Gaussian baseline velocity
             derived from Bayes' rule. The network then only learns the residual.
             Only used when estimator_type="flow". Defaults to False.
+        x_dist: Distribution used for unconstrained transformation when
+            z_score_x="transform_to_unconstrained". Must have a `.support` attribute.
         **kwargs: Additional arguments forwarded to the estimator and network
             constructors.  Valid keys are defined by ``ScoreEstimatorConfig``
             and ``FlowEstimatorConfig``; validation happens in the upstream
@@ -163,12 +178,20 @@ def build_vector_field_estimator(
     """
     # Check inputs and device
     check_data_device(batch_x, batch_y)
-    assert_transform_to_unconstrained_supported(
-        z_score_x,
-        "build_vector_field_estimator",
-        "Vector field estimators (flow matching / score matching) do not implement "
-        "it; use one of 'none', 'independent', or 'structured' instead.",
-    )
+
+    # Build input transform for transform_to_unconstrained
+    input_transform: Optional[TorchTransform] = None
+    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'."
+            )
+        if not hasattr(x_dist, "support"):
+            raise ValueError(
+                "x_dist requires a `.support` attribute when using "
+                "z_score_x='transform_to_unconstrained'."
+            )
+        input_transform = mcmc_transform(x_dist, device=batch_x.device)
 
     # Build network if not provided
     if net == "mlp":
@@ -227,7 +250,12 @@ def build_vector_field_estimator(
 
     # Z-score setup
     z_score_x_bool, structured_x = z_score_parser(z_score_x)
-    if z_score_x_bool:
+    if input_transform is not None:
+        # Transform data to unconstrained space, then compute stats for
+        # time-dependent z-scoring in that space.
+        batch_x_transformed = input_transform(batch_x)
+        mean_0, std_0 = z_standardization(batch_x_transformed, structured_x)
+    elif z_score_x_bool:
         mean_0, std_0 = z_standardization(batch_x, structured_x)
     else:
         mean_0, std_0 = 0, 1
@@ -248,6 +276,7 @@ def build_vector_field_estimator(
             mean_0=mean_0,
             std_0=std_0,
             gaussian_baseline=gaussian_baseline,
+            input_transform=input_transform,
         )
     elif estimator_type == "score":
         # Choose the appropriate score estimator based on SDE type
@@ -286,6 +315,7 @@ def build_vector_field_estimator(
             embedding_net=embedding_net_y,
             mean_0=mean_0,
             std_0=std_0,
+            input_transform=input_transform,
             **estimator_kwargs,
         )
     else: