Skip to content

feat: support frozen dataclasses in apply_to_collection(s) via reconstruction#499

Open
deependujha wants to merge 7 commits into
Lightning-AI:mainfrom
deependujha:fix/apply-to-collection-frozen-dataclass-reconstruct
Open

feat: support frozen dataclasses in apply_to_collection(s) via reconstruction#499
deependujha wants to merge 7 commits into
Lightning-AI:mainfrom
deependujha:fix/apply-to-collection-frozen-dataclass-reconstruct

Conversation

@deependujha

Copy link
Copy Markdown
Contributor
Before submitting
  • Was this discussed/agreed via a Github issue? (no need for typos and docs improvements)
  • Did you read the contributor guideline, Pull Request section?
  • Did you make sure to update the docs?
  • Did all existing and newly added tests pass locally?

What does this PR do?

Fixes # (issue).

Makes apply_to_collection and apply_to_collections work with frozen dataclasses
by 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 raises FrozenInstanceError on frozen
instances. 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:

  • Mutable dataclasses keep the existing deepcopy + setattr path, completely
    unchanged. This preserves current behavior for InitVar, init=False, and
    cached_property reset.
  • Frozen dataclasses are reconstructed: the function is applied recursively to each
    init=True field, then a new instance is built via cls(**init_fields). Fields with
    init=False are written back from the original via object.__setattr__ (matching the
    retain-old behavior of the mutable path, rather than letting __post_init__ recompute them).

A shared helper _reconstruct_frozen_dataclass is used by both apply_to_collection
and apply_to_collections, so the singular/plural asymmetry is gone — both now support
frozen dataclasses.

Unsupported case: frozen + InitVar

Frozen dataclasses that declare an InitVar field cannot be reconstructed: InitVar
values are consumed by __post_init__ and not stored on the instance, so they're
unavailable to pass back through the constructor. These raise a clear ValueError
rather than failing with a confusing constructor TypeError. (ClassVar is correctly
not treated as InitVar — detection uses _field_type is _FIELD_INITVAR, not a field-count diff.)

Mutable + InitVar is unaffected, since the mutable path never reconstructs.

Breaking change: allow_frozen removed

The allow_frozen parameter is removed from the public signature. Its old behavior
(break on the first frozen field and return a partially transformed copy) is obsolete
now that frozen dataclasses reconstruct fully. To avoid hard-breaking callers, it is
absorbed via **kwargs, popped before forwarding to function, and emits a
DeprecationWarning. Passing allow_frozen=True is 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.

@codecov-commenter

codecov-commenter commented Jun 28, 2026

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 97.22222% with 1 line in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@c89c402). Learn more about missing BASE report.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

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:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Frozen dataclass as a batch

2 participants