Fix torch.compile _tree_map recompilations - #30
Merged
Merged
Conversation
added 6 commits
February 25, 2026 22:34
Add a compile-time fast path to avoid error-wrapping overhead and document expected torch.compile behavior with new tests.
Add a detailed investigation write-up for _tree_map recompilation and a test script that exercises frame guard behavior across compiled scenarios.
Add is_compiling fast paths to call pytree.tree_map directly and avoid Dynamo recompile limits, plus tests and docs to verify behavior.
document why `_tree_map_fn()` and `substitute_in_graph` approaches do not prevent recompiles in real workloads, even when unit tests pass. clarify the proven fix as a two-layer defense: per-method `is_compiling()` fast paths as the primary solution, plus module-level `_tree_map_impl` polyfill support as structural backup. update the investigation matrix and test strategy section to reflect the new failure modes and current validation approach.
add author metadata entry in project configuration
Delete the malformed `author` entry in `pyproject.toml` to keep project metadata valid for packaging tools. Also reformat long compile-related call sites in shape operations and test files for readability, with no intended behavior changes.
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
1 participant
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.
Summary
This PR includes the commits on
recompilesthat are not yet inorigin/main, focused on resolving Dynamo recompilation behavior around_tree_mapand finalizing release metadata._tree_mapand usepytree.tree_mapdirectly in affected TensorContainer operations, preventing callable-identity guard churn._tree_mapbehavior so existing error wrapping and diagnostics remain available outside compiled regions.docs/compile-recompilation-investigation.mdwith root cause analysis, alternatives considered, and rationale for the selected fix.0.9.0and fixes invalidpyproject.tomlauthor metadata.Why
Real workloads were hitting
torch._dynamorecompile limits due to_tree_mapbeing compiled as a separate frame with changing callable identities. The compile-path bypass removes this failure mode while preserving eager-path diagnostics and keeping behavior stable.Validation
_tree_mapbypass behavior undertorch.compile