diff --git a/sbi/inference/posteriors/vector_field_posterior.py b/sbi/inference/posteriors/vector_field_posterior.py
index 20c1990f0..d638bb01d 100644
--- a/sbi/inference/posteriors/vector_field_posterior.py
+++ b/sbi/inference/posteriors/vector_field_posterior.py
@@ -235,6 +235,18 @@ def sample(
x = self._x_else_default_x(x)
x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
is_iid = x.shape[0] > 1
+ # Composed standardization integrates the SDE/ODE in z-space; the iid and
+ # guidance score combinations evaluate the (theta-space) prior on the z-space
+ # state, which is not yet transformed. Guard explicitly instead of returning
+ # silently-wrong samples. (Single-observation sampling is fully supported.)
+ if self.vector_field_estimator._compose_standardization and (
+ is_iid or guidance_method is not None
+ ):
+ raise NotImplementedError(
+ "compose_standardization does not yet support iid (x with batch>1) or "
+ "guided sampling. Use a single observation, or disable "
+ "compose_standardization."
+ )
self.potential_fn.set_x(
x,
x_is_iid=is_iid,
@@ -401,6 +413,13 @@ def _sample_via_diffusion(
"This may indicate numerical instability in the vector field."
)
+ # Composed standardization (opt-in): the estimator samples in z-space.
+ # Unstandardize theta = shift + scale * z before returning so that
+ # rejection (within_support on the prior) and the final samples are in
+ # original theta space. No-op when the flag is off.
+ if self.vector_field_estimator._compose_standardization:
+ samples = self.vector_field_estimator.from_z(samples)
+
return samples
def sample_via_ode(
@@ -429,6 +448,13 @@ def sample_via_ode(
torch.Size((num_samples,))
)
+ # Composed standardization (opt-in): the ODE samples in z-space.
+ # Unstandardize theta = shift + scale * z before returning so that
+ # rejection (within_support on the prior) and the final samples are in
+ # original theta space. No-op when the flag is off.
+ if self.vector_field_estimator._compose_standardization:
+ samples = self.vector_field_estimator.from_z(samples)
+
return samples
def log_prob(
@@ -457,6 +483,15 @@ def log_prob(
x = self._x_else_default_x(x)
x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
is_iid = x.shape[0] > 1
+ # See sample(): the iid score combination evaluates the theta-space prior on the
+ # z-space state under composed standardization. Guard rather than return wrong
+ # values; single-observation log_prob is exact (affine Jacobian correction).
+ if self.vector_field_estimator._compose_standardization and is_iid:
+ raise NotImplementedError(
+ "compose_standardization does not yet support iid (x with batch>1) "
+ "log_prob. Use a single observation, or disable "
+ "compose_standardization."
+ )
self.potential_fn.set_x(x, x_is_iid=is_iid, **(ode_kwargs or {}))
theta = ensure_theta_batched(torch.as_tensor(theta))
@@ -516,6 +551,17 @@ def sample_batched(
Returns:
Samples from the posteriors of shape (*sample_shape, B, *input_shape)
"""
+ # Composed standardization integrates the SDE/ODE in z-space. Batched
+ # sampling (and the AutoGaussCorrectedScoreFn it feeds, which builds
+ # theta-space precisions) is not yet transformed for z-space and would
+ # return silently-wrong samples. Guard explicitly. (Single-observation
+ # sampling via sample() is fully supported.)
+ if self.vector_field_estimator._compose_standardization:
+ raise NotImplementedError(
+ "compose_standardization does not yet support sample_batched "
+ "(batched / multi-observation sampling). Use a single observation "
+ "via sample(), or disable compose_standardization."
+ )
num_samples = torch.Size(sample_shape).numel()
x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
condition_dim = len(self.vector_field_estimator.condition_shape)
@@ -647,6 +693,13 @@ def map(
Returns:
The MAP estimate.
"""
+ if self.vector_field_estimator._compose_standardization:
+ raise NotImplementedError(
+ "MAP is not yet supported with compose_standardization. "
+ "The potential gradient is computed in standardized z-space, so "
+ "gradient ascent in theta-space would be incorrect."
+ )
+
if x is not None:
raise ValueError(
"Passing `x` directly to `.map()` has been deprecated."
diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py
index b7ac4ba45..abd509ae9 100644
--- a/sbi/inference/potentials/vector_field_potential.py
+++ b/sbi/inference/potentials/vector_field_potential.py
@@ -118,6 +118,25 @@ def set_x(
`IIDScoreFunction`.
ode_kwargs: Additional keyword arguments for the neural ODE.
"""
+ # Composed standardization runs the SDE/ODE in z-space; the iid and guided
+ # score combinations evaluate the (theta-space) prior on the z-space state,
+ # which is not transformed. Guard at this chokepoint so the broken combination
+ # cannot be reached even via direct potential use (not just via
+ # VectorFieldPosterior.sample/log_prob).
+ if bool(
+ getattr(self.vector_field_estimator, "_compose_standardization", False)
+ ):
+ if x_is_iid:
+ raise NotImplementedError(
+ "compose_standardization does not yet support iid (x with "
+ "batch>1). Use a single observation, or disable "
+ "compose_standardization."
+ )
+ if guidance_method is not None:
+ raise NotImplementedError(
+ "compose_standardization does not yet support guided sampling. "
+ "Disable guidance, or disable compose_standardization."
+ )
super().set_x(x_o, x_is_iid)
self.iid_method = iid_method or self.iid_method
self.iid_params = iid_params
@@ -155,8 +174,40 @@ def __call__(
)
self.vector_field_estimator.eval()
+ # Composed standardization (opt-in): the flow operates in z-space. Map the
+ # incoming theta -> z for the flow, and correct the log-prob with the
+ # affine Jacobian: log p_theta(theta) = log p_z(z) - sum(log scale).
+ # within_support below still uses the ORIGINAL theta. No-op when off.
+ compose = self.vector_field_estimator._compose_standardization
+ if compose:
+ flow_input = self.vector_field_estimator.to_z(theta_density_estimator)
+ log_abs_det = self.vector_field_estimator.log_abs_det()
+ else:
+ flow_input = theta_density_estimator
+
with torch.set_grad_enabled(track_gradients):
if self.x_is_iid:
+ if compose:
+ # Backstop for a directly-forced `_x_is_iid`: compose + iid is
+ # unreachable through the public API (guarded in `set_x`).
+ # Reaching here means `_x_is_iid` was set directly, and continuing
+ # would return a partially-incorrect log-prob -- the affine
+ # Jacobian is applied but the prior is not yet expressed in
+ # z-space. Raise rather than return quiet wrong numbers, matching
+ # the set_x guards.
+ raise NotImplementedError(
+ "compose_standardization does not yet support iid (x with "
+ "batch>1) log_prob. This path is guarded in set_x; reaching "
+ "it means _x_is_iid was set directly, which would return a "
+ "partially-incorrect log-prob (affine Jacobian applied, prior "
+ "not yet expressed in z). Use a single observation, or disable "
+ "compose_standardization."
+ )
+ # future iid support: lift this backstop AND the set_x guard
+ # once the prior is transformed to z (see
+ # followup/compose-iid-guidance-map). Retained z-space affine-
+ # Jacobian correction to restore then:
+ # iid_posteriors_prob -= num_iid * log_abs_det
assert self.prior is not None, (
"Prior is required for evaluating log_prob with iid observations."
)
@@ -166,10 +217,7 @@ def __call__(
num_iid = self.x_o.shape[0] # number of iid samples
iid_posteriors_prob = torch.sum(
torch.stack(
- [
- flow.log_prob(theta_density_estimator).squeeze(-1)
- for flow in self.flows
- ],
+ [flow.log_prob(flow_input).squeeze(-1) for flow in self.flows],
dim=0,
),
dim=0,
@@ -180,7 +228,9 @@ def __call__(
theta_density_estimator
).squeeze(-1)
else:
- log_probs = self.flow.log_prob(theta_density_estimator).squeeze(-1)
+ log_probs = self.flow.log_prob(flow_input).squeeze(-1)
+ if compose:
+ log_probs = log_probs - log_abs_det
# Force probability to be zero outside prior support.
in_prior_support = within_support(self.prior, theta)
@@ -199,6 +249,11 @@ def gradient(
) -> Tensor:
r"""Returns the potential function gradient for score-based methods.
+ Coordinate convention: under compose_standardization the SDE sampler calls
+ gradient() on a z-space state, while MAP (which would need theta-space
+ gradients) is guarded; so gradient() always operates in the estimator's
+ current coordinate.
+
Args:
theta: The parameters at which to evaluate the potential gradient.
time: The diffusion time. If None, then `t_min` of the
diff --git a/sbi/neural_nets/estimators/base.py b/sbi/neural_nets/estimators/base.py
index 68eb7f996..4d5fefed2 100644
--- a/sbi/neural_nets/estimators/base.py
+++ b/sbi/neural_nets/estimators/base.py
@@ -386,6 +386,65 @@ def __init__(
embedding_net if embedding_net is not None else nn.Identity()
)
+ # ---- Composed standardization (opt-in; default OFF) ----
+ # When enabled, the estimator stays PURE z-space and an invertible per-dim
+ # affine transform theta = shift + scale * z is applied ONLY at the
+ # boundaries (loss input, sample output, log_prob input). Buffers default
+ # to the identity transform so the flag-OFF path is byte-identical.
+ self.register_buffer(
+ "_theta_shift", torch.zeros(1, *self.input_shape, dtype=torch.float32)
+ )
+ self.register_buffer(
+ "_theta_scale", torch.ones(1, *self.input_shape, dtype=torch.float32)
+ )
+ # Persisted as a buffer (not a plain attribute) so that compose mode survives
+ # save/load: a checkpoint trained with composition reloads as compose-enabled.
+ self.register_buffer(
+ "_compose_standardization", torch.tensor(False), persistent=True
+ )
+ # Plain Python bool mirror — NOT a buffer, NOT persisted. Allows hot-path
+ # guards to short-circuit before touching the tensor buffer (a tensor
+ # truth-test forces a host/device sync on CUDA). Kept in sync by
+ # _wire_compose (build path) and _load_from_state_dict (checkpoint path).
+ self._compose_enabled: bool = False
+
+ def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+ r"""Gracefully load checkpoints from before composed standardization existed.
+
+ The three compose buffers (``_theta_shift``/``_theta_scale``/
+ ``_compose_standardization``) are handled atomically:
+
+ * NONE present (legacy pre-compose checkpoint): inject the current
+ identity / disabled defaults so the estimator loads as composition-OFF
+ (its original behavior).
+ * ALL present (compose-aware checkpoint): load normally.
+ * PARTIAL (some present, some missing): raise rather than silently inject
+ identity defaults — a checkpoint with ``_compose_standardization=True``
+ but a missing ``_theta_shift``/``_theta_scale`` would otherwise load a
+ silently-wrong identity affine.
+ """
+ compose_names = ("_theta_shift", "_theta_scale", "_compose_standardization")
+ present = [n for n in compose_names if prefix + n in state_dict]
+ if present and len(present) != len(compose_names):
+ missing = [n for n in compose_names if prefix + n not in state_dict]
+ raise RuntimeError(
+ "Partial composed-standardization checkpoint: present "
+ f"{[prefix + n for n in present]} but missing "
+ f"{[prefix + n for n in missing]}. Refusing to inject identity "
+ "defaults for the missing buffers (that would silently produce a "
+ "wrong affine). Load a checkpoint with either all three compose "
+ "buffers or none of them."
+ )
+ if not present:
+ # Legacy checkpoint: inject identity / disabled defaults (compose-OFF).
+ for name in compose_names:
+ if hasattr(self, name):
+ state_dict[prefix + name] = getattr(self, name).clone()
+ super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+ # Sync the plain-Python mirror AFTER the buffer is loaded so the hot-path
+ # short-circuit in _check_compose_internal_stats_unit stays correct.
+ self._compose_enabled = bool(self._compose_standardization)
+
@property
def embedding_net(self) -> nn.Module:
r"""Return the embedding network if it exists."""
@@ -420,6 +479,59 @@ def std_base(self) -> Tensor:
(the initial noise at time t=T)."""
return self._std_base
+ # -------------------------- COMPOSED STANDARDIZATION HELPERS ----------------
+
+ def to_z(self, theta: Tensor) -> Tensor:
+ r"""Map original-space parameters to standardized z-space.
+
+ ``z = (theta - shift) / scale``. Identity when composed standardization
+ is disabled (shift=0, scale=1).
+ """
+ return (theta - self._theta_shift) / self._theta_scale
+
+ def from_z(self, z: Tensor) -> Tensor:
+ r"""Map standardized z-space parameters back to original space.
+
+ ``theta = shift + scale * z``. Identity when composed standardization
+ is disabled (shift=0, scale=1).
+ """
+ return self._theta_shift + self._theta_scale * z
+
+ def log_abs_det(self) -> Tensor:
+ r"""Log absolute determinant of the affine z->theta Jacobian.
+
+ ``sum(log scale)`` over the input dimensions. Used to correct the
+ flow log-prob (which is computed in z-space) back to theta-space:
+ ``log p_theta(theta) = log p_z(z) - sum(log scale)``.
+ """
+ return torch.log(self._theta_scale).sum()
+
+ def _check_compose_internal_stats_unit(self) -> None:
+ r"""Enforce the composed-standardization invariant: unit internal stats.
+
+ When composition is enabled the estimator must stay PURE z-space, which
+ requires the internal input-norm stats ``mean_0``/``std_0`` to be exactly
+ ``0``/``1`` (the network sees unit-z input; the boundary affine carries the
+ original-space scale). The build path forces this, but a manually assembled
+ estimator or a tampered checkpoint could re-enable composition with
+ non-unit stats, silently producing an inconsistent model. Reject it at
+ use-time. Short-circuits cheaply when composition is off (the common path).
+ """
+ # Short-circuit on the plain Python bool mirror (no tensor truth-test,
+ # no host/device sync). The compose-OFF path (the common case, called on
+ # every ODE/SDE step) pays zero tensor access.
+ if not self._compose_enabled:
+ return
+ if not ((self.mean_0 == 0).all() and (self.std_0 == 1).all()):
+ raise ValueError(
+ "compose_standardization is enabled but the internal input-norm "
+ "stats mean_0/std_0 are not unit (0/1). Under composition the "
+ "estimator must stay pure z-space (the boundary affine carries the "
+ "original-space scale). This estimator is internally inconsistent "
+ "— rebuild it via build_vector_field_estimator(..., "
+ "compose_standardization=True), which forces unit stats."
+ )
+
# -------------------------- ODE METHODS --------------------------
@abstractmethod
diff --git a/sbi/neural_nets/estimators/flowmatching_estimator.py b/sbi/neural_nets/estimators/flowmatching_estimator.py
index ad76d1e80..366527044 100644
--- a/sbi/neural_nets/estimators/flowmatching_estimator.py
+++ b/sbi/neural_nets/estimators/flowmatching_estimator.py
@@ -109,6 +109,33 @@ def __init__(
self.register_buffer("mean_0", mean_0_tensor)
self.register_buffer("std_0", std_0_tensor)
+ def _check_compose_baseline_compatible(self) -> None:
+ """Guard the unsupported gaussian_baseline + composed-standardization pair.
+
+ ``gaussian_baseline`` derives the velocity from the (original-space) data
+ statistics ``mean_0``/``std_0``; composed standardization instead maps
+ ``theta -> z`` so the network sees standardized inputs. The two are
+ mutually exclusive (see also the build-time guard in ``_wire_compose``).
+
+ The public build path refuses the pair up front; this runtime check is the
+ backstop for estimators assembled manually, or loaded from a checkpoint
+ that predates the build-time guard, where ``_compose_standardization`` is
+ re-activated alongside ``gaussian_baseline``.
+ """
+ # ``gaussian_baseline`` is a plain Python bool, so checking it first lets
+ # the common path (and every ODE step) short-circuit before touching the
+ # ``_compose_standardization`` tensor buffer (whose truth test would force a
+ # device sync on CUDA).
+ if self.gaussian_baseline and bool(self._compose_standardization):
+ raise ValueError(
+ "gaussian_baseline and composed standardization cannot be used "
+ "together: the analytical baseline velocity is derived from "
+ "mean_0/std_0 in the original data space, while composed "
+ "standardization feeds the network standardized z (theta -> z). "
+ "Enable only one (set gaussian_baseline=False to use composed "
+ "standardization)."
+ )
+
def _get_time_dependent_stats(self, time: Tensor) -> Tuple[Tensor, Tensor]:
"""Compute time-dependent mean and std for z-scoring.
@@ -209,6 +236,9 @@ def forward(self, input: Tensor, condition: Tensor, time: Tensor) -> Tensor:
Returns:
The estimated vector field in original space.
"""
+ self._check_compose_baseline_compatible()
+ self._check_compose_internal_stats_unit()
+
# Continue with standard processing (broadcast shapes etc.)
batch_shape_input = input.shape[: -len(self.input_shape)]
batch_shape_cond = condition.shape[: -len(self.condition_shape)]
@@ -283,6 +313,15 @@ def loss(
Returns:
Loss value.
"""
+ self._check_compose_baseline_compatible()
+ self._check_compose_internal_stats_unit()
+
+ # Composed standardization (opt-in): the estimator is trained PURELY in
+ # z-space. Standardize theta -> z at the very top; everything below is
+ # unchanged. When the flag is off, this is a no-op (shift=0, scale=1).
+ if self._compose_standardization:
+ input = self.to_z(input)
+
# Randomly sample time steps
if times is None:
times = torch.rand(input.shape[:-1], device=input.device, dtype=input.dtype)
diff --git a/sbi/neural_nets/estimators/score_estimator.py b/sbi/neural_nets/estimators/score_estimator.py
index 87bead1ee..fe83557a0 100644
--- a/sbi/neural_nets/estimators/score_estimator.py
+++ b/sbi/neural_nets/estimators/score_estimator.py
@@ -152,6 +152,7 @@ def forward(self, input: Tensor, condition: Tensor, time: Tensor) -> Tensor:
Returns:
Score (gradient of the density) at a given time, matches input shape.
"""
+ self._check_compose_internal_stats_unit()
# Continue with standard processing (broadcast shapes etc.)
batch_shape_input = input.shape[: -len(self.input_shape)]
@@ -248,6 +249,13 @@ def loss(
MSE between target score and network output, scaled by the weight function.
"""
+ # Composed standardization (opt-in): the estimator is trained PURELY in
+ # z-space. Standardize theta -> z at the very top; everything below is
+ # unchanged. When the flag is off, this is a no-op (shift=0, scale=1).
+ self._check_compose_internal_stats_unit()
+ if self._compose_standardization:
+ input = self.to_z(input)
+
# Sample times from the Markov chain, use batch dimension
if times is None:
times = self.train_schedule(input.shape[0])
diff --git a/sbi/neural_nets/factory.py b/sbi/neural_nets/factory.py
index 14b12bb33..4ef1d6fca 100644
--- a/sbi/neural_nets/factory.py
+++ b/sbi/neural_nets/factory.py
@@ -358,6 +358,7 @@ def posterior_score_nn(
embedding_net: nn.Module = nn.Identity(),
time_emb_type: Literal["sinusoidal", "fourier"] = "sinusoidal",
t_embedding_dim: int = 32,
+ compose_standardization: bool = False,
**kwargs: Any,
) -> Callable:
"""Build util function that builds a ScoreEstimator object for score-based
@@ -391,6 +392,8 @@ def posterior_score_nn(
nn.Identity().
time_emb_type: Type of time embedding. Defaults to 'sinusoidal'.
t_embedding_dim: Embedding dimension of diffusion time. Defaults to 32.
+ compose_standardization: Opt-in per-dim affine standardization theta<->z
+ for scale-equivariant calibration. Defaults to False.
**kwargs: Additional estimator / network arguments. Valid keys are
defined by ``ScoreEstimatorConfig``; unknown keys raise
``TypeError``.
@@ -408,6 +411,7 @@ def posterior_score_nn(
time_embedding_dim=t_embedding_dim,
time_emb_type=time_emb_type,
net=model,
+ compose_standardization=compose_standardization,
)
# Validate against known fields — warns on unknown kwargs (typos)
@@ -443,6 +447,7 @@ def posterior_flow_nn(
time_emb_type: Literal["sinusoidal", "fourier"] = "sinusoidal",
t_embedding_dim: int = 32,
gaussian_baseline: bool = False,
+ compose_standardization: bool = False,
**kwargs: Any,
) -> Callable:
"""Build util function that builds a FlowMatchingEstimator object for flow-based
@@ -471,6 +476,8 @@ def posterior_flow_nn(
gaussian_baseline: If True, use analytical Gaussian baseline velocity
derived from Bayes' rule. The network then only learns the residual.
Defaults to False.
+ compose_standardization: Opt-in per-dim affine standardization theta<->z
+ for scale-equivariant calibration. Defaults to False.
**kwargs: Additional estimator / network arguments. Valid keys are
defined by ``FlowEstimatorConfig``; unknown keys raise
``TypeError``.
@@ -489,6 +496,7 @@ def posterior_flow_nn(
time_emb_type=time_emb_type,
net=model,
gaussian_baseline=gaussian_baseline,
+ compose_standardization=compose_standardization,
)
# Validate against known fields — warns on unknown kwargs (typos)
diff --git a/sbi/neural_nets/net_builders/vector_field_nets.py b/sbi/neural_nets/net_builders/vector_field_nets.py
index ab829974a..8f8a8f957 100644
--- a/sbi/neural_nets/net_builders/vector_field_nets.py
+++ b/sbi/neural_nets/net_builders/vector_field_nets.py
@@ -41,6 +41,11 @@ class _VectorFieldBaseConfig(_EstimatorConfigBase):
sinusoidal_max_freq: Optional[float] = None
fourier_scale: Optional[float] = None
+ # Composed standardization (opt-in scale-equivariance for FMPE / NPSE).
+ # When True, the estimator is trained/sampled in standardized z-space with an
+ # invertible per-dim affine transform applied only at the boundaries.
+ compose_standardization: Optional[bool] = None
+
# MLP-specific
layer_norm: Optional[bool] = None
skip_connections: Optional[bool] = None
@@ -108,6 +113,54 @@ class FlowEstimatorConfig(_VectorFieldBaseConfig):
gaussian_baseline: Optional[bool] = None
+def _compute_theta_standardization(
+ batch_x: Tensor,
+ z_score_x: Optional[str],
+ compose_standardization: bool,
+):
+ """Single source of truth for the theta-standardization stats.
+
+ This is the ONLY place that computes the standardization statistics from the
+ training theta, replacing the two previously-parallel ``z_standardization``
+ call sites (maintainer feedback: "don't maintain two parallel mechanisms").
+
+ Returns ``(mean_0, std_0, compose_shift, compose_scale)``:
+
+ * compose OFF: ``mean_0``/``std_0`` are the per-dim z-score stats (or scalar
+ ``0``/``1`` when z-scoring is disabled); ``compose_shift``/``compose_scale``
+ are ``None`` (no boundary affine).
+ * compose ON: the estimator stays PURE z-space, so the boundary affine
+ ``theta = shift + scale * z`` carries the per-dim z-score stats
+ (``compose_shift``/``compose_scale``, scale clamped to ``1e-20``), and the
+ internal input-norm stats are forced to unit (``mean_0=0``/``std_0=1``) so
+ the network sees unit-z input. The unit invariant is asserted here.
+ """
+ z_score_x_bool, structured_x = z_score_parser(z_score_x)
+ if z_score_x_bool:
+ mean_0, std_0 = z_standardization(batch_x, structured_x)
+ else:
+ mean_0, std_0 = 0, 1
+
+ compose_shift = None
+ compose_scale = None
+ if compose_standardization:
+ # Per-dim standardization (independent), matching the validated reference.
+ compose_shift, compose_scale = z_standardization(batch_x, structured_dims=False)
+ compose_scale = compose_scale.clamp_min(1e-20)
+ # Force the internal input-norm stats to unit (the estimator now sees z).
+ # Use PER-DIM tensors (not scalars) so the VE base buffers (mean_t/std_t,
+ # computed via broadcast_to) stay contiguous and load_state_dict-safe.
+ mean_0 = torch.zeros_like(compose_shift)
+ std_0 = torch.ones_like(compose_scale)
+ # Invariant: under composition the internal stats MUST be unit-z.
+ assert (mean_0 == 0).all() and (std_0 == 1).all(), (
+ "compose_standardization invariant violated: internal mean_0/std_0 "
+ "must be 0/1 (unit z) when composition is enabled."
+ )
+
+ return mean_0, std_0, compose_shift, compose_scale
+
+
def build_vector_field_estimator(
batch_x: Tensor,
batch_y: Tensor,
@@ -126,6 +179,7 @@ def build_vector_field_estimator(
VectorFieldNet,
] = "mlp",
gaussian_baseline: bool = False,
+ compose_standardization: bool = False,
**kwargs,
) -> Union[FlowMatchingEstimator, ConditionalScoreEstimator]:
"""Builds a vector field estimator (flow matching or score matching) with the given
@@ -151,6 +205,12 @@ 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.
+ compose_standardization: If True (opt-in), compose an invertible per-dim
+ affine standardization with the flow so the estimator becomes
+ scale-equivariant. The estimator is trained and sampled in
+ standardized z-space (per-dim mean/std computed from training theta),
+ and theta = shift + scale * z is applied only at the boundaries.
+ Defaults to False (behavior identical to current sbi).
**kwargs: Additional arguments forwarded to the estimator and network
constructors. Valid keys are defined by ``ScoreEstimatorConfig``
and ``FlowEstimatorConfig``; validation happens in the upstream
@@ -218,12 +278,17 @@ def build_vector_field_estimator(
else:
raise ValueError(f"Unknown architecture: {net}")
- # Z-score setup
- z_score_x_bool, structured_x = z_score_parser(z_score_x)
- if z_score_x_bool:
- mean_0, std_0 = z_standardization(batch_x, structured_x)
- else:
- mean_0, std_0 = 0, 1
+ # Z-score + composed-standardization setup (single source of truth). The
+ # helper computes the internal input-norm stats (mean_0/std_0) and, when
+ # composition is enabled, the per-dim boundary affine theta = shift + scale*z
+ # (forcing mean_0=0/std_0=1 so the network sees unit-z input). The base
+ # distribution is left in z-space: FMPE uses N(0, 1), and VE uses
+ # approx_marginal_std(t_max) = sqrt(std_0**2 + sigma_max**2) = sqrt(1 +
+ # sigma_max**2) since compose forces std_0=1. The base is NOT scaled by
+ # `scale` — sampling happens in z, and from_z then maps z -> theta.
+ mean_0, std_0, compose_shift, compose_scale = _compute_theta_standardization(
+ batch_x, z_score_x, compose_standardization
+ )
z_score_y_bool, structured_y = z_score_parser(z_score_y)
embedding_net_y = (
@@ -232,15 +297,47 @@ def build_vector_field_estimator(
else embedding_net
)
+ def _wire_compose(estimator):
+ """Wire the per-dim affine standardization onto the estimator (opt-in)."""
+ if compose_shift is not None and compose_scale is not None:
+ # gaussian_baseline derives the network velocity from the
+ # (original-space) data statistics mean_0/std_0; composed
+ # standardization instead maps theta -> z so the network sees
+ # unit-scale inputs. The two are mutually exclusive. Refuse the
+ # (untested) pair explicitly at this public-API chokepoint rather than
+ # silently producing an inconsistent estimator. Score estimators have
+ # no gaussian_baseline attribute (getattr default False).
+ if getattr(estimator, "gaussian_baseline", False):
+ raise ValueError(
+ "gaussian_baseline and compose_standardization cannot be "
+ "enabled together: gaussian_baseline assumes non-standardized "
+ "data (its analytical velocity is derived from mean_0/std_0), "
+ "while compose_standardization standardizes theta -> z before "
+ "the network. Enable only one (set gaussian_baseline=False to "
+ "use composed standardization)."
+ )
+ shift = compose_shift.reshape(1, *estimator.input_shape).float()
+ scale = compose_scale.reshape(1, *estimator.input_shape).float()
+ estimator._theta_shift.copy_(shift)
+ estimator._theta_scale.copy_(scale)
+ estimator._compose_standardization.fill_(True)
+ # Sync the plain-Python mirror so the hot-path guard in
+ # _check_compose_internal_stats_unit short-circuits without a
+ # tensor truth-test on the compose-OFF (default) path.
+ estimator._compose_enabled = True
+ return estimator
+
if estimator_type == "flow":
- return FlowMatchingEstimator(
- net=vectorfield_net,
- input_shape=batch_x[0].shape,
- condition_shape=batch_y[0].shape,
- embedding_net=embedding_net_y,
- mean_0=mean_0,
- std_0=std_0,
- gaussian_baseline=gaussian_baseline,
+ return _wire_compose(
+ FlowMatchingEstimator(
+ net=vectorfield_net,
+ input_shape=batch_x[0].shape,
+ condition_shape=batch_y[0].shape,
+ embedding_net=embedding_net_y,
+ mean_0=mean_0,
+ std_0=std_0,
+ gaussian_baseline=gaussian_baseline,
+ )
)
elif estimator_type == "score":
# Choose the appropriate score estimator based on SDE type
@@ -272,14 +369,16 @@ def build_vector_field_estimator(
vp_keys = ["beta_min", "beta_max"]
estimator_kwargs = {k: kwargs[k] for k in vp_keys if k in kwargs}
- return estimator_cls(
- net=vectorfield_net,
- input_shape=batch_x[0].shape,
- condition_shape=batch_y[0].shape,
- embedding_net=embedding_net_y,
- mean_0=mean_0,
- std_0=std_0,
- **estimator_kwargs,
+ return _wire_compose(
+ estimator_cls(
+ net=vectorfield_net,
+ input_shape=batch_x[0].shape,
+ condition_shape=batch_y[0].shape,
+ embedding_net=embedding_net_y,
+ mean_0=mean_0,
+ std_0=std_0,
+ **estimator_kwargs,
+ )
)
else:
raise ValueError(f"Unknown estimator type: {estimator_type}")
diff --git a/tests/test_compose_baseline_guard.py b/tests/test_compose_baseline_guard.py
new file mode 100644
index 000000000..2753c714a
--- /dev/null
+++ b/tests/test_compose_baseline_guard.py
@@ -0,0 +1,227 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see
+
+"""Guard tests for the ``gaussian_baseline`` + ``compose_standardization`` pair.
+
+``gaussian_baseline`` (PR #1752) derives the network velocity from the
+original-space data statistics ``mean_0``/``std_0``. Composed standardization
+(#1680) instead maps ``theta -> z`` so the network sees unit-scale inputs. The
+two are mutually exclusive; the maintainer flagged the combination as untested.
+
+Two guards make the behavior defined instead of silently inconsistent:
+ * BUILD-TIME (``_wire_compose``): the public API (``build_vector_field_estimator``
+ / ``posterior_flow_nn``) refuses the pair up front.
+ * RUNTIME (``FlowMatchingEstimator._check_compose_baseline_compatible``): backstop
+ that rejects the pair on loss()/forward() for estimators assembled manually, or
+ loaded from a checkpoint predating the build-time guard, where composition is
+ re-activated alongside ``gaussian_baseline`` (regardless of the stored stats).
+
+These are fast unit tests (no training); the equivariance behavior itself lives
+in ``tests/test_scale_equivariance.py``.
+"""
+
+import pytest
+import torch
+
+from sbi.neural_nets.estimators.flowmatching_estimator import FlowMatchingEstimator
+from sbi.neural_nets.net_builders.vector_field_nets import (
+ build_vector_field_estimator,
+)
+
+NUM_DIM = 2
+BATCH = 8
+
+
+class _ZeroNet(torch.nn.Module):
+ """Minimal VectorFieldNet stub: returns zeros, trainable param keeps autograd."""
+
+ def __init__(self):
+ super().__init__()
+ self.dummy = torch.nn.Parameter(torch.zeros(1))
+
+ def forward(self, input, condition, time):
+ return torch.zeros_like(input) * self.dummy
+
+
+def _batches():
+ torch.manual_seed(0)
+ # Non-unit theta stats on purpose: exercises the build path that zeroes them.
+ batch_x = 100.0 + 5.0 * torch.randn(16, NUM_DIM)
+ batch_y = torch.randn(16, NUM_DIM)
+ return batch_x, batch_y
+
+
+def _manual_estimator(mean_0, std_0, gaussian_baseline, compose):
+ """Assemble a FlowMatchingEstimator directly and (optionally) flip the private
+ composition buffers on — mimicking misuse that bypasses the build guard."""
+ est = FlowMatchingEstimator(
+ net=_ZeroNet(),
+ input_shape=torch.Size([NUM_DIM]),
+ condition_shape=torch.Size([NUM_DIM]),
+ mean_0=mean_0,
+ std_0=std_0,
+ gaussian_baseline=gaussian_baseline,
+ )
+ if compose:
+ est._theta_shift.copy_(torch.zeros(1, NUM_DIM))
+ est._theta_scale.copy_(torch.ones(1, NUM_DIM))
+ est._compose_standardization.fill_(True)
+ est._compose_enabled = True # sync plain-Python mirror (as _wire_compose does)
+ return est
+
+
+# --------------------------------------------------------------------------
+# Guard II — build-time, public API
+# --------------------------------------------------------------------------
+
+
+def test_build_rejects_compose_plus_baseline():
+ """Public build path refuses compose_standardization + gaussian_baseline."""
+ batch_x, batch_y = _batches()
+ with pytest.raises(
+ ValueError, match="gaussian_baseline and compose_standardization"
+ ):
+ build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ gaussian_baseline=True,
+ compose_standardization=True,
+ )
+
+
+def test_build_compose_without_baseline_ok():
+ """Normal compose path (gaussian_baseline=False) builds and enables compose;
+ build forces mean_0=0/std_0=1 (unit z-stats)."""
+ batch_x, batch_y = _batches()
+ est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ gaussian_baseline=False,
+ compose_standardization=True,
+ )
+ assert bool(est._compose_standardization)
+ assert not est.gaussian_baseline
+ assert torch.allclose(est.mean_0, torch.zeros_like(est.mean_0))
+ assert torch.allclose(est.std_0, torch.ones_like(est.std_0))
+
+
+def test_build_baseline_without_compose_ok():
+ """Normal baseline path (compose off) builds; compose flag stays False."""
+ batch_x, batch_y = _batches()
+ est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ gaussian_baseline=True,
+ compose_standardization=False,
+ )
+ assert est.gaussian_baseline
+ assert not bool(est._compose_standardization)
+
+
+# --------------------------------------------------------------------------
+# Guard I — runtime, manual assembly that bypasses the build guard
+# --------------------------------------------------------------------------
+
+
+def _inputs():
+ torch.manual_seed(1)
+ inp = torch.randn(BATCH, NUM_DIM)
+ cond = torch.randn(BATCH, NUM_DIM)
+ t = torch.rand(BATCH)
+ return inp, cond, t
+
+
+def test_runtime_rejects_manual_compose_baseline_nonunit_stats():
+ """Manual estimator: compose + baseline + NON-unit mean_0/std_0 -> raise on
+ both loss() and forward()."""
+ est = _manual_estimator(
+ mean_0=torch.tensor([100.0, 100.0]),
+ std_0=torch.tensor([5.0, 5.0]),
+ gaussian_baseline=True,
+ compose=True,
+ )
+ inp, cond, t = _inputs()
+ with pytest.raises(ValueError, match="cannot be used together"):
+ est.loss(inp, cond, t)
+ with pytest.raises(ValueError, match="cannot be used together"):
+ est.forward(inp, cond, t)
+
+
+def test_runtime_rejects_compose_baseline_unit_stats():
+ """Policy: the pair is rejected ENTIRELY, not only when stats are non-unit.
+
+ The build path forces mean_0=0/std_0=1 under composition, so a checkpoint
+ produced before the build-time guard existed (or a manual assembly) carries
+ unit stats. The runtime guard must still reject it on loss() and forward(),
+ consistent with the build-time guard that refuses the pair up front."""
+ est = _manual_estimator(
+ mean_0=torch.zeros(NUM_DIM),
+ std_0=torch.ones(NUM_DIM),
+ gaussian_baseline=True,
+ compose=True,
+ )
+ inp, cond, t = _inputs()
+ with pytest.raises(ValueError, match="cannot be used together"):
+ est.loss(inp, cond, t)
+ with pytest.raises(ValueError, match="cannot be used together"):
+ est.forward(inp, cond, t)
+
+
+def test_runtime_rejects_compose_with_nonunit_internal_stats():
+ """T1.2 invariant: compose ON requires unit internal stats (mean_0=0/std_0=1).
+
+ The build path forces unit z-stats under composition; a manual assembly (or a
+ tampered checkpoint) with compose ON but NON-unit mean_0/std_0 is internally
+ inconsistent (the network would not see unit-z input) and must RAISE on both
+ loss() and forward()."""
+ est = _manual_estimator(
+ mean_0=torch.tensor([100.0, 100.0]),
+ std_0=torch.tensor([5.0, 5.0]),
+ gaussian_baseline=False,
+ compose=True,
+ )
+ inp, cond, t = _inputs()
+ with pytest.raises(ValueError, match="unit"):
+ est.loss(inp, cond, t)
+ with pytest.raises(ValueError, match="unit"):
+ est.forward(inp, cond, t)
+
+
+def test_runtime_allows_compose_with_unit_internal_stats():
+ """No over-fire: compose ON with unit internal stats (mean_0=0/std_0=1) runs
+ fine on both loss() and forward()."""
+ est = _manual_estimator(
+ mean_0=torch.zeros(NUM_DIM),
+ std_0=torch.ones(NUM_DIM),
+ gaussian_baseline=False,
+ compose=True,
+ )
+ inp, cond, t = _inputs()
+ loss = est.loss(inp, cond, t)
+ v = est.forward(inp, cond, t)
+ assert torch.isfinite(loss).all()
+ assert v.shape == (BATCH, NUM_DIM)
+ assert torch.isfinite(v).all()
+
+
+def test_runtime_allows_baseline_without_compose():
+ """Guard does not over-fire: baseline ON, compose OFF, non-unit stats -> runs
+ (plain gaussian_baseline FMPE, no composition)."""
+ est = _manual_estimator(
+ mean_0=torch.tensor([100.0, 100.0]),
+ std_0=torch.tensor([5.0, 5.0]),
+ gaussian_baseline=True,
+ compose=False,
+ )
+ inp, cond, t = _inputs()
+ loss = est.loss(inp, cond, t)
+ v = est.forward(inp, cond, t)
+ assert torch.isfinite(loss).all()
+ assert v.shape == (BATCH, NUM_DIM)
+ assert torch.isfinite(v).all()
diff --git a/tests/test_compose_scope_guards.py b/tests/test_compose_scope_guards.py
new file mode 100644
index 000000000..efd6a15cd
--- /dev/null
+++ b/tests/test_compose_scope_guards.py
@@ -0,0 +1,227 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see
+
+"""Scope-guard and boundary tests for opt-in ``compose_standardization``.
+
+These cover the patch lines that the build/checkpoint tests in
+``test_compose_standardization_internal.py`` and the baseline guards in
+``test_compose_baseline_guard.py`` do not reach:
+
+ * the ``NotImplementedError`` scope guards for iid / guided sampling, iid
+ ``log_prob``, MAP, and the potential ``set_x`` chokepoint (single-obs is the
+ only supported path under composition);
+ * the ``from_z`` sample-output boundary on both the ODE (FMPE) and SDE (NPSE)
+ paths, i.e. samples are returned in original theta-space, not z-space;
+ * the score-estimator ``loss`` composition hook (``theta -> z`` at the top);
+ * the compose-OFF ``else`` branch of the single-obs ``log_prob`` potential.
+
+All tests are training-free and fast: the guards raise up front, and the two
+sampling tests run an untrained estimator for a handful of steps (the numbers
+are meaningless, only the coordinate space and finiteness are asserted).
+"""
+
+import pytest
+import torch
+from torch.distributions import Independent, Normal
+
+from sbi.inference.posteriors.vector_field_posterior import VectorFieldPosterior
+from sbi.inference.potentials.vector_field_potential import VectorFieldBasedPotential
+from sbi.neural_nets.net_builders.vector_field_nets import (
+ build_vector_field_estimator,
+)
+
+NUM_DIM = 2
+
+
+def _batches():
+ torch.manual_seed(0)
+ # Non-unit theta stats: the compose shift is ~100, so from_z-corrected samples
+ # sit far from the z-space origin (used by the boundary assertions below).
+ batch_x = 100.0 + 5.0 * torch.randn(64, NUM_DIM)
+ batch_y = torch.randn(64, NUM_DIM)
+ return batch_x, batch_y
+
+
+def _wide_prior():
+ return Independent(Normal(torch.zeros(NUM_DIM), 200.0 * torch.ones(NUM_DIM)), 1)
+
+
+def _compose_estimator(estimator_type="flow", **kwargs):
+ batch_x, batch_y = _batches()
+ return build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type=estimator_type,
+ z_score_x="independent",
+ compose_standardization=True,
+ **kwargs,
+ )
+
+
+def _compose_posterior(estimator_type="flow", sample_with="ode", **kwargs):
+ est = _compose_estimator(estimator_type=estimator_type, **kwargs)
+ return VectorFieldPosterior(
+ vector_field_estimator=est, prior=_wide_prior(), sample_with=sample_with
+ )
+
+
+# --------------------------------------------------------------------------
+# Scope guards — raise NotImplementedError rather than return wrong results
+# --------------------------------------------------------------------------
+
+
+def test_sample_iid_rejects_compose():
+ """sample() with batched (iid) x under compose -> NotImplementedError."""
+ posterior = _compose_posterior()
+ x_iid = torch.randn(3, NUM_DIM) # batch > 1 => iid
+ with pytest.raises(NotImplementedError, match="iid"):
+ posterior.sample((2,), x=x_iid, show_progress_bars=False)
+
+
+def test_sample_guidance_rejects_compose():
+ """sample() with a guidance method under compose -> NotImplementedError."""
+ posterior = _compose_posterior()
+ x_o = torch.zeros(1, NUM_DIM)
+ with pytest.raises(NotImplementedError, match="guided"):
+ posterior.sample(
+ (2,), x=x_o, guidance_method="classifier_free", show_progress_bars=False
+ )
+
+
+def test_log_prob_iid_rejects_compose():
+ """log_prob() with batched (iid) x under compose -> NotImplementedError."""
+ posterior = _compose_posterior()
+ x_iid = torch.randn(3, NUM_DIM)
+ with pytest.raises(NotImplementedError, match="iid"):
+ posterior.log_prob(torch.zeros(1, NUM_DIM), x=x_iid)
+
+
+def test_map_rejects_compose():
+ """map() under compose -> NotImplementedError (guard is the first statement,
+ so it fires before the default-x check)."""
+ posterior = _compose_posterior()
+ with pytest.raises(NotImplementedError, match="MAP"):
+ posterior.map(show_progress_bars=False)
+
+
+def test_potential_set_x_iid_rejects_compose():
+ """Direct potential use: set_x(x_is_iid=True) under compose -> NotImplementedError
+ (the chokepoint guard, reachable without going through the posterior)."""
+ est = _compose_estimator()
+ potential = VectorFieldBasedPotential(
+ est, prior=_wide_prior(), x_o=None, device="cpu"
+ )
+ with pytest.raises(NotImplementedError, match="iid"):
+ potential.set_x(torch.randn(2, NUM_DIM), x_is_iid=True)
+
+
+def test_potential_set_x_guidance_rejects_compose():
+ """Direct potential use: set_x(guidance_method=...) under compose ->
+ NotImplementedError."""
+ est = _compose_estimator()
+ potential = VectorFieldBasedPotential(
+ est, prior=_wide_prior(), x_o=None, device="cpu"
+ )
+ with pytest.raises(NotImplementedError, match="guided"):
+ potential.set_x(
+ torch.zeros(1, NUM_DIM), x_is_iid=False, guidance_method="classifier_free"
+ )
+
+
+# --------------------------------------------------------------------------
+# from_z sample-output boundary — samples returned in ORIGINAL theta-space
+# --------------------------------------------------------------------------
+
+
+def test_sample_ode_returns_theta_space_under_compose():
+ """FMPE/ODE single-obs sample under compose applies from_z: output lives in
+ original theta-space (shift ~100), not z-space. Untrained -> only coordinate
+ space + finiteness asserted."""
+ posterior = _compose_posterior(estimator_type="flow", sample_with="ode")
+ samples = posterior.sample(
+ (5,),
+ x=torch.zeros(1, NUM_DIM),
+ reject_outside_prior=False,
+ show_progress_bars=False,
+ )
+ assert samples.shape == (5, NUM_DIM)
+ assert torch.isfinite(samples).all()
+ # from_z applied => samples carry the ~100 shift; z-space samples would be O(1).
+ assert samples.abs().mean() > 10.0
+
+
+def test_sample_sde_returns_theta_space_under_compose():
+ """NPSE/SDE single-obs sample under compose applies from_z on the diffusion
+ path: output lives in original theta-space."""
+ posterior = _compose_posterior(
+ estimator_type="score", sample_with="sde", sde_type="ve"
+ )
+ samples = posterior.sample(
+ (5,),
+ x=torch.zeros(1, NUM_DIM),
+ reject_outside_prior=False,
+ show_progress_bars=False,
+ )
+ assert samples.shape == (5, NUM_DIM)
+ assert torch.isfinite(samples).all()
+ assert samples.abs().mean() > 10.0
+
+
+# --------------------------------------------------------------------------
+# score-estimator loss composition hook (theta -> z at the top of loss)
+# --------------------------------------------------------------------------
+
+
+def test_score_estimator_loss_compose_hook():
+ """Building a compose NPSE and calling loss() exercises the score-estimator
+ composition hook (unit-stats check + theta->z standardization of the input)."""
+ est = _compose_estimator(estimator_type="score", sde_type="ve")
+ theta = 100.0 + 5.0 * torch.randn(16, NUM_DIM)
+ x = torch.randn(16, NUM_DIM)
+ loss = est.loss(theta, x)
+ assert torch.isfinite(loss).all()
+
+
+# --------------------------------------------------------------------------
+# compose-OFF else branch of the single-obs log_prob potential
+# --------------------------------------------------------------------------
+
+
+def test_single_obs_log_prob_compose_off_else_branch():
+ """compose OFF: the potential takes the identity (else) branch of the
+ single-obs log_prob path (no to_z, no Jacobian). Runs and returns finite."""
+ batch_x, batch_y = _batches()
+ est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=False,
+ )
+ prior = _wide_prior()
+ potential = VectorFieldBasedPotential(est, prior=prior, x_o=None, device="cpu")
+ potential.set_x(torch.zeros(1, NUM_DIM), x_is_iid=False) # builds flow, no training
+ out = potential(torch.zeros(1, NUM_DIM))
+ assert torch.isfinite(out).all()
+
+
+# --------------------------------------------------------------------------
+# R4 backstop — directly-forced iid under compose raises in __call__
+# --------------------------------------------------------------------------
+
+
+def test_iid_log_prob_compose_backstop_raises_on_direct_force():
+ """R4 backstop: forcing ``_x_is_iid=True`` under compose (bypassing the set_x
+ guard) makes the potential ``__call__`` raise ``NotImplementedError`` instead of
+ returning a partially-incorrect log-prob (affine Jacobian applied, prior not yet
+ in z). This is exactly the branch the set_x guard cannot structurally reach."""
+ est = _compose_estimator()
+ potential = VectorFieldBasedPotential(
+ est, prior=_wide_prior(), x_o=None, device="cpu"
+ )
+ # Reach the iid+compose branch WITHOUT going through set_x (which guards it):
+ potential._x_o = torch.randn(3, NUM_DIM)
+ potential._x_is_iid = True
+ potential.flows = potential.rebuild_flows_for_batch()
+ with pytest.raises(NotImplementedError, match="iid"):
+ potential(torch.zeros(1, NUM_DIM))
diff --git a/tests/test_compose_standardization_internal.py b/tests/test_compose_standardization_internal.py
new file mode 100644
index 000000000..6d2041fb9
--- /dev/null
+++ b/tests/test_compose_standardization_internal.py
@@ -0,0 +1,451 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see
+
+"""Tier-1 hardening unit tests for the opt-in ``compose_standardization`` feature.
+
+All tests here are training-free (build / wire only) and fast:
+
+ * T1.1 single-source theta-standardization: building an FMPE / NPSE-ve with
+ ``compose_standardization=True`` zeroes the internal input-norm stats
+ (mean_0=0/std_0=1) and stores the per-dim theta affine in
+ ``_theta_shift``/``_theta_scale``; with compose OFF the affine is identity.
+ * T1.5 dtype/device roundtrip: ``from_z(to_z(theta)) == theta`` in float32 and
+ float64, and ``log_abs_det() == sum(log scale)``.
+
+The build-time / runtime gaussian_baseline guards live in
+``tests/test_compose_baseline_guard.py``.
+"""
+
+import pytest
+import torch
+
+from sbi.neural_nets.net_builders.vector_field_nets import (
+ _compute_theta_standardization,
+ build_vector_field_estimator,
+)
+from sbi.utils.sbiutils import z_standardization
+
+NUM_DIM = 3
+
+
+def _batches():
+ torch.manual_seed(0)
+ # Non-unit theta stats on purpose: exercises the build path that zeroes them.
+ batch_x = 100.0 + 5.0 * torch.randn(32, NUM_DIM)
+ batch_y = torch.randn(32, NUM_DIM)
+ return batch_x, batch_y
+
+
+# --------------------------------------------------------------------------
+# T1.1 — single-source theta-standardization helper
+# --------------------------------------------------------------------------
+
+
+def test_helper_compose_off_zscore():
+ """compose OFF + z_score on -> z-score stats, no compose affine."""
+ batch_x, _ = _batches()
+ mean_0, std_0, shift, scale = _compute_theta_standardization(
+ batch_x, z_score_x="independent", compose_standardization=False
+ )
+ exp_mean, exp_std = z_standardization(batch_x, structured_dims=False)
+ assert torch.allclose(mean_0, exp_mean)
+ assert torch.allclose(std_0, exp_std)
+ assert shift is None and scale is None
+
+
+def test_helper_compose_off_no_zscore():
+ """compose OFF + z_score off -> identity stats (0/1), no compose affine."""
+ batch_x, _ = _batches()
+ mean_0, std_0, shift, scale = _compute_theta_standardization(
+ batch_x, z_score_x="none", compose_standardization=False
+ )
+ assert mean_0 == 0 and std_0 == 1
+ assert shift is None and scale is None
+
+
+def test_helper_compose_on_invariant():
+ """compose ON -> internal stats unit, compose affine = per-dim z stats."""
+ batch_x, _ = _batches()
+ mean_0, std_0, shift, scale = _compute_theta_standardization(
+ batch_x, z_score_x="independent", compose_standardization=True
+ )
+ exp_mean, exp_std = z_standardization(batch_x, structured_dims=False)
+ assert (mean_0 == 0).all() and (std_0 == 1).all()
+ assert torch.allclose(shift, exp_mean)
+ assert torch.allclose(scale, exp_std.clamp_min(1e-20))
+
+
+@pytest.mark.parametrize(
+ "estimator_type,kwargs",
+ [
+ ("flow", {}),
+ ("score", {"sde_type": "ve"}),
+ ],
+)
+def test_build_compose_on_unit_stats_and_affine(estimator_type, kwargs):
+ """End-to-end build (no training): compose ON => unit internal stats and
+ _theta_shift/_theta_scale match per-dim z_standardization of the batch."""
+ batch_x, batch_y = _batches()
+ est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type=estimator_type,
+ z_score_x="independent",
+ compose_standardization=True,
+ **kwargs,
+ )
+ exp_mean, exp_std = z_standardization(batch_x, structured_dims=False)
+ assert bool(est._compose_standardization) is True
+ assert torch.allclose(est.mean_0, torch.zeros_like(est.mean_0))
+ assert torch.allclose(est.std_0, torch.ones_like(est.std_0))
+ assert torch.allclose(est._theta_shift.reshape(-1), exp_mean.float(), atol=1e-5)
+ assert torch.allclose(
+ est._theta_scale.reshape(-1), exp_std.clamp_min(1e-20).float(), atol=1e-5
+ )
+
+
+@pytest.mark.parametrize(
+ "estimator_type,kwargs",
+ [
+ ("flow", {}),
+ ("score", {"sde_type": "ve"}),
+ ],
+)
+def test_build_compose_off_identity(estimator_type, kwargs):
+ """compose OFF => flag False, identity theta affine (byte-identical OFF path)."""
+ batch_x, batch_y = _batches()
+ est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type=estimator_type,
+ z_score_x="independent",
+ compose_standardization=False,
+ **kwargs,
+ )
+ assert bool(est._compose_standardization) is False
+ assert torch.allclose(est._theta_shift, torch.zeros_like(est._theta_shift))
+ assert torch.allclose(est._theta_scale, torch.ones_like(est._theta_scale))
+
+
+# --------------------------------------------------------------------------
+# T1.5 — dtype/device roundtrip of the boundary affine
+# --------------------------------------------------------------------------
+
+
+def _compose_estimator():
+ batch_x, batch_y = _batches()
+ return build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=True,
+ )
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_affine_roundtrip_dtype(dtype):
+ """from_z(to_z(theta)) == theta in float32 and float64."""
+ est = _compose_estimator()
+ if dtype == torch.float64:
+ est = est.double()
+ theta = (100.0 + 5.0 * torch.randn(10, NUM_DIM)).to(dtype)
+ recon = est.from_z(est.to_z(theta))
+ assert recon.dtype == dtype
+ assert torch.allclose(recon, theta, atol=1e-5 if dtype == torch.float32 else 1e-10)
+
+
+def test_log_abs_det_equals_sum_log_scale():
+ est = _compose_estimator()
+ assert torch.allclose(est.log_abs_det(), torch.log(est._theta_scale).sum())
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+def test_affine_roundtrip_cuda():
+ est = _compose_estimator().cuda()
+ theta = (100.0 + 5.0 * torch.randn(10, NUM_DIM)).cuda()
+ recon = est.from_z(est.to_z(theta))
+ assert recon.is_cuda
+ assert torch.allclose(recon, theta, atol=1e-5)
+
+
+# --------------------------------------------------------------------------
+# T1.3 — partial-checkpoint guard in base._load_from_state_dict
+# --------------------------------------------------------------------------
+
+_COMPOSE_KEYS = ("_theta_shift", "_theta_scale", "_compose_standardization")
+
+
+def test_checkpoint_legacy_no_compose_keys_loads_as_off():
+ """All 3 compose buffers absent (legacy pre-compose checkpoint) -> loads fine
+ as compose-OFF with identity affine."""
+ batch_x, batch_y = _batches()
+ src = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=False,
+ )
+ sd = src.state_dict()
+ for k in _COMPOSE_KEYS:
+ sd.pop(k, None)
+
+ dst = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=False,
+ )
+ dst.load_state_dict(sd) # must not raise
+ assert bool(dst._compose_standardization) is False
+ assert torch.allclose(dst._theta_shift, torch.zeros_like(dst._theta_shift))
+ assert torch.allclose(dst._theta_scale, torch.ones_like(dst._theta_scale))
+
+
+def test_checkpoint_partial_compose_keys_raises():
+ """compose=True checkpoint MISSING _theta_scale -> raise on load (no silent
+ identity injection)."""
+ batch_x, batch_y = _batches()
+ src = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=True,
+ )
+ sd = src.state_dict()
+ sd.pop("_theta_scale", None) # partial: shift + flag present, scale missing
+
+ dst = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=False,
+ )
+ with pytest.raises((RuntimeError, ValueError), match="_theta_scale"):
+ dst.load_state_dict(sd)
+
+
+def test_checkpoint_all_compose_keys_loads():
+ """All 3 compose buffers present (compose checkpoint) -> normal load."""
+ batch_x, batch_y = _batches()
+ src = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=True,
+ )
+ sd = src.state_dict()
+ dst = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=False,
+ )
+ dst.load_state_dict(sd) # must not raise
+ assert bool(dst._compose_standardization) is True
+ assert torch.allclose(dst._theta_shift, src._theta_shift)
+ assert torch.allclose(dst._theta_scale, src._theta_scale)
+
+
+# --------------------------------------------------------------------------
+# T1.4 — sample_batched compose guard
+# --------------------------------------------------------------------------
+
+
+def test_sample_batched_rejects_compose():
+ """sample_batched is untested/unsupported under compose (used by the batched
+ score correction in theta-space) -> raise NotImplementedError up front."""
+ from torch.distributions import Independent, Normal
+
+ from sbi.inference.posteriors.vector_field_posterior import VectorFieldPosterior
+
+ batch_x, batch_y = _batches()
+ est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=True,
+ )
+ prior = Independent(Normal(torch.zeros(NUM_DIM), torch.ones(NUM_DIM)), 1)
+ posterior = VectorFieldPosterior(vector_field_estimator=est, prior=prior)
+
+ x_batch = torch.randn(2, NUM_DIM) # batch of 2 observations
+ with pytest.raises(NotImplementedError, match="compose_standardization"):
+ posterior.sample_batched(torch.Size([3]), x=x_batch)
+
+
+# --------------------------------------------------------------------------
+# Nit 1 — compose-OFF preservation for structured z-scoring
+# --------------------------------------------------------------------------
+
+
+def test_helper_compose_off_structured_zscore():
+ """compose OFF + z_score structured -> scalar/structured stats match
+ z_standardization(batch, structured_dims=True). The structured branch is
+ not covered by the other T1.1 helper tests, which use independent/none."""
+ batch_x, _ = _batches()
+ mean_0, std_0, shift, scale = _compute_theta_standardization(
+ batch_x, z_score_x="structured", compose_standardization=False
+ )
+ exp_mean, exp_std = z_standardization(batch_x, structured_dims=True)
+ assert torch.allclose(torch.as_tensor(mean_0), exp_mean)
+ assert torch.allclose(torch.as_tensor(std_0), exp_std)
+ assert shift is None and scale is None
+
+
+# --------------------------------------------------------------------------
+# Nit 2 — nested/prefixed checkpoint load exercises prefix + name path
+# --------------------------------------------------------------------------
+
+
+class _WrappedEstimator(torch.nn.Module):
+ """Minimal wrapper that places the estimator at a non-root prefix."""
+
+ def __init__(self, est):
+ super().__init__()
+ self.est = est
+
+
+def test_checkpoint_prefixed_legacy_loads_as_off():
+ """Wrapped estimator (prefix='est.'): all 3 compose buffers stripped ->
+ loads as compose-OFF with identity affine (exercises prefix + name path)."""
+ batch_x, batch_y = _batches()
+ src_est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=False,
+ )
+ src = _WrappedEstimator(src_est)
+ sd = src.state_dict()
+ _compose_keys = ("_theta_shift", "_theta_scale", "_compose_standardization")
+ for k in list(sd.keys()):
+ if any(k.endswith(n) for n in _compose_keys):
+ del sd[k]
+
+ dst_est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=False,
+ )
+ dst = _WrappedEstimator(dst_est)
+ dst.load_state_dict(sd) # must not raise
+ assert bool(dst.est._compose_standardization) is False
+ assert torch.allclose(dst.est._theta_shift, torch.zeros_like(dst.est._theta_shift))
+ assert torch.allclose(dst.est._theta_scale, torch.ones_like(dst.est._theta_scale))
+
+
+def test_checkpoint_prefixed_partial_raises():
+ """Wrapped estimator (prefix='est.'): compose flag present, _theta_scale
+ stripped -> raises naming the missing key (exercises prefix + name path)."""
+ batch_x, batch_y = _batches()
+ src_est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=True,
+ )
+ src = _WrappedEstimator(src_est)
+ sd = src.state_dict()
+ # Remove only the scale; keep shift + flag (partial).
+ for k in list(sd.keys()):
+ if k.endswith("_theta_scale"):
+ del sd[k]
+
+ dst_est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=False,
+ )
+ dst = _WrappedEstimator(dst_est)
+ with pytest.raises((RuntimeError, ValueError), match="_theta_scale"):
+ dst.load_state_dict(sd)
+
+
+# --------------------------------------------------------------------------
+# Nit 3 — sharper log_abs_det test with hand-set known scale
+# --------------------------------------------------------------------------
+
+
+def test_log_abs_det_known_scale():
+ """Non-tautological: hand-set _theta_scale to [2.0, 3.0, 5.0] and assert
+ log_abs_det() == log(2)+log(3)+log(5) exactly."""
+ batch_x, batch_y = _batches()
+ est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=True,
+ )
+ known_scale = torch.tensor([[[2.0, 3.0, 5.0]]]) # shape (1, 1, NUM_DIM)
+ est._theta_scale.copy_(known_scale.reshape(1, NUM_DIM))
+ expected = torch.log(torch.tensor([2.0, 3.0, 5.0])).sum()
+ assert torch.allclose(est.log_abs_det(), expected)
+
+
+# --------------------------------------------------------------------------
+# Deterministic single-obs log_prob Jacobian correction (no training)
+# --------------------------------------------------------------------------
+
+
+def test_single_obs_log_prob_jacobian_exact():
+ """Deterministic (training-free) pin of the affine-Jacobian log_prob
+ correction in the single-obs branch of ``VectorFieldBasedPotential.__call__``.
+
+ Wire a tiny compose estimator into the potential, set a single observation
+ (builds the neural ODE -- no training), set ``_theta_scale`` to a known
+ [2.0, 3.0], and monkeypatch the underlying ``flow.log_prob`` to return a known
+ constant. This exercises the REAL ``log_abs_det()`` and the REAL subtraction
+ ``log_probs = log_probs - log_abs_det`` in the code path (not a
+ reimplementation), so the returned potential must equal
+ ``known_constant - log(2) - log(3)`` to tight tolerance.
+ """
+ from torch.distributions import Independent, Normal
+
+ from sbi.inference.potentials.vector_field_potential import (
+ VectorFieldBasedPotential,
+ )
+
+ torch.manual_seed(0)
+ dim = 2
+ batch_x = 100.0 + 5.0 * torch.randn(32, dim)
+ batch_y = torch.randn(32, dim)
+ est = build_vector_field_estimator(
+ batch_x,
+ batch_y,
+ estimator_type="flow",
+ z_score_x="independent",
+ compose_standardization=True,
+ )
+ # Wide prior so theta=0 lies in support (within_support uses original theta).
+ prior = Independent(Normal(torch.zeros(dim), 100.0 * torch.ones(dim)), 1)
+ potential = VectorFieldBasedPotential(est, prior=prior, x_o=None, device="cpu")
+ potential.set_x(torch.zeros(1, dim), x_is_iid=False) # builds flow, no training
+
+ # Known per-dim scale -> log_abs_det = log(2) + log(3).
+ est._theta_scale.copy_(torch.tensor([[2.0, 3.0]]))
+
+ known_constant = 7.0
+ # Monkeypatch the z-space flow log_prob used in the single-obs branch.
+ potential.flow.log_prob = lambda z: torch.full(z.shape[:-1], known_constant)
+
+ out = potential(torch.zeros(1, dim))
+ expected = known_constant - torch.log(torch.tensor([2.0, 3.0])).sum()
+ assert torch.allclose(out.reshape(-1), expected.reshape(-1), atol=1e-6), (
+ f"log_prob Jacobian correction off: got {out.item()}, expected "
+ f"{expected.item()}"
+ )
diff --git a/tests/test_scale_equivariance.py b/tests/test_scale_equivariance.py
new file mode 100644
index 000000000..202c4b863
--- /dev/null
+++ b/tests/test_scale_equivariance.py
@@ -0,0 +1,306 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see
+
+"""Regression tests for scale-equivariance of FMPE / NPSE posteriors (#1680).
+
+Companion to ``tests/linearGaussian_vector_field_test.py``. A calibrated amortized
+posterior must be invariant under an exact reparameterization ``theta -> c * theta``
+(the simulator unscales, so the x-distribution is identical). FMPE and NPSE default
+posteriors are NOT equivariant when parameters are far from O(1) scale (the flow/SDE
+integrate against a unit base), so calibration collapses despite the default
+``z_score_theta="independent"``.
+
+The fix validated here is the OPT-IN ``compose_standardization`` flag: the estimator
+trains and samples in standardized z-space with an invertible PER-DIM affine
+``theta = shift + scale * z`` composed at the boundaries (loss input standardized;
+samples unstandardized; log_prob corrected by the affine Jacobian). Per-dim handles
+HETEROGENEOUS scales (mixed O(1) and O(1e-6) dims).
+
+Scope: single observation. iid / guided sampling are guarded (NotImplementedError)
+and not tested here. Runtime: marked slow.
+"""
+
+import contextlib
+
+import numpy as np
+import pytest
+import torch
+from torch.distributions import Independent, Normal
+
+from sbi.inference import FMPE, NPSE
+from sbi.neural_nets.factory import posterior_flow_nn, posterior_score_nn
+from sbi.utils import BoxUniform
+
+N_TRAIN, N_POST, MAX_EPOCHS = 2000, 2000, 150
+
+
+def _linear_gaussian(s, seed=0):
+ """Homogeneous linear-Gaussian: theta ~ N(0, s^2 I_2), x = sum(theta) + s * noise.
+ Posterior coupling rho = -0.5, tightness ~0.816 (both scale-free)."""
+ torch.manual_seed(seed)
+ np.random.seed(seed)
+ prior = Independent(Normal(torch.zeros(2), s * torch.ones(2)), 1)
+ theta = prior.sample((N_TRAIN,))
+ x = theta.sum(1, keepdim=True) + s * torch.randn(N_TRAIN, 1)
+ return prior, theta, x
+
+
+def _heterogeneous(seed=0):
+ """Mixed-scale prior: two O(1) dims and two O(1e-6) dims (cfg62-like). A scalar
+ noise-schedule rescale cannot bracket both scales; per-dim composed
+ standardization can. The simulator's second output pins (theta2 - theta3) at the
+ small scale, so at the observation x=[*, 0] the posterior must place
+ theta2 ~ theta3 -> STRONG POSITIVE coupling of the small-scale block."""
+ torch.manual_seed(seed)
+ np.random.seed(seed)
+ scale = torch.tensor([1.0, 1.0, 1e-6, 1e-6])
+ prior = Independent(Normal(torch.zeros(4), scale), 1)
+ theta = prior.sample((N_TRAIN,))
+ x = torch.stack(
+ [theta[:, 0] + theta[:, 1], (theta[:, 2] - theta[:, 3]) / 1e-6], dim=-1
+ ) + 0.01 * torch.randn(N_TRAIN, 2)
+ return prior, theta, x, scale
+
+
+def _fit(kind, prior, theta, x, compose):
+ if kind == "FMPE":
+ de = posterior_flow_nn(compose_standardization=True) if compose else "mlp"
+ tr = FMPE(prior=prior, vf_estimator=de)
+ else:
+ de = (
+ posterior_score_nn(sde_type="ve", compose_standardization=True)
+ if compose
+ else "mlp"
+ )
+ tr = NPSE(prior=prior, sde_type="ve", vf_estimator=de)
+ tr.append_simulations(theta, x).train(
+ max_num_epochs=MAX_EPOCHS, show_train_summary=False
+ )
+ return tr.build_posterior()
+
+
+def _coupling_tightness(post, s, x_o):
+ a = post.sample((N_POST,), x=x_o, show_progress_bars=False).numpy()
+ rho = float(np.corrcoef(a[:, 0], a[:, 1])[0, 1])
+ tight = float(a[:, 0].std()) / s
+ return rho, tight
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("kind", ["FMPE", "NPSE"])
+@pytest.mark.parametrize("seed", [0, 1, 2])
+def test_compose_standardization_homogeneous_equivariant(kind, seed):
+ """GREEN (multi-seed): composed standardization recovers calibration at small
+ scale and matches the unit-scale reference (coupling + tightness), for both
+ FMPE and NPSE."""
+ prior_u, th_u, x_u = _linear_gaussian(1.0, seed=seed)
+ rho_u, _ = _coupling_tightness(
+ _fit(kind, prior_u, th_u, x_u, True), 1.0, torch.zeros(1)
+ )
+
+ s = 1e-5
+ prior_i, th_i, x_i = _linear_gaussian(s, seed=seed)
+ rho_i, tight_i = _coupling_tightness(
+ _fit(kind, prior_i, th_i, x_i, True), s, torch.zeros(1)
+ )
+ assert tight_i < 3.0, f"[seed={seed}] ill-scale tightness inflated: {tight_i}"
+ assert rho_i < -0.15, f"[seed={seed}] ill-scale coupling collapsed: {rho_i}"
+ assert abs(rho_i - rho_u) < 0.3, (
+ f"[seed={seed}] ill vs unit mismatch: {rho_i} vs {rho_u}"
+ )
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("kind", ["FMPE", "NPSE"])
+def test_default_collapses_without_compose(kind):
+ """RED guard: with the flag OFF (default) the posterior collapses at small scale
+ -- documents the bug #1681 left unfixed and that the fix is opt-in (no behavior
+ change)."""
+ s = 1e-5
+ prior_i, th_i, x_i = _linear_gaussian(s)
+ rho_i, tight_i = _coupling_tightness(
+ _fit(kind, prior_i, th_i, x_i, False), s, torch.zeros(1)
+ )
+ assert tight_i > 3.0 or rho_i > -0.15, (
+ f"expected collapse without compose: rho={rho_i}, tight={tight_i}"
+ )
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("kind", ["FMPE", "NPSE"])
+def test_compose_standardization_heterogeneous_recovers(kind):
+ """GREEN + comparator: on a mixed-scale (O(1)+O(1e-6)) problem, composed
+ standardization (per-dim) recovers the SIGN-CORRECT strong POSITIVE coupling of
+ the small-scale block (x pins theta2 ~ theta3), where the default collapses it
+ (rho ~ 0)."""
+ prior, theta, x, scale = _heterogeneous()
+ x_o = torch.tensor([[0.0, 0.0]])
+
+ a_fix = (
+ _fit(kind, prior, theta, x, True)
+ .sample((N_POST,), x=x_o, show_progress_bars=False)
+ .numpy()
+ )
+ rho_fix = float(np.corrcoef(a_fix[:, 2], a_fix[:, 3])[0, 1])
+ tight_fix = float(a_fix[:, 2].std()) / float(scale[2])
+
+ a_def = (
+ _fit(kind, prior, theta, x, False)
+ .sample((N_POST,), x=x_o, show_progress_bars=False, reject_outside_prior=False)
+ .numpy()
+ )
+ rho_def = float(np.corrcoef(a_def[:, 2], a_def[:, 3])[0, 1])
+
+ # Recovery: strong POSITIVE coupling (sign-correct), contracted small-scale block.
+ assert tight_fix < 5.0, f"small-scale block not contracted: {tight_fix}"
+ assert rho_fix > 0.7, f"composed std failed to recover positive coupling: {rho_fix}"
+ # Comparator: default leaves the small-scale block far less coupled.
+ assert rho_fix - rho_def > 0.3, (
+ f"compose did not beat default on coupling: fix={rho_fix} def={rho_def}"
+ )
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("kind", ["FMPE", "NPSE"])
+def test_compose_log_prob_jacobian_scaling(kind):
+ """Affine-Jacobian correctness: an equivariant posterior satisfies
+ p_s(theta) = s^{-d} p_1(theta/s), so at the (symmetric) mode theta=0,
+ log p_s(0) - log p_1(0) = -d * log(s). This checks the log_prob Jacobian
+ correction (-sum log scale) is applied with the right magnitude/sign (d=2)."""
+ d, s = 2, 1e-3
+ prior_u, th_u, x_u = _linear_gaussian(1.0)
+ lp_u = float(
+ _fit(kind, prior_u, th_u, x_u, True)
+ .log_prob(torch.zeros(1, d), x=torch.zeros(1))
+ .item()
+ )
+ prior_s, th_s, x_s = _linear_gaussian(s)
+ lp_s = float(
+ _fit(kind, prior_s, th_s, x_s, True)
+ .log_prob(torch.zeros(1, d), x=torch.zeros(1))
+ .item()
+ )
+
+ expected = -d * np.log(s) # ~ +13.8
+ assert np.isfinite(lp_u) and np.isfinite(lp_s), (
+ f"non-finite log_prob: {lp_u}, {lp_s}"
+ )
+ assert abs((lp_s - lp_u) - expected) < 4.0, (
+ f"Jacobian scaling off: lp_s-lp_u={lp_s - lp_u:.2f} vs expected {expected:.2f}"
+ )
+
+
+class _RecordingSupport:
+ """Wraps a distribution's ``support`` and records every tensor passed to
+ ``check`` (the call the VFPE rejection path makes via ``within_support``)."""
+
+ def __init__(self, support, sink):
+ self._support = support
+ self._sink = sink
+
+ def check(self, value):
+ self._sink.append(value.detach().clone())
+ return self._support.check(value)
+
+
+class _RecordingPrior:
+ """Thin proxy around a prior that records the coordinates handed to the
+ prior-support check during rejection sampling. ``within_support`` reaches the
+ prior via ``prior.support.check(theta)``; we intercept exactly that so the test
+ can assert the recorded coordinate is in THETA space (post ``from_z``), not
+ z-space. ``log_prob`` is also recorded as a fallback path."""
+
+ def __init__(self, prior, sink):
+ self._prior = prior
+ self._sink = sink
+ self.support = _RecordingSupport(prior.support, sink)
+
+ def log_prob(self, value):
+ self._sink.append(value.detach().clone())
+ return self._prior.log_prob(value)
+
+ def __getattr__(self, name):
+ return getattr(self._prior, name)
+
+
+@pytest.mark.slow
+def test_compose_boxuniform_rejection_in_original_theta_space():
+ """Direct coordinate oracle: the prior-support check during rejection must
+ receive THETA-space coordinates (post ``from_z``), not z-space ones.
+
+ The other sample tests use Normal priors, whose ``within_support`` is always
+ True, so they cannot detect whether prior rejection happens in z-space or in
+ original-theta space. Under composed standardization the SDE/ODE proposal
+ samples in z-space (~O(1)); the sampler must map z -> theta (``from_z``) BEFORE
+ rejecting against the prior support.
+
+ We choose a BoxUniform support FAR from unit scale (|theta| <= 3e-4) so that
+ theta-space coordinates (O(1e-4)) and z-space coordinates (O(1)) are numerically
+ disjoint. A recording proxy on the prior captures the exact tensor passed to the
+ support check; we assert its magnitude is consistent with the THETA box, not z.
+ A regression that rejected in z-space (``from_z`` removed before rejection)
+ would feed O(1) values here and fail by a clear assertion (not a timeout).
+
+ Tiny train (few sims/epochs); marked slow.
+ """
+ torch.manual_seed(0)
+ np.random.seed(0)
+ bound = 3e-4
+ low = -bound * torch.ones(2)
+ high = bound * torch.ones(2)
+ prior = BoxUniform(low=low, high=high)
+
+ n_train = 500
+ theta = prior.sample((n_train,))
+ x = theta.sum(1, keepdim=True) + bound * torch.randn(n_train, 1)
+
+ de = posterior_flow_nn(compose_standardization=True)
+ tr = FMPE(prior=prior, vf_estimator=de)
+ tr.append_simulations(theta, x).train(max_num_epochs=20, show_train_summary=False)
+ post = tr.build_posterior()
+
+ # Spy on the coordinates handed to the prior-support check during rejection.
+ recorded: list = []
+ post.prior = _RecordingPrior(prior, recorded)
+
+ # The proxy records the coordinate on the FIRST rejection batch, so the oracle
+ # below does not depend on collecting any accepted samples. A z-space regression
+ # would reject almost everything (O(1) proposals vs the |theta|<=3e-4 box) and
+ # could otherwise time out instead of asserting; cap the time and tolerate a
+ # raise so detection is the DIRECT coordinate assertion, never a timeout.
+ samples = None
+ # A z-space regression rejects almost everything and may raise on timeout; the
+ # recorded coordinate is already captured, so suppress and let the assertion judge.
+ with contextlib.suppress(RuntimeError):
+ samples = post.sample(
+ (8,),
+ x=torch.zeros(1),
+ show_progress_bars=False,
+ max_sampling_time=20.0,
+ return_partial_on_timeout=True,
+ )
+
+ # The oracle: the rejection path must have queried the prior at least once,
+ # and the queried coordinate must be in the THETA box, NOT at z scale.
+ assert recorded, "prior-support check was never invoked during rejection"
+ max_recorded = max(float(t.abs().max()) for t in recorded)
+ # Theta-space magnitude is ~bound (here ~5e-4); z-space would be ~O(1) (>> bound).
+ # Allow slack above `bound` for the proposal's pre-rejection tails, but stay far
+ # below the O(1) z scale so a z-space regression fails by THIS assertion.
+ z_scale_floor = 1e-1
+ assert max_recorded < z_scale_floor, (
+ "prior-support check received z-space coordinates "
+ f"(max|recorded|={max_recorded:.3e} >= {z_scale_floor:g}); "
+ "from_z must be applied BEFORE prior rejection"
+ )
+
+ # And any returned (theta-space) samples lie within the ORIGINAL BoxUniform box.
+ if samples is not None and samples.numel() > 0:
+ assert torch.all(samples >= low), (
+ "sample below BoxUniform low (z-space rejection?): "
+ f"min={samples.min().item()}"
+ )
+ assert torch.all(samples <= high), (
+ "sample above BoxUniform high (z-space rejection?): "
+ f"max={samples.max().item()}"
+ )