fix: transform_to_unconstrained device mdn device handling (#1893)

sbi/neural_nets/estimators/mixture_density_estimator.py

@@ -22,6 +22,7 @@
 from sbi.neural_nets.estimators.base import ConditionalDensityEstimator
 from sbi.neural_nets.estimators.mog import MoG
 from sbi.sbi_types import TorchTransform
+from sbi.utils.sbiutils import _apply_to_transform
 
 
 class MultivariateGaussianMDN(nn.Module):
@@ -405,6 +406,12 @@ def __init__(
             self.register_buffer("_transform_shift", None)
             self.register_buffer("_transform_scale", None)
 
+    def _apply(self, fn):
+        super()._apply(fn)
+        if self._prior_transform is not None:
+            _apply_to_transform(self._prior_transform, fn)
+        return self
+
     @property
     def embedding_net(self) -> nn.Module:
         """Return the embedding network."""

sbi/utils/sbiutils.py

@@ -811,6 +811,69 @@ def match_theta_and_x_batch_shapes(theta: Tensor, x: Tensor) -> Tuple[Tensor, Te
     return theta_repeated, x_repeated
 
 
+def _apply_to_transform(
+    transform: TorchTransform, fn: Callable[[Tensor], Tensor]
+) -> None:
+    """Apply fn to all tensors in a transform tree.
+
+    Walks ComposeTransform.parts, IndependentTransform.base_transform,
+    etc., so .to() calls propagate into the transform.
+
+    Args:
+        transform: Root of the transform tree.
+        fn: Callable applied to each tensor (e.g. ``lambda t: t.to(device)``).
+    """
+    seen = set()
+
+    def _walk(t):
+        if id(t) in seen:
+            return
+        seen.add(id(t))
+        for key, val in list(t.__dict__.items()):
+            if isinstance(val, Tensor):
+                object.__setattr__(t, key, fn(val))
+            elif isinstance(val, (list, tuple)):
+                for item in val:
+                    if isinstance(item, TorchTransform):
+                        _walk(item)
+            elif isinstance(val, TorchTransform):
+                _walk(val)
+
+    _walk(transform)
+
+
+def _transform_tensors(
+    transform: TorchTransform,
+) -> List[Tensor]:
+    """Collect every tensor in a transform tree (mirror of _apply_to_transform).
+
+    Args:
+        transform: Root of the transform tree.
+
+    Returns:
+        List of all tensors found in the transform tree.
+    """
+    seen = set()
+    tensors: List[Tensor] = []
+
+    def _walk(t):
+        if id(t) in seen:
+            return
+        seen.add(id(t))
+        for val in t.__dict__.values():
+            if isinstance(val, Tensor):
+                tensors.append(val)
+            elif isinstance(val, (list, tuple)):
+                for item in val:
+                    if isinstance(item, TorchTransform):
+                        _walk(item)
+            elif isinstance(val, TorchTransform):
+                _walk(val)
+
+    _walk(transform)
+    return tensors
+
+
 def mcmc_transform(
     prior: Distribution,
     num_prior_samples_for_zscoring: int = 1000,
@@ -925,6 +988,9 @@ def prior_mean_std_transform(prior, device):
 
     check_transform(prior, transform)  # type: ignore
 
+    if enable_transform:
+        _apply_to_transform(transform, lambda t: t.to(device))
+
     return transform.inv  # type: ignore
 
 

tests/inference_on_device_test.py

@@ -891,3 +891,56 @@ def test_npe_pfn_on_device(prior_device):
     assert posterior.estimator._context_input.device.type == "cpu", (
         "TabPFN context must always remain on CPU."
     )
+
+
+@pytest.mark.gpu
+def test_mdn_device_transform():
+    """MDN with transform_to_unconstrained moves transform tensors on .to()."""
+    from sbi.neural_nets.net_builders.mdn import build_mdn
+    from sbi.utils.sbiutils import _transform_tensors
+
+    device = process_device("gpu")
+    prior = BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
+    bx, by = prior.sample((512,)), torch.randn(512, 3)
+    est = build_mdn(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
+    est.to(device)
+
+    transform_tensors = _transform_tensors(est._prior_transform)
+    assert transform_tensors, "expected the prior transform to hold tensors"
+    for t in transform_tensors:
+        assert t.device.type == device.split(":")[0], (
+            f"transform tensor on {t.device}, expected {device}"
+        )
+
+    theta = prior.sample((5,)).to(device)
+    cond = torch.randn(1, 3).to(device)
+    lp = est.log_prob(theta.unsqueeze(1), cond)
+    assert lp.device.type == device.split(":")[0]
+    s = est.sample((10,), cond)
+    assert s.device.type == device.split(":")[0]
+
+
+def test_mdn_transform_follows_dtype():
+    """transform_to_unconstrained transform follows dtype casts on the MDN.
+
+    Guards the callable-fn design in _apply_to_transform: a rebuild-from-prior
+    would leave the transform in float32 and silently desync it from the weights.
+    Runs on every CI (no GPU needed).
+    """
+    from sbi.neural_nets.net_builders.mdn import build_mdn
+    from sbi.utils.sbiutils import _transform_tensors
+
+    prior = BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
+    bx, by = prior.sample((256,)), torch.randn(256, 3)
+    est = build_mdn(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
+
+    est.double()
+
+    transform_tensors = _transform_tensors(est._prior_transform)
+    assert transform_tensors, "expected the prior transform to hold tensors"
+    assert all(t.dtype == torch.float64 for t in transform_tensors)
+
+    theta = prior.sample((5,)).double()
+    cond = torch.randn(1, 3).double()
+    lp = est.log_prob(theta.unsqueeze(1), cond)
+    assert lp.dtype == torch.float64