feat: support frozen dataclasses in apply_to_collection(s) via reconstruction#499
Open
deependujha wants to merge 7 commits into
Open
Conversation
|
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #499 +/- ##
=====================================
Coverage ? 73%
=====================================
Files ? 17
Lines ? 773
Branches ? 0
=====================================
Hits ? 565
Misses ? 208
Partials ? 0 🚀 New features to boost your workflow:
|
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
This was referenced Jun 28, 2026
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
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.
Before submitting
What does this PR do?
Fixes # (issue).
Makes
apply_to_collectionandapply_to_collectionswork with frozen dataclassesby reconstructing a new instance from transformed fields, instead of raising
ValueError: A frozen dataclass was passed to apply_to_collection but this is not allowed.Previously, frozen dataclasses were unsupported: the implementation deep-copies the
input and mutates fields via
setattr, which raisesFrozenInstanceErroron frozeninstances. The prior decision (Lightning-AI/pytorch-lightning#10873) was to fail loudly,
because the only alternatives considered were mutating the instance (violates the point
of
frozen=True) or skipping the function (silently leaves tensors on the wrong device).This PR takes a third approach that neither earlier discussion considered:
reconstruct a new instance from the transformed fields. This honors immutability
(the original is never touched, a new instance is returned) and actually applies the
transformation (no silent skip), so it sidesteps both objections.
Fixes Lightning-AI/pytorch-lightning#21577
Related: Lightning-AI/pytorch-lightning#10873
How it works
The dataclass branch now splits on frozen-ness:
setattrpath, completelyunchanged. This preserves current behavior for
InitVar,init=False, andcached_propertyreset.init=Truefield, then a new instance is built viacls(**init_fields). Fields withinit=Falseare written back from the original viaobject.__setattr__(matching theretain-old behavior of the mutable path, rather than letting
__post_init__recompute them).A shared helper
_reconstruct_frozen_dataclassis used by bothapply_to_collectionand
apply_to_collections, so the singular/plural asymmetry is gone — both now supportfrozen dataclasses.
Unsupported case: frozen +
InitVarFrozen dataclasses that declare an
InitVarfield cannot be reconstructed:InitVarvalues are consumed by
__post_init__and not stored on the instance, so they'reunavailable to pass back through the constructor. These raise a clear
ValueErrorrather than failing with a confusing constructor
TypeError. (ClassVaris correctlynot treated as
InitVar— detection uses_field_type is _FIELD_INITVAR, not a field-count diff.)Mutable +
InitVaris unaffected, since the mutable path never reconstructs.Breaking change:
allow_frozenremovedThe
allow_frozenparameter is removed from the public signature. Its old behavior(
breakon the first frozen field and return a partially transformed copy) is obsoletenow that frozen dataclasses reconstruct fully. To avoid hard-breaking callers, it is
absorbed via
**kwargs, popped before forwarding tofunction, and emits aDeprecationWarning. Passingallow_frozen=Trueis now a no-op (frozen support is automatic).PR review
Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in GitHub issues there's a high chance it will not be merged.