harden+docs: guard zuko pickle re-link; document transform persistence (#1893)

- ZukoFlow.__setstate__ now fails loudly (RuntimeError) if a prior transform
  existed before pickling but cannot be recovered from the flow on load (e.g. if
  a future zuko version changes the internal `Partial.f` structure the re-link
  relies on), instead of silently degrading to None and dropping device moves.
- Note the prior transform's pickle/device behavior and state_dict omission in
  the build_mdn and ZukoFlow docstrings.

Author: janfb

sbi/neural_nets/estimators/zuko_flow.py

@@ -41,7 +41,10 @@ def __init__(
             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()``; its tensors are
-                buried in a plain attribute that ``nn.Module`` cannot reach.
+                buried in a plain attribute that ``nn.Module`` cannot reach. It 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()``.
         """
 
         # assert len(condition_shape) == 1, "Zuko Flows require 1D conditions."
@@ -61,18 +64,34 @@ 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__).
+        # transform correctly, and we re-link to it on load (see __setstate__). Keep
+        # a flag so __setstate__ can fail loudly if the re-link ever comes up empty.
         state = dict(super().__getstate__())
-        if self._prior_transform is not None:
-            state["_prior_transform"] = None
+        state["_had_prior_transform"] = self._prior_transform is not None
+        state["_prior_transform"] = None
         return state
 
     def __setstate__(self, state):
+        had_prior_transform = state.pop("_had_prior_transform", False)
         super().__setstate__(state)
         self._prior_transform = self._find_prior_transform()
+        if had_prior_transform and 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."""
+        """Recover the prior transform from the flow's CallableTransform, if any.
+
+        Assumes at most one `CallableTransform` in the flow (true today: only the
+        single `transform_to_unconstrained` prior transform is wrapped this way).
+        Returns the first match; revisit if `_prepare_x_transforms` ever composes
+        multiple `CallableTransform`s.
+        """
         for module in self.net.modules():
             wrapped = getattr(module, "f", None)
             if isinstance(wrapped, CallableTransform):

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.