fix: make transform_to_unconstrained picklable (#1893)#1933
Open
janfb wants to merge 5 commits into
Open
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## fix/1893-zuko-device-statedict #1933 +/- ##
==================================================================
+ Coverage 87.88% 87.90% +0.01%
==================================================================
Files 143 143
Lines 13420 13450 +30
==================================================================
+ Hits 11794 11823 +29
- Misses 1626 1627 +1
Flags with carried forward coverage won't be shown. Click here to find out more.
|
janfb
force-pushed
the
fix/1893-zuko-device-statedict
branch
from
July 13, 2026 15:22
5ea4bad to
e7b32c3
Compare
janfb
force-pushed
the
fix/1893-transform-pickle
branch
from
July 13, 2026 15:22
9b1e311 to
a353e7d
Compare
dgedon
reviewed
Jul 14, 2026
Comment on lines
+70
to
+71
| state["_had_prior_transform"] = self._prior_transform is not None | ||
| state["_prior_transform"] = None |
Collaborator
There was a problem hiding this comment.
can we not do it like for the MDN estimator to set the prior transform if it is not none? Then we could get rid of the the had_prior_transform flag?
| "`_find_prior_transform`)." | ||
| ) | ||
|
|
||
| def _find_prior_transform(self) -> Optional[TorchTransform]: |
Collaborator
There was a problem hiding this comment.
Is this used anywhere else than for __set_state__()?
The prior transform is stored as an inverse transform whose data lives behind its `_inv` link, which torch's Transform.__getstate__ nulls on pickling — so a pickled MDN lost the transform entirely and raised "_inv must not be None" on first use. Store the forward transform in __getstate__ (it pickles cleanly, preserving device and dtype) and re-invert it in __setstate__.
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.
#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.
…m flag Store the forward transform as the signal a prior transform existed (as the MDN estimator does) instead of a separate _had_prior_transform flag; the flag was redundant since the flow already pickles the transform via CallableTransform. Also trim verbose inline comments and docstrings per review (#1893).
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.
janfb
force-pushed
the
fix/1893-transform-pickle
branch
from
July 15, 2026 06:02
a353e7d to
260fb1f
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #1932 (review/merge that first; this diff is only the last 3 commits).
Fixes pickling of MDN and Zuko estimators that use
transform_to_unconstrained. Currently the reloaded object crashes withAssertionError: _inv must not be Noneon first use: the transform is stored as an inverse transform whose data lives behind the_invlink that torch'sTransform.__getstate__nulls on pickling, so the tensors are silently dropped.Fix: store the forward transform (which pickles cleanly, keeping device and dtype) and re-invert on load — via
__getstate__/__setstate__on the MDN estimator and onCallableTransform. For Zuko,ZukoFlowre-links its reference to the flow's transform on load, and fails loudly if it can't (e.g. a future zuko internal change) rather than silently dropping device moves.Pickling the whole posterior is sbi's supported save path, so this matters more than the
state_dictgap — which we document instead of fixing (the transform is reconstructed via the builder; it's intentionally not carried instate_dict()).Tests: pickle round-trip for both estimators.
Closes #1893.
Written with Claude Code assistance.