fix: make Zuko transform_to_unconstrained picklable (#1893)

Same _inv-nulling issue as the MDN case, but the transform is wrapped in a
CallableTransform inside the flow and is also referenced by
ZukoFlow._prior_transform. Give CallableTransform __getstate__/__setstate__ that
store the forward transform (pickles cleanly, keeps device/dtype), and have
ZukoFlow drop its redundant reference on pickling and re-link it to the flow's
CallableTransform on load, so both stay the same object and device moves after a
round-trip still work.

Author: janfb

sbi/neural_nets/estimators/zuko_flow.py

@@ -12,7 +12,7 @@
     UnconditionalDensityEstimator,
 )
 from sbi.sbi_types import Shape, TorchTransform
-from sbi.utils.sbiutils import _apply_to_transform
+from sbi.utils.sbiutils import CallableTransform, _apply_to_transform
 
 
 class ZukoFlow(ConditionalDensityEstimator):
@@ -57,6 +57,28 @@ def _apply(self, fn):
             _apply_to_transform(self._prior_transform, fn)
         return self
 
+    def __getstate__(self):
+        # `_prior_transform` is the same inverse-transform object the flow holds
+        # (inside a CallableTransform). Torch would drop its data on pickling, so
+        # don't pickle this redundant reference — the CallableTransform pickles the
+        # transform correctly, and we re-link to it on load (see __setstate__).
+        state = dict(super().__getstate__())
+        if self._prior_transform is not None:
+            state["_prior_transform"] = None
+        return state
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        self._prior_transform = self._find_prior_transform()
+
+    def _find_prior_transform(self) -> Optional[TorchTransform]:
+        """Recover the prior transform from the flow's CallableTransform, if any."""
+        for module in self.net.modules():
+            wrapped = getattr(module, "f", None)
+            if isinstance(wrapped, CallableTransform):
+                return wrapped.transform
+        return None
+
     @property
     def embedding_net(self) -> nn.Module:
         r"""Return the embedding network."""

sbi/utils/sbiutils.py

@@ -295,6 +295,16 @@ def __init__(self, transform):
     def __call__(self):
         return self.transform
 
+    def __getstate__(self):
+        # `transform` is an inverse transform whose data lives behind its `_inv`
+        # link, which torch's Transform.__getstate__ nulls on pickling. Store the
+        # forward instead (it pickles cleanly, keeping device/dtype) and re-invert
+        # on load, so a pickled Zuko flow keeps a working transform.
+        return {"transform": self.transform.inv}
+
+    def __setstate__(self, state):
+        self.transform = state["transform"].inv
+
 
 def biject_transform_zuko(
     transform: TorchTransform,

tests/save_and_load_test.py

@@ -92,3 +92,26 @@ def test_mdn_transform_to_unconstrained_picklable():
     actual = reloaded.log_prob(theta, cond)
 
     assert torch.allclose(expected, actual)
+
+
+def test_zuko_transform_to_unconstrained_picklable():
+    """A Zuko flow using transform_to_unconstrained survives a pickle round-trip.
+
+    Same root cause as the MDN case, but the transform is wrapped in a
+    CallableTransform inside the flow, and is additionally referenced by
+    ZukoFlow._prior_transform — both must come back consistent.
+    """
+    from sbi.neural_nets.net_builders.flow import build_zuko_maf
+    from sbi.utils import BoxUniform
+
+    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)
+
+    theta, cond = prior.sample((5,)).unsqueeze(1), torch.randn(1, 3)
+    expected = est.log_prob(theta, cond)
+
+    reloaded = pickle.loads(pickle.dumps(est))
+    actual = reloaded.log_prob(theta, cond)
+
+    assert torch.allclose(expected, actual)