fix: only invert inverse prior transforms in MDN pickle (#1893 review)

The MDN __getstate__ stored self._prior_transform.inv unconditionally, which is
correct for the inverse transform mcmc_transform produces but silently corrupts a
concrete forward transform (as the __init__ contract allows): inverting it wraps it
in an inverse whose tensors torch then drops on pickle. Guard on _InverseTransform
so concrete transforms pickle as-is. Adds a regression test.

Author: janfb

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
@@ -413,17 +414,20 @@ def _apply(self, fn):
         return self
 
     def __getstate__(self):
-        # torch drops an inverse transform's tensors on pickling; store the forward
-        # and re-invert on load.
         state = dict(super().__getstate__())
-        if self._prior_transform is not None:
-            state["_prior_transform"] = self._prior_transform.inv
+        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 self._prior_transform is not None:
-            self._prior_transform = self._prior_transform.inv
+        if forward is not None:
+            self._prior_transform = forward.inv
 
     @property
     def embedding_net(self) -> nn.Module:

tests/save_and_load_test.py

@@ -94,6 +94,35 @@ def test_mdn_transform_to_unconstrained_picklable():
     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.