tree.traverse can't be captured by torch.compile on the torch backend
On the torch backend, keras.tree.traverse and anything routed through _dict_to_ordered_dict (flatten_with_path, pack_sequence_as, map_shape_structure, lists_to_tuples) cannot be captured into a single graph by torch.compile, because keras/src/tree/torchtree_impl.py uses two patterns that Dynamo refuses to trace:
_traverseflattens exactly one level withis_leaf=lambda x: id(x) != structure_id. Dynamo modelsid()as a compile-time-only value and graph-breaks on the comparison._dict_to_ordered_dictdecides whether to rebuild an internal node withif ordered_child is not child:. On dicts this raisestorch._dynamo.exc.Unsupported: Failed to trace builtin operator ... builtinis_notwith argument types ['dict', 'dict'].
Neither is data-dependent, so the break fires on every traced call over a structure containing a dict or a nested container. Measured on torch 2.11.0 with torch._dynamo.explain over tree.traverse on [tensor, {"a": tensor}]: 2 graphs, 1 graph break, and fullgraph=True raises outright.
Both checks exist only to answer "is this the root?" and "did anything below me change?", which can be answered without identity comparisons:
- Walk the one level directly through
torch_tree.SUPPORTED_NODES[...]'sflatten_fn/unflatten_fninstead of going throughtree_flattenwith anis_leafpredicate. This is the pattern_tree_is_leafand_dict_to_ordered_dictalready use in the same module, and it also skips building aTreeSpecthat is discarded immediately. - Return an explicit
changedflag from the dict-ordering recursion instead of comparing rebuilt children by identity.
Part of the per-call Python-dispatch overhead series in #22561.
Source: keras-team/keras