fix: transform_to_unconstrained device mdn device handling (#1893)

[BHARATH0153]

Fixes MDN part of #1893

Two changes:

1. mcmc_transform bounded branch now moves transform tensor internals to the requested device (was CPU-only via biject_to)
2. MixtureDensityEstimator._apply overridden so _prior_transform follows .to(), .cuda(), .mps() moves

[codecov]

Codecov Report

❌ Patch coverage is 57.14286% with 18 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.82%. Comparing base (dbf432b) to head (4a1e306).
⚠️ Report is 20 commits behind head on main.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
sbi/utils/sbiutils.py	52.77%	17 Missing ⚠️
...eural_nets/estimators/mixture_density_estimator.py	83.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1896      +/-   ##
==========================================
- Coverage   87.88%   87.82%   -0.07%     
==========================================
  Files         143      143              
  Lines       13387    13474      +87     
==========================================
+ Hits        11765    11833      +68     
- Misses       1622     1641      +19     

Flag	Coverage Δ
fast	81.30% <57.14%> (?)

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

Files with missing lines	Coverage Δ
...eural_nets/estimators/mixture_density_estimator.py	91.47% <83.33%> (-0.29%)	⬇️
sbi/utils/sbiutils.py	84.98% <52.77%> (-3.45%)	⬇️

... and 1 file with indirect coverage changes

[BHARATH0153]

@janfb please review when you have time thanks! and anything needed

[janfb]

Thanks @BHARATH0153 , looks good! It's focused on MDN for now and it works nicely. I just suggest a bit of refactoring towards a more general _move_tensor function.

The zuko and state-dict problems are not fixed here, which is fine. We can do it in a follow up.

[janfb]

One more thing: could you add test coverage for this? Codecov has it at ~20% atm. The test should be @pytest.mark.gpu (it'll auto-skip locally and run on the GPU job, like the other tests in tests/inference_on_device_test.py): build an MDN with z_score_x="transform_to_unconstrained", est.to("cuda"), then assert log_prob runs and the transform's tensors are on CUDA or MPS.
Thanks!

[BHARATH0153]

@janfb once look into this please

[janfb]

Thanks @BHARATH0153 , this is exactly the refactor I had in mind, the core code looks good now. Now there are only things on the new tests that need to be fixed, and one small cleanup, then we can merge soon.

[BHARATH0153]

@janfb thanks! for the review done please check.

[janfb]

Looks good, well done @BHARATH0153 ! 👏