[pctablet505]

Description

Fixes #23648

Split out of #23201 at @hertschuh's request. On that PR the torch.compile hunk was reviewed
alongside the dict-ordering fix, and the ask was:

It looks like this is a separate issue of compilation. Make this a separate PR.

— #23201 (comment)

#23201 (the nested dict key-order fix, #23281) has since merged without it. This PR is that
carve-out, rebased on current master. No behavior from #23201 is changed here.

The problem

keras/src/tree/torchtree_impl.py has two identity checks that Dynamo refuses to trace, so
keras.tree.traverse — and anything else routed through _dict_to_ordered_dict
(flatten_with_path, pack_sequence_as, map_shape_structure, lists_to_tuples) — always
graph-breaks:

1. _traverse flattens one level with is_leaf=lambda x: id(x) != structure_id. Dynamo models
   id() as a compile-time-only value and breaks on the comparison.
2. _dict_to_ordered_dict decides whether to rebuild a node with if ordered_child is not child:,
   which on dicts raises Unsupported: ... builtin `is_not` with argument types ['dict', 'dict'].

Neither is data-dependent, so the break fires on every traced call over a structure containing a
dict or nested container.

The fix

Both checks only answer "is this the root?" and "did anything below me change?", which do not
need identity comparisons:

1. _traverse now walks one level through torch_tree.SUPPORTED_NODES[...]'s
   flatten_fn/unflatten_fn directly. This is the pattern _tree_is_leaf and
   _dict_to_ordered_dict already use in this same module, and it removes the tree_flatten
   round trip along with the TreeSpec that was built and discarded immediately.
2. The dict-ordering recursion moved into _dict_to_ordered_dict_with_changed, which returns an
   explicit (ordered_structure, changed) pair. _dict_to_ordered_dict stays as the thin
   public-facing wrapper, so its callers are untouched.

Note on (1): an earlier revision of this patch kept tree_flatten and passed a stateful is_leaf
closure that returned False for the first node it saw and True afterwards. That works, but
only because tree_flatten happens to call is_leaf exactly once per node in pre-order with the
root first — an undocumented implementation detail, and one that would silently produce a
full flatten rather than an error if it ever changed. Going through the node type directly has
no such dependency.

Benchmarks

torch 2.11.0. torch._dynamo.explain over tree.traverse on [tensor, {"b": tensor, "a": tensor}]:

	master	this PR
graphs	2	1
graph breaks	1	0
fullgraph=True	raises Unsupported	passes

Dropping the tree_flatten/tree_unflatten round trip also helps eager, median of 10k calls over
[{"b": 1, "a": [2, 3]}, (4, 5), 6]:

	master	this PR
tree.traverse	14.29 µs	5.25 µs	2.7x

Tests

test_traverse_torch_compile_fullgraph in keras/src/tree/tree_test.py compiles a tree.traverse
over a nested list/dict with fullgraph=True and checks both the values and the canonical key
order. It raises torch._dynamo.exc.Unsupported on master and passes here. It is skipped for the
non-torch tree implementations.

The rest of tree_test.py is unchanged and passes (29 passed, 21 skipped), which is what pins the
equivalence of the new one-level walk. I also diffed the old and new _traverse across 112
combinations (14 structure shapes — leaf, None, empty containers, OrderedDict, defaultdict,
deque, namedtuple, nested dicts — x 4 callbacks x top_down both ways) with identical results.

Contributor Agreement

• This PR is linked to an issue that has been assigned to me (see Fixes #xxx above).
• I am a human, and not a bot.
• I will be responsible for responding to review comments in a timely manner.
• I will work with the maintainers to push this PR forward until submission.

[gemini-code-assist]

Code Review

This pull request refactors torchtree_impl.py to make traverse compatible with torch.compile(..., fullgraph=True). It replaces id() comparisons in _traverse with a stateful closure _make_root_is_leaf to prevent graph breaks in PyTorch Dynamo, and optimizes _dict_to_ordered_dict to avoid rebuilding dictionaries if they are already ordered. A corresponding unit test is added to verify full-graph compilation. The feedback suggests merging the duplicate sorting and traversal logic for dict and defaultdict to reduce code duplication and improve maintainability.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 84.67%. Comparing base (58db6d6) to head (0764bf0).
⚠️ Report is 29 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #23647      +/-   ##
==========================================
- Coverage   86.95%   84.67%   -2.28%     
==========================================
  Files         492      492              
  Lines       71781    78225    +6444     
  Branches    11838    13595    +1757     
==========================================
+ Hits        62416    66236    +3820     
- Misses       6243     8897    +2654     
+ Partials     3122     3092      -30     

Flag	Coverage Δ
keras	84.48% <100.00%> (-2.29%)	⬇️
keras-cpu	84.22% <100.00%> (-1.90%)	⬇️
keras-gpu	70.88% <100.00%> (+0.08%)	⬆️
keras-jax	60.15% <0.00%> (-1.03%)	⬇️
keras-numpy	54.08% <0.00%> (-0.16%)	⬇️
keras-openvino	60.53% <0.00%> (-0.17%)	⬇️
keras-tensorflow	61.77% <0.00%> (-0.99%)	⬇️
keras-torch	61.40% <100.00%> (-1.02%)	⬇️
keras-tpu	58.22% <0.00%> (+0.02%)	⬆️
keras.applications	?
keras.applications-jax	?
keras.applications-numpy	?
keras.applications-openvino	?
keras.applications-tensorflow	?
keras.applications-torch	?
keras.wrappers	?
keras.wrappers-jax	?
keras.wrappers-numpy	?
keras.wrappers-openvino	?
keras.wrappers-tensorflow	?
keras.wrappers-torch	?

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hertschuh]

Thanks!