Fix handling of OrderedDict with optree and related documentation.
#20481
Conversation
The tree API had specific but contradicting documentation calling out the handling of `OrderedDict`s. However, the behavior of the `optree` implementation did not honor this documentation (using the key order, not the sequence order) for `flatten`, although it did for `pack_sequence_as`. The result was that not only did `flatten` not behave the same with `optree` and `dm-tree`, but also `pack_sequence_as(flatten(...))` was not idempotent. The `optree` implementation did have all the machinery needed to handle `OrderedDict`s per spec, which was used for `pack_sequence_as`, but not `flatten`. This also fixes the discrepancy in the behavior for `namedtuple`s. - Fixed contradicting documentation in `flatten` and `pack_sequence_as` related to the handling of `OrderedDict`s. - Fixed references to `unflatten_as`, which doesn't exist. - Removed most `if optree` tests in `tree_test.py`, which should not exist for consistency between `optree` and `dm-tree`. - Fixed unit tests which were incorrectly flattening the result of `flatten_with_path`. - Fixed unintented use of `tree` instead of `keras.tree` in unit test. - Ran unit tests for all backends with `dm-tree` uninstalled.
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## master #20481 +/- ##
==========================================
- Coverage 82.07% 76.40% -5.68%
==========================================
Files 515 515
Lines 47504 47512 +8
Branches 7454 7457 +3
==========================================
- Hits 38991 36303 -2688
- Misses 6703 9452 +2749
+ Partials 1810 1757 -53
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. |
| from absl.testing import parameterized | ||
|
|
||
| from keras.src import backend | ||
| from keras.src import metrics as losses_module | ||
| from keras.src import metrics as metrics_module | ||
| from keras.src import ops | ||
| from keras.src import testing | ||
| from keras.src import tree |
There was a problem hiding this comment.
This seems to be triggering a failure on the torch side https://github.com/keras-team/keras/actions/runs/11785192949/job/32825931002?pr=20481
(inability to recognize a NamedTuple as a loss tuple)
There was a problem hiding this comment.
Looking into it. For some reason, this test required dm-tree.
Testing further, I tested with optree (the default) uninstalled and before this change, and a number of tests fail.
So some tests (and maybe some code) are still tree implementation dependent.
The tree API had specific but contradicting documentation calling out the handling of
OrderedDicts. However, the behavior of theoptreeimplementation did not honor this documentation (using the key order, not the sequence order) forflatten, although it did forpack_sequence_as. The result was that not only didflattennot behave the same withoptreeanddm-tree, but alsopack_sequence_as(flatten(...))was not idempotent. Theoptreeimplementation did have all the machinery needed to handleOrderedDicts per spec, which was used forpack_sequence_as, but notflatten. This also fixes the discrepancy in the behavior fornamedtuples.flattenandpack_sequence_asrelated to the handling ofOrderedDicts.unflatten_as, which doesn't exist.if optreetests intree_test.py, which should not exist for consistency betweenoptreeanddm-tree.flatten_with_path.treeinstead ofkeras.treein unit test.dm-treeuninstalled.