Repository navigation
fix(torch): canonicalize nested dict key order in the torch tree implementation - #23201
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a short-circuit optimization in the PyTorch tree implementation's _dict_to_ordered_dict function to return leaves immediately, avoiding unnecessary round-trips. It also adds corresponding unit tests for various data structures. The feedback suggests extending the test suite to include a nested dictionary test case to ensure recursive, deterministic sorting is thoroughly verified.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Codecov Reportโ
All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #23201 +/- ##
==========================================
- Coverage 86.88% 84.15% -2.74%
==========================================
Files 477 478 +1
Lines 71493 77360 +5867
Branches 11825 13535 +1710
==========================================
+ Hits 62120 65099 +2979
- Misses 6239 9216 +2977
+ Partials 3134 3045 -89
Flags with carried forward coverage won't be shown. Click here to find out more. โ View full report in Codecov by Harness. ๐ New features to boost your workflow:
|
e85aa77 to
b451829
Compare
b451829 to
674e82e
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds a short-circuit optimization to _dict_to_ordered_dict in torchtree_impl.py to return leaf structures immediately, avoiding an unnecessary round-trip. A comprehensive set of unit tests has been added in tree_test.py to validate this behavior across various structures such as tensors, lists, dicts, and nested structures. No review comments were provided, so there is no additional feedback.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
โฆile-fix fold Reworked eager A/B charts for keras-team#23189 (P4+C1), keras-team#23201 (P18+C2), keras-team#23198 (P13+C3) post-fold, plus the combined-13-base cumulative C1/C1+C2/ C1+C2+C3 graph-break and torch.compile-speedup decomposition chart.
โฆs full stack, eager vs compile)
|
This PR is stale because it has been open for 14 days with no activity. It will be closed if no further activity occurs. Thank you. |
test_dict_to_ordered_dict_leaf_shortcircuit previously only checked a top-level dict and a nested dict-of-lists whose keys happened to already be sorted, so it would not catch a sorting/reassociation regression at a nested level. Add a case with a dict-of-dict where keys are out of order at both the outer and inner level, flatten and repack it, and verify each leaf lands back at its correct nested key, matching an equivalent structure whose keys are already sorted at every level.
_dict_to_ordered_dict only sorted the outermost dict level: a dict match short-circuited before the recursive traversal that would have processed a nested dict value, so nested dicts kept insertion order instead of the deterministic order other tree implementations (optree, dmtree) provide. This produced different leaf orders for structurally-equivalent nested dicts with different insertion order, and pack_sequence_as could silently swap values when packing one structure's leaves into a differently-ordered equivalent structure. Values are now recursed into while building the sorted dict, so every level is canonicalized, not just the top one.
The dict-canonicalization comment in _dict_to_ordered_dict had a duplicated "that" across the line wrap and joined two clauses with a bare hyphen instead of a full sentence. The leaf-shortcircuit test's comment referenced an internal root-cause tracking id that isn't meaningful outside that tracker.
โฆuit test test_dict_to_ordered_dict_leaf_shortcircuit's dict-of-dict block checked that each of two differently-ordered but structurally-equivalent nested dicts round-trips correctly on its own, which holds even when a nested dict level is left unsorted, since flatten and pack_sequence_as apply the same canonicalization consistently within a single round-trip. Add a direct comparison of the two structures' flattened leaf lists, which is the property that actually depends on every nesting level being canonicalized and not just the outermost one.
torch.compile/dynamo cannot trace id() comparisons, so the one-level flatten in _dict_to_ordered_dict and traverse compared child identity against the root's id() and graph-broke under compilation. Replace it with a root-flag closure that tracks whether the root node has been visited yet, with no id() compare at all. The closure is deduplicated into a single module-level factory, _make_root_is_leaf, used by both traverse_children call sites (previously each defined its own copy). Each call to the factory returns a fresh closure instance, so nested/ repeated traversal calls do not leak seen_root state across each other -- semantics are unchanged from the per-call id(structure) scoping this replaces. Adds regression coverage in tree_test.py: traverse/ _dict_to_ordered_dict equivalence over nested dict/list/namedtuple structures, and a root-visited-exactly-once check guarding against the root being mistakenly re-split into its own children.
095726f to
3c48dbd
Compare
pctablet505
changed the title
โฆpile fix _dict_to_ordered_dict now recurses into dict values, so key order is canonical at every level rather than only at the root. traverse canonicalizes once at the root and walks the already-canonical structure, so the recursion is not repeated per level. The id()-comparison change for torch.compile is removed from this branch and will be sent as its own PR, as requested.
|
cannot mark the PR ready for review due to new limits enforced on number of open pull requests. |
google-ml-butler
Bot
added
kokoro:force-run
ready to pull
google-ml-butler
Bot
removed
awaiting review
ready to pull
Closes #23281
Summary
In the torch tree implementation,
_dict_to_ordered_dictonly sorted the keys of the outermost dict.func(structure)returns a dict for any dict input, so theif ret is None: return traverse_children()branch below it was unreachable for dicts and the values were never visited.Nested dicts therefore kept their insertion order, which makes two structurally identical inputs flatten differently:
The last one is the damaging case: values are silently paired with the wrong keys.
_dict_to_ordered_dictnow recurses, so key order is canonical at every level.traversecanonicalizes once at the root and walks the already-canonical structure, so the recursion is not repeated per level.Cost
Sorting at every level is work master skips, so dict-heavy structures get slower. But the rewrite also stops round-tripping every structure through
torch_tree.tree_flatten/tree_unflatten, and an internal node is now rebuilt only when a dict below it actually changed, so non-dict structures get faster. Measured against master, median of 9 rounds:flattenlist of 12map_structurenestedpack_sequence_asnestedflattennested dictflattenflat dict of 16flattenleafTests
tree_test.pycovers nested flatten order,pack_sequence_asround-tripping, and equality of the two structurally identical dicts above. The nested-order test fails on master and passes here.Contributor Agreement