fix: make transform_to_unconstrained picklable (#1893)

sbi/neural_nets/estimators/mixture_density_estimator.py

@@ -17,6 +17,7 @@
 import numpy as np
 import torch
 from torch import Tensor, nn
+from torch.distributions.transforms import _InverseTransform
 from torch.nn import functional as F
 
 from sbi.neural_nets.estimators.base import ConditionalDensityEstimator
@@ -412,6 +413,22 @@ def _apply(self, fn):
             _apply_to_transform(self._prior_transform, fn)
         return self
 
+    def __getstate__(self):
+        state = dict(super().__getstate__())
+        t = self._prior_transform
+        # torch drops an inverse transform's data (its `_inv`) on pickling; stash the
+        # wrapped forward and rebuild on load. Concrete transforms pickle fine as-is.
+        if isinstance(t, _InverseTransform):
+            state["_prior_transform"] = None
+            state["_prior_transform_forward"] = t.inv
+        return state
+
+    def __setstate__(self, state):
+        forward = state.pop("_prior_transform_forward", None)
+        super().__setstate__(state)
+        if forward is not None:
+            self._prior_transform = forward.inv
+
     @property
     def embedding_net(self) -> nn.Module:
         """Return the embedding network."""

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):
@@ -40,7 +40,9 @@ def __init__(
             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()``.
+                reference so it follows ``.to()``/``.double()``. Preserved by
+                pickling the estimator (sbi's save path), not by ``state_dict()``;
+                reconstruct via the builder before ``load_state_dict()``.
         """
 
         # assert len(condition_shape) == 1, "Zuko Flows require 1D conditions."
@@ -56,6 +58,38 @@ def _apply(self, fn):
             _apply_to_transform(self._prior_transform, fn)
         return self
 
+    def __getstate__(self):
+        # Store the forward (picklable); on load, re-link to the flow's own copy
+        # rather than this one, so device moves stay in sync (see __setstate__).
+        state = dict(super().__getstate__())
+        if self._prior_transform is not None:
+            state["_prior_transform"] = self._prior_transform.inv
+        return state
+
+    def __setstate__(self, state):
+        super().__setstate__(state)
+        if self._prior_transform is not None:
+            self._prior_transform = self._find_prior_transform()
+            if self._prior_transform is None:
+                raise RuntimeError(
+                    "ZukoFlow lost its prior transform while unpickling: it could "
+                    "not be recovered from the flow. This likely means the installed "
+                    "`zuko` version changed the flow's internal structure (the "
+                    "transform is located via `zuko`'s `Partial.f` attribute in "
+                    "`_find_prior_transform`)."
+                )
+
+    def _find_prior_transform(self) -> Optional[TorchTransform]:
+        """Recover the prior transform from the flow's CallableTransform, if any.
+
+        Assumes at most one CallableTransform in the flow (true today).
+        """
+        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/neural_nets/net_builders/mdn.py

@@ -44,7 +44,10 @@ def build_mdn(
             over the entire batch, instead of per-dimension. Should be used when each
             sample is, for example, a time series or an image.
             - `transform_to_unconstrained`: transform inputs to unconstrained space
-            using the prior's support. Requires `x_dist` to be provided.
+            using the prior's support. Requires `x_dist` to be provided. This
+            transform follows `.to()`/`.cuda()`/`.double()` and is preserved by
+            pickling the estimator (the standard sbi save path), but is not part of
+            `state_dict()`; reconstruct via the builder before `load_state_dict()`.
         z_score_y: Whether to z-score ys passing into the network, same options as
             z_score_x.
         hidden_features: Number of hidden features.

sbi/utils/sbiutils.py

@@ -295,6 +295,14 @@ def __init__(self, transform):
     def __call__(self):
         return self.transform
 
+    def __getstate__(self):
+        # torch drops an inverse transform's tensors on pickling; store the forward
+        # and re-invert on load.
+        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

@@ -68,3 +68,79 @@ def test_picklability(
         pickle.dump(inference, handle)
     with open(f"{tmp_path}/saved_inference.pickle", "rb") as handle:
         _ = pickle.load(handle)
+
+
+def test_mdn_transform_to_unconstrained_picklable():
+    """An MDN using transform_to_unconstrained survives a pickle round-trip.
+
+    Regression: sbi stores the prior transform as an inverse transform, whose data
+    lives behind the `_inv` link that torch's Transform.__getstate__ nulls on
+    pickling. Naively pickling therefore discards the transform's tensors and yields
+    an object that raises "_inv must not be None" on first use.
+    """
+    from sbi.neural_nets.net_builders.mdn import build_mdn
+    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_mdn(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)
+
+
+def test_mdn_pickle_preserves_concrete_prior_transform():
+    """A concrete (non-inverse) prior transform on an MDN survives pickling.
+
+    The `_inv`-dropping workaround only applies to inverse transforms. A concrete
+    forward transform (e.g. a user-supplied AffineTransform, as the __init__ contract
+    allows) must be pickled as-is; inverting it would wrap it in an inverse whose data
+    torch then drops on pickle.
+    """
+    from torch.distributions.transforms import AffineTransform
+
+    from sbi.neural_nets.net_builders.mdn import build_mdn
+    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_mdn(bx, by)
+    # A concrete forward transform (maps constrained -> unconstrained), not the
+    # inverse wrapper that mcmc_transform produces.
+    est._prior_transform = AffineTransform(torch.zeros(2), 2 * torch.ones(2))
+
+    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)
+
+
+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)