fix: transform_to_unconstrained follows device moves for Zuko flows (#1893)

sbi/neural_nets/estimators/zuko_flow.py

@@ -1,7 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
-from typing import Tuple
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor, nn
@@ -11,7 +11,8 @@
     ConditionalDensityEstimator,
     UnconditionalDensityEstimator,
 )
-from sbi.sbi_types import Shape
+from sbi.sbi_types import Shape, TorchTransform
+from sbi.utils.sbiutils import _apply_to_transform
 
 
 class ZukoFlow(ConditionalDensityEstimator):
@@ -27,6 +28,7 @@ def __init__(
         embedding_net: nn.Module,
         input_shape: torch.Size,
         condition_shape: torch.Size,
+        prior_transform: Optional[TorchTransform] = None,
     ):
         r"""Initialize the conditional density estimator.
 
@@ -36,13 +38,23 @@ def __init__(
                 density is being evaluated (and which is also the event_shape
                 of samples).
             condition_shape: Event shape of the condition.
+            prior_transform: Optional unconstrained prior transform prepended to the
+                flow (for ``z_score_x="transform_to_unconstrained"``). Kept as a
+                reference so it follows ``.to()``/``.double()``.
         """
 
         # assert len(condition_shape) == 1, "Zuko Flows require 1D conditions."
         super().__init__(
             net=net, input_shape=input_shape, condition_shape=condition_shape
         )
         self._embedding_net = embedding_net
+        self._prior_transform = prior_transform
+
+    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:

sbi/neural_nets/net_builders/flow.py

@@ -18,6 +18,7 @@
 
 from sbi.neural_nets.estimators import NFlowsFlow, ZukoFlow, ZukoUnconditionalFlow
 from sbi.neural_nets.estimators.tabpfn_flow import TabPFNFlow
+from sbi.sbi_types import TorchTransform
 from sbi.utils.nn_utils import MADEMoGWrapper, get_numel
 from sbi.utils.sbiutils import (
     assert_transform_to_unconstrained_supported,
@@ -1148,7 +1149,7 @@ def build_zuko_flow(
     )
 
     # Get x transforms (z-score or logit transform)
-    x_transforms = _prepare_x_transforms(z_score_x, batch_x, x_dist)
+    x_transforms, prior_transform = _prepare_x_transforms(z_score_x, batch_x, x_dist)
 
     # Combine all transforms
     transforms = x_transforms + base_transforms
@@ -1164,6 +1165,7 @@ def build_zuko_flow(
         embedding_net,
         input_shape=batch_x[0].shape,
         condition_shape=batch_y[0].shape,
+        prior_transform=prior_transform,
     )
 
     return flow
@@ -1281,7 +1283,7 @@ def _prepare_x_transforms(
     ],
     batch_x: Tensor,
     x_dist: Optional[Distribution],
-) -> tuple:
+) -> Tuple[Tuple, Optional[TorchTransform]]:
     """
     Prepare transforms to prepend for x processing.
 
@@ -1291,9 +1293,12 @@ def _prepare_x_transforms(
         x_dist: Distribution for unconstrained transformation.
 
     Returns:
-        Tuple of transforms to prepend (empty tuple if no preprocessing).
+        Tuple ``(transforms, prior_transform)``: the transforms to prepend (empty
+        tuple if no preprocessing) and the raw unconstrained prior transform
+        (``None`` unless ``z_score_x == "transform_to_unconstrained"``).
     """
     transforms = ()
+    prior_transform = None
     z_score_x_bool, structured_x = z_score_parser(z_score_x)
     if z_score_x == "transform_to_unconstrained":
         if x_dist is None:
@@ -1306,13 +1311,13 @@ def _prepare_x_transforms(
                 "`x_dist` requires a `.support` attribute for"
                 "an unconstrained transformation."
             )
-        transform_to_unconstrained = biject_transform_zuko(mcmc_transform(x_dist))
-        transforms = (transform_to_unconstrained,)
+        prior_transform = mcmc_transform(x_dist, device=batch_x.device)
+        transforms = (biject_transform_zuko(prior_transform),)
     elif z_score_x_bool:
         z_score_transform = standardizing_transform_zuko(batch_x, structured_x)
         transforms = (z_score_transform,)
 
-    return transforms
+    return transforms, prior_transform
 
 
 def build_zuko_unconditional_flow(

sbi/utils/sbiutils.py

@@ -814,10 +814,12 @@ def match_theta_and_x_batch_shapes(theta: Tensor, x: Tensor) -> Tuple[Tensor, Te
 def _apply_to_transform(
     transform: TorchTransform, fn: Callable[[Tensor], Tensor]
 ) -> None:
-    """Apply fn to all tensors in a transform tree.
+    """Apply fn to all tensors in a transform tree, in place.
 
-    Walks ComposeTransform.parts, IndependentTransform.base_transform,
-    etc., so .to() calls propagate into the transform.
+    So ``.to()``/``.double()`` calls propagate into a transform held as a plain
+    attribute (e.g. a prior transform on an estimator), which is otherwise
+    invisible to ``nn.Module._apply``. Traverses ``ComposeTransform.parts``,
+    ``IndependentTransform.base_transform`` and nested inverses.
 
     Args:
         transform: Root of the transform tree.
@@ -842,38 +844,6 @@ def _walk(t):
     _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,

tests/inference_on_device_test.py

@@ -58,6 +58,25 @@
 )
 
 
+def _collect_transform_tensors(transform):
+    """Collect every tensor in a transform tree, reusing the production walk.
+
+    Test helper for the transform_to_unconstrained device/dtype checks: drives
+    ``_apply_to_transform`` with a recording identity fn so we don't duplicate the
+    tree-walk in the tests.
+    """
+    from sbi.utils.sbiutils import _apply_to_transform
+
+    collected = []
+
+    def _record(tensor):
+        collected.append(tensor)
+        return tensor
+
+    _apply_to_transform(transform, _record)
+    return collected
+
+
 @pytest.mark.slow
 @pytest.mark.gpu
 @pytest.mark.parametrize(
@@ -893,19 +912,21 @@ def test_npe_pfn_on_device(prior_device):
     )
 
 
-@pytest.mark.gpu
 def test_mdn_device_transform():
-    """MDN with transform_to_unconstrained moves transform tensors on .to()."""
+    """MDN with transform_to_unconstrained moves transform tensors on .to().
+
+    Runs on any available accelerator (CUDA or MPS); the module-level skipif
+    gates it off on CPU-only machines.
+    """
     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)
+    transform_tensors = _collect_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], (
@@ -928,19 +949,68 @@ def test_mdn_transform_follows_dtype():
     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)
+    transform_tensors = _collect_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
+
+
+def test_zuko_device_transform():
+    """Zuko flow with transform_to_unconstrained moves transform tensors on .to().
+
+    Runs on any available accelerator (CUDA or MPS); the module-level skipif
+    gates it off on CPU-only machines.
+    """
+    from sbi.neural_nets.net_builders.flow import build_zuko_maf
+
+    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_zuko_maf(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
+    est.to(device)
+
+    transform_tensors = _collect_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_zuko_transform_follows_dtype():
+    """Zuko transform_to_unconstrained transform follows dtype casts.
+
+    Runs on every CI (no GPU needed): guards that the prior transform is reachable
+    by nn.Module._apply through the estimator's _apply override. Unlike the network
+    weights, the transform lives behind a plain attribute, so without the override a
+    .double()/.to() cast would silently leave it on the old dtype/device.
+    """
+    from sbi.neural_nets.net_builders.flow import build_zuko_maf
+
+    prior = BoxUniform(-2 * torch.ones(2), 2 * torch.ones(2))
+    bx, by = prior.sample((256,)), torch.randn(256, 3)
+    est = build_zuko_maf(bx, by, z_score_x="transform_to_unconstrained", x_dist=prior)
+
+    est.double()
+
+    transform_tensors = _collect_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)

tests/sbiutils_test.py

@@ -709,3 +709,31 @@ def test_mdn_transform_to_unconstrained():
     assert torch.allclose(lp, mog.log_prob(z) + ldj, atol=1e-5)
     s = est.sample((10,), cond)
     assert s.shape[0] == 10 and torch.isfinite(s).all()
+
+
+def test_apply_to_transform_reaches_nested_tensors():
+    """`_apply_to_transform` reaches tensors behind lists, nesting and .inv.
+
+    The leaf `loc`/`scale` sit behind `_InverseTransform -> IndependentTransform ->
+    ComposeTransform.parts (a list) -> AffineTransform`, so a walk that missed the
+    list branch or the inverse would silently touch nothing. A dtype cast is the
+    observable effect: both leaves must change.
+    """
+    from torch.distributions.transforms import (
+        AffineTransform,
+        ComposeTransform,
+        SigmoidTransform,
+    )
+
+    from sbi.utils.sbiutils import _apply_to_transform
+
+    affine = AffineTransform(torch.zeros(3), torch.ones(3))
+    # matches the shape mcmc_transform produces for a bounded prior, then inverted.
+    transform = IndependentTransform(
+        ComposeTransform([SigmoidTransform(), affine]), 1
+    ).inv
+
+    assert affine.loc.dtype == torch.float32 and affine.scale.dtype == torch.float32
+    _apply_to_transform(transform, lambda t: t.double())
+    assert affine.loc.dtype == torch.float64
+    assert affine.scale.dtype == torch.float64