#23648·keras

tree.traverse can't be captured by torch.compile on the torch backend

Author: pctablet505Created Sep 16, 2026Updated Sep 16, 2026
Labelsbackend:torch

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:

  1. _traverse flattens exactly one level with is_leaf=lambda x: id(x) != structure_id. Dynamo models id() as a compile-time-only value and graph-breaks on the comparison.
  2. _dict_to_ordered_dict decides whether to rebuild an internal node with if ordered_child is not child:. On dicts this raises torch._dynamo.exc.Unsupported: Failed to trace builtin operator ... 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 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[...]'s flatten_fn/unflatten_fn instead of going through tree_flatten with an is_leaf predicate. This is the pattern _tree_is_leaf and _dict_to_ordered_dict already use in the same module, and it also skips building a TreeSpec that is discarded immediately.
  • Return an explicit changed flag from the dict-ordering recursion instead of comparing rebuilt children by identity.

Part of the per-call Python-dispatch overhead series in #22561.