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

[BHARATH0153]

Summary

Remove all vector field estimator changes from this PR. Keep only the transform_to_unconstrained z-scoring for MDN.

Changes

• MixtureDensityEstimator: Added prior_transform parameter for change-of-variables via mcmc_transform. log_prob includes the Jacobian correction (log_prob(z) + log|det J⁻¹|). sample applies the inverse transform to map back to constrained space. The mutual-exclusion guard transform_input ↔ prior_transform ensures only one is active.
• build_mdn: Wires prior_transform = mcmc_transform(x_dist) when z_score_x="transform_to_unconstrained". Added x_dist as explicit parameter.
• Bugfix: Fixed double-sum in _log_det_jacobian — the IndependentTransform(reinterpreted_batch_ndims=1) already sums over the event dimension, so the explicit .sum(dim=-1) was duplicating the sum.
• Naming: Renamed from input_transform → prior_transform to avoid confusion with the existing transform_input parameter.
• Docs: Restored omitted docstrings for has_input_transform, _transform_input, _inverse_transform_input, _log_det_jacobian_forward.
• Type hints: Uses TorchTransform from sbi.sbi_types.
• Tests: Added test_mdn_transform_to_unconstrained that tests the change-of-variables formula directly against the underlying MoG.

Closes #1885

---

The VF estimator changes will be pursued in a separate PR by @janfb

[BHARATH0153]

@janfb please review at your convenience thanks!

[janfb]

Thanks for picking this up, @BHARATH0153!

I suggest to split this into two PRs, because the two halves are at very different stages:

1. MDN: The structure is clean, mcmc_transform direction, the change-of-variables in log_prob, the inverse on sample, and the transform_input/input_transform mutual-exclusion guard are all good. There's one correctness bug and a test gap (details inline). Both are small; if you fix them this half is mergeable.

2. Vector field (deferred). Please drop the vector-field changes from this PR for now. I checked it locally and it currently crashes at sampling, and more importantly it overlaps with #1884 (composed standardization for scale-equivariant FMPE/NPSE), which changes the same standardization path. I'm assigned to #1885 and will take the VF part myself once #1884 is settled and the ODE-space design is worked out.

So overall, please trim this PR to MDN-only with the fix + a test below, and we can get it in. Thanks again, this is a useful contribution!

[codecov]

Codecov Report

❌ Patch coverage is 78.57143% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.86%. Comparing base (c444a92) to head (2545833).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
sbi/inference/trainers/npe/npe_a.py	0.00%	2 Missing ⚠️
sbi/inference/trainers/npe/npe_c.py	50.00%	1 Missing ⚠️
...eural_nets/estimators/mixture_density_estimator.py	93.33%	1 Missing ⚠️
sbi/neural_nets/net_builders/mdn.py	85.71%	1 Missing ⚠️
sbi/utils/conditional_density_utils.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1888      +/-   ##
==========================================
- Coverage   87.90%   87.86%   -0.04%     
==========================================
  Files         143      143              
  Lines       13366    13387      +21     
==========================================
+ Hits        11749    11763      +14     
- Misses       1617     1624       +7     

Flag	Coverage Δ
fast	81.39% <78.57%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
sbi/inference/posteriors/npe_a_posterior.py	56.81% <ø> (ø)
sbi/inference/trainers/npe/npe_c.py	93.20% <50.00%> (-0.55%)	⬇️
...eural_nets/estimators/mixture_density_estimator.py	91.76% <93.33%> (-0.12%)	⬇️
sbi/neural_nets/net_builders/mdn.py	96.42% <85.71%> (-3.58%)	⬇️
sbi/utils/conditional_density_utils.py	92.72% <50.00%> (-0.53%)	⬇️
sbi/inference/trainers/npe/npe_a.py	77.70% <0.00%> (-1.65%)	⬇️

[janfb]

Thanks again for the PR @BHARATH0153 . Here the inline review.

Please revert all changes to the vector field related methods. And please revert the removed docstrings. The remaining todos are variable naming and testing issues, and the jac summing bug.
Looking forward to your fixes!

[BHARATH0153]

@janfb thanks for the review by today 12 I will make all changes as per review thanks for the review again

[BHARATH0153]

@janfb as requested all changes are done please review thanks!

[BHARATH0153]

@janfb please rerun the checks all are resolved thanks!

[janfb]

Thanks for the quick quixes, @BHARATH0153! The MDN estimator itself now looks right: the Jacobian fix is correct (I verified log_prob returns the right shape), the prior_transform rename is clean, the docstrings are back, and the new test is also in 👍

When reviewing this again I found two two more blocking bugs, both from the same root cause: the PR changes two things that the MDN shares with the NPE-A / NPE-C machinery, without updating the other side. They don't show up in the diff because the broken code is in files the PR doesn't touch. Details inline, summary here:

1. _log_det_jacobian_forward now takes a required transformed_input, but npe_a_posterior.py:173 still calls it with one argument. That's an unconditional TypeError in NPE-A's corrected log_prob — and note it breaks the default affine z-scoring too, not just the new option. To reproduce, try: NPE_A(prior).append_simulations(theta, x).train(); post.log_prob(theta, x=x_o) → TypeError: ... missing 1 required positional argument: 'transformed_input'.

2. has_input_transform now returns True for the prior_transform path, but _transform_shift/_transform_scale are None there. npe_a.py:~358, npe_c.py (_set_maybe_z_scored_prior), and conditional_density_utils.py:~196 all read those buffers when has_input_transform is true. And this leads to a TypeError: unsupported operand type(s) for -: 'Tensor' and 'NoneType'. The property used to mean "affine z-score is on" (since Refactor MDN implementation: port MDN from pyknos, improve numerical stability #1724) and those three consumers rely on that.

The reason bug 2 hides is that the analytic NPE-A / NPE-C correction is mathematically only valid for an affine z-score, it can't represent a nonlinear mcmc_transform in closed form. So the right fix isn't to make those paths handle prior_transform; it's to let them raise a clear error, and keep transform_to_unconstrained scoped to what the MDN can do correctly on its own: direct log_prob/sample and single-round NPE (those go through the estimator's own methods and already work).

So, I'd suggest:

• fix the npe_a_posterior.py caller (one line),
• add a narrow has_affine_z_score property and have the three analytic consumers gate on it, raising NotImplementedError for the prior_transform case,
• plus the device/buffer point and a ruff pass (inline).

More generally, this is now reaching into the NPE-A/NPE-C correction internals, which is a bit beyond where I wanted this PR to go. Do let me know if you prefer that I implement those fixes, I can commit to your branch.

[BHARATH0153]

@janfb done I implemented once re check please at your convenience and please commit if anything going wrong

and thanks for the great review
So the right fix isn't to make those paths handle prior_transform; it's to let them raise a clear error, and keep transform_to_unconstrained scoped to what the MDN can do correctly on its own: direct log_prob/sample and single-round NPE (those go through the estimator's own methods and already work). and this is true and I have fixed by your suggestion thanks please commit if you have any thing

[janfb]

Thanks again @BHARATH0153 , I made specific suggestions to fix the new bugs. let me know how it goes.

[janfb]

@janfb done I implemented once re check please at your convenience and please commit if anything going wrong

@BHARATH0153 you were too fast, I posted my inline comments just know and I suggest to not use a new property, but just raise an error in the caller sites.

[BHARATH0153]

@janfb sorry I had terrible network and trying to do fast by the data finish and two comments are adressed dint used a new property sorry for using before

[BHARATH0153]

@janfb now all tests are passing tested locally

[janfb]

Converging! one last fix and we are good to merge @BHARATH0153 🚀

[BHARATH0153]

@janfb done please check thanks!

[janfb]

All good, well done @BHARATH0153! First PR merged 🎉