Added support for pytree types that inherit from tuple and typing.namedtuple (#2845)

2025-12-04 11:06:45 -08:00
parent 50d3914c67
commit 9abb0b8123
6 changed files with 196 additions and 3 deletions
@@ -44,10 +44,11 @@ def tree_map(
        return fn(tree, *rest)
    elif isinstance(tree, (list, tuple)):
        TreeType = type(tree)
-        return TreeType(
+        subtrees = (
            tree_map(fn, child, *(r[i] for r in rest), is_leaf=is_leaf)
            for i, child in enumerate(tree)
        )
+        return TreeType(*subtrees) if hasattr(tree, "_fields") else TreeType(subtrees)
    elif isinstance(tree, dict):
        return {
            k: tree_map(fn, child, *(r[k] for r in rest), is_leaf=is_leaf)