image-match/python/py/Lib/site-packages/torch/_dynamo/polyfills/pytree.py

"""
Python polyfills for torch.utils.pytree
"""

from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field
from typing import Any, TYPE_CHECKING, TypeVar

import optree
import optree._C
import optree.utils
from optree import (
    is_namedtuple,
    is_namedtuple_class,
    is_namedtuple_instance,
    is_structseq,
    is_structseq_class,
    is_structseq_instance,
    namedtuple_fields,
    structseq_fields,
)

import torch.utils._cxx_pytree as cxx_pytree  # noqa: F401
import torch.utils._pytree as python_pytree
from torch.utils._pytree import BUILTIN_TYPES, STANDARD_DICT_TYPES

from ..decorators import substitute_in_graph


if TYPE_CHECKING:
    import builtins
    from collections.abc import Callable, Iterable, Mapping
    from typing_extensions import Self, TypeIs

    from torch.utils._cxx_pytree import PyTree


__all__ = [
    "is_namedtuple",
    "is_namedtuple_class",
    "is_namedtuple_instance",
    "is_structseq",
    "is_structseq_class",
    "is_structseq_instance",
    "namedtuple_fields",
    "structseq_fields",
    "treespec_leaf",
    "treespec_tuple",
    "treespec_dict",
    "tree_is_leaf",
    "tree_iter",
    "tree_leaves",
    "tree_flatten",
    "tree_flatten_with_path",
    "tree_structure",
    "tree_unflatten",
]


_T = TypeVar("_T")
_KT = TypeVar("_KT")
_VT = TypeVar("_VT")


@substitute_in_graph(
    optree._C.is_dict_insertion_ordered,
    can_constant_fold_through=True,
)
def _(*args: Any, **kwargs: Any) -> bool:
    # In namespace 'torch', the dictionary is always traversed in insertion order.
    # This function returns True.
    raise ValueError(
        "Should not be called directly "
        "because the original function will be called in the constant fold path."
    )


__name = ""
for __name, __func in (
    ("is_namedtuple", is_namedtuple),
    ("is_namedtuple_class", is_namedtuple_class),
    ("is_namedtuple_instance", is_namedtuple_instance),
    ("is_structseq", is_structseq),
    ("is_structseq_class", is_structseq_class),
    ("is_structseq_instance", is_structseq_instance),
    ("namedtuple_fields", namedtuple_fields),
    ("structseq_fields", structseq_fields),
):
    globals()[__name] = substitute_in_graph(
        __func,  # type: ignore[arg-type]
        can_constant_fold_through=True,
    )(__func.__python_implementation__)  # type: ignore[attr-defined]
    del __func
del __name


@substitute_in_graph(optree.tree_is_leaf, can_constant_fold_through=True)  # type: ignore[arg-type]
def tree_is_leaf(
    tree: PyTree,
    /,
    is_leaf: Callable[[PyTree], bool] | None = None,
    *,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> bool:
    if (tree is None and none_is_leaf) or (is_leaf is not None and is_leaf(tree)):
        return True
    if optree.register_pytree_node.get(type(tree), namespace=namespace) is None:
        return True
    return False


@substitute_in_graph(optree.tree_iter, can_constant_fold_through=False)  # type: ignore[arg-type]
def tree_iter(
    tree: PyTree,
    /,
    is_leaf: Callable[[PyTree], bool] | None = None,
    *,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> Iterable[Any]:
    stack = [tree]
    while stack:
        node = stack.pop()
        if tree_is_leaf(
            node,
            is_leaf=is_leaf,
            none_is_leaf=none_is_leaf,
            namespace=namespace,
        ):
            yield node
            continue

        children, *_ = optree.tree_flatten_one_level(
            node,
            is_leaf=is_leaf,
            none_is_leaf=none_is_leaf,
            namespace=namespace,
        )
        stack.extend(reversed(children))


@substitute_in_graph(optree.tree_leaves, can_constant_fold_through=True)  # type: ignore[arg-type]
def tree_leaves(
    tree: PyTree,
    /,
    is_leaf: Callable[[PyTree], bool] | None = None,
    *,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> list[Any]:
    return list(
        tree_iter(
            tree,
            is_leaf=is_leaf,
            none_is_leaf=none_is_leaf,
            namespace=namespace,
        )
    )


class _Asterisk(str):
    __slots__ = ()

    def __new__(cls) -> Self:
        return super().__new__(cls, "*")

    def __repr__(self) -> str:
        return "*"  # no quotes


_asterisk = _Asterisk()
del _Asterisk


@dataclass(frozen=True, slots=True)
class PyTreeSpec:
    """Analog for :class:`optree.PyTreeSpec` in Python."""

    _children: tuple[PyTreeSpec, ...]
    _type: builtins.type | None
    _metadata: Any
    _entries: tuple[Any, ...]
    _unflatten_func: Callable[[Any | None, Iterable[PyTree]], PyTree] | None
    none_is_leaf: bool
    namespace: str

    num_nodes: int = field(init=False)
    num_leaves: int = field(init=False)
    num_children: int = field(init=False)

    def __post_init__(self, /) -> None:
        if self._type is None:
            assert len(self._children) == 0
            assert self._metadata is None
            assert self._entries == ()
            assert self._unflatten_func is None
            num_nodes = 1
            num_leaves = 1
            num_children = 0
        else:
            assert callable(self._unflatten_func)
            num_nodes = 1
            num_leaves = 0
            for child in self._children:
                num_nodes += child.num_nodes
                num_leaves += child.num_leaves
            num_children = len(self._children)

        object.__setattr__(self, "num_nodes", num_nodes)
        object.__setattr__(self, "num_leaves", num_leaves)
        object.__setattr__(self, "num_children", num_children)

    def __repr__(self, /) -> str:
        def helper(treespec: PyTreeSpec) -> str:
            if treespec.is_leaf():
                assert treespec.type is None
                return _asterisk

            assert treespec.type is not None
            assert callable(treespec._unflatten_func)
            children_representations = [
                helper(subspec) for subspec in treespec._children
            ]
            if (
                treespec.type in BUILTIN_TYPES
                or (treespec.type is type(None) and not self.none_is_leaf)
                or optree.is_namedtuple_class(treespec.type)
                or optree.is_structseq_class(treespec.type)
            ):
                # pyrefly: ignore [bad-return]
                return treespec._unflatten_func(
                    treespec._metadata,
                    children_representations,
                )
            return (
                f"CustomTreeNode({treespec.type.__name__}[{treespec._metadata!r}], "
                f"[{', '.join(children_representations)}])"
            )

        inner = [
            str(helper(self)),
            *(["NoneIsLeaf"] if self.none_is_leaf else []),
            f"namespace={self.namespace!r}",
        ]
        return f"PyTreeSpec({', '.join(inner)})"

    def __len__(self, /) -> int:
        return self.num_leaves

    @property
    def type(self, /) -> builtins.type | None:
        return self._type

    def is_leaf(self, /) -> bool:
        return self.num_nodes == 1 and self.num_leaves == 1

    def paths(self, /) -> list[tuple[Any, ...]]:
        def helper(treespec: PyTreeSpec, path_prefix: list[Any]) -> None:
            if treespec.is_leaf():
                paths.append(path_prefix)
                return

            for entry, subspec in zip(
                treespec._entries,
                treespec._children,
                strict=True,
            ):
                helper(subspec, path_prefix + [entry])

        paths: list[list[Any]] = []
        helper(self, [])
        return [tuple(path) for path in paths]

    def accessors(self, /) -> list[optree.PyTreeAccessor]:
        def helper(
            treespec: PyTreeSpec,
            entry_path_prefix: list[optree.PyTreeEntry],
        ) -> None:
            if treespec.is_leaf():
                entry_paths.append(entry_path_prefix)
                return

            node_type = treespec.type
            assert node_type is not None
            handler = optree.register_pytree_node.get(
                node_type, namespace=treespec.namespace
            )
            assert handler is not None
            kind: optree.PyTreeKind = handler.kind
            path_entry_type: type[optree.PyTreeEntry] = handler.path_entry_type

            for entry, subspec in zip(
                treespec._entries,
                treespec._children,
                strict=True,
            ):
                helper(
                    subspec,
                    entry_path_prefix + [path_entry_type(entry, node_type, kind)],
                )

        entry_paths: list[list[optree.PyTreeEntry]] = []
        helper(self, [])
        return [optree.PyTreeAccessor(path) for path in entry_paths]

    def children(self, /) -> list[PyTreeSpec]:
        return list(self._children)

    def child(self, index: int, /) -> PyTreeSpec:
        return self._children[index]

    def entries(self, /) -> list[Any]:
        return list(self._entries)

    def entry(self, index: int, /) -> Any:
        return self._entries[index]

    def flatten_up_to(self, tree: PyTree, /) -> list[PyTree]:
        def helper(
            treespec: PyTreeSpec,
            node: PyTree,
            subtrees: list[PyTree],
        ) -> None:
            if treespec.is_leaf():
                subtrees.append(node)
                return

            node_type = type(node)
            if treespec.type not in BUILTIN_TYPES:
                # Always require custom node types to match exactly
                if node_type != treespec.type:
                    raise ValueError(
                        f"Type mismatch; "
                        f"expected {treespec.type!r}, but got {node_type!r}.",
                    )

                children, metadata, *_ = optree.tree_flatten_one_level(
                    node,
                    none_is_leaf=self.none_is_leaf,
                    namespace=self.namespace,
                )
                if len(children) != treespec.num_children:
                    raise ValueError(
                        f"Node arity mismatch; "
                        f"expected {treespec.num_children}, but got {len(children)}.",
                    )
                if metadata != treespec._metadata:
                    raise ValueError(
                        f"Node context mismatch for custom node type {treespec.type!r}.",
                    )
            else:
                # For builtin dictionary types, we allow some flexibility
                # Otherwise, we require exact matches
                both_standard_dict = (
                    treespec.type in STANDARD_DICT_TYPES
                    and node_type in STANDARD_DICT_TYPES
                )
                if not both_standard_dict and node_type != treespec.type:
                    raise ValueError(
                        f"Node type mismatch; "
                        f"expected {treespec.type!r}, but got {node_type!r}.",
                    )
                if len(node) != treespec.num_children:
                    raise ValueError(
                        f"Node arity mismatch; "
                        f"expected {treespec.num_children}, but got {len(node)}.",
                    )

                if both_standard_dict:
                    # dictionary types are compatible with each other
                    expected_keys = treespec.entries()
                    got_key_set = set(node)
                    expected_key_set = set(expected_keys)
                    if got_key_set != expected_key_set:
                        missing_keys = expected_key_set.difference(got_key_set)
                        extra_keys = got_key_set.difference(expected_key_set)
                        message = ""
                        if missing_keys:
                            message += f"; missing key(s): {missing_keys}"
                        if extra_keys:
                            message += f"; extra key(s): {extra_keys}"
                        raise ValueError(f"Node keys mismatch{message}.")
                    children = [node[key] for key in expected_keys]
                else:
                    # node_type is treespec.type
                    children, metadata, *_ = optree.tree_flatten_one_level(
                        node,
                        none_is_leaf=self.none_is_leaf,
                        namespace=self.namespace,
                    )
                    if (
                        node_type is not deque  # ignore mismatch of `maxlen` for deque
                    ) and metadata != treespec._metadata:
                        raise ValueError(
                            f"Node metadata mismatch for node type {treespec.type!r}; "
                            f"expected {treespec._metadata!r}, but got {metadata!r}.",  # namedtuple type mismatch
                        )

            for subtree, subspec in zip(children, treespec._children, strict=True):
                helper(subspec, subtree, subtrees)

        subtrees: list[PyTree] = []
        helper(self, tree, subtrees)
        return subtrees

    def unflatten(self, leaves: Iterable[Any], /) -> PyTree:
        if not isinstance(leaves, (list, tuple)):
            leaves = list(leaves)
        if len(leaves) != self.num_leaves:
            raise ValueError(
                f"treespec.unflatten(leaves): `leaves` has length {len(leaves)} "
                f"but the spec refers to a pytree that holds {self.num_leaves} "
                f"items ({self}).",
            )
        if self.is_leaf():
            return leaves[0]

        # Recursively unflatten the children
        start = 0
        end = 0
        subtrees = []
        for subspec in self._children:
            end += subspec.num_leaves
            subtrees.append(subspec.unflatten(leaves[start:end]))
            start = end

        assert callable(self._unflatten_func)
        return self._unflatten_func(self._metadata, subtrees)


def _is_pytreespec_instance(obj: Any, /) -> TypeIs[PyTreeSpec | python_pytree.TreeSpec]:
    return isinstance(obj, (PyTreeSpec, python_pytree.TreeSpec))


@substitute_in_graph(  # type: ignore[arg-type]
    optree.treespec_leaf,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def treespec_leaf(
    *,
    none_is_leaf: bool = False,
    namespace: str = "",  # unused
) -> PyTreeSpec:
    return PyTreeSpec(
        (),
        None,
        None,
        (),
        None,
        none_is_leaf=none_is_leaf,
        namespace="",
    )


@substitute_in_graph(  # type: ignore[arg-type]
    optree.treespec_tuple,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def treespec_tuple(
    iterable: Iterable[PyTreeSpec] = (),
    /,
    *,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> PyTreeSpec:
    children = tuple(iterable)
    if any(not _is_pytreespec_instance(child) for child in children):
        raise ValueError(f"Expected a tuple of PyTreeSpecs, got: {children!r}.")
    if any(child.none_is_leaf != none_is_leaf for child in children):
        raise ValueError(
            "All children PyTreeSpecs must have the same `none_is_leaf` value "
            f"as the parent; expected {none_is_leaf}, got: {children!r}.",
        )
    if any(child.namespace not in (namespace, "") for child in children):
        raise ValueError(
            "All children PyTreeSpecs must have the same `namespace` value "
            f"as the parent; expected {namespace!r}, got: {children!r}.",
        )
    handler = optree.register_pytree_node.get(tuple, namespace=namespace)
    assert handler is not None
    return PyTreeSpec(
        tuple(children),
        tuple,
        None,
        tuple(range(len(children))),
        handler.unflatten_func,
        none_is_leaf=none_is_leaf,
        namespace=namespace,
    )


@substitute_in_graph(  # type: ignore[arg-type]
    optree.treespec_dict,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def treespec_dict(
    mapping: Mapping[Any, PyTreeSpec] | Iterable[tuple[Any, PyTreeSpec]] = (),
    /,
    *,
    none_is_leaf: bool = False,
    namespace: str = "",
    **kwargs: PyTreeSpec,
) -> PyTreeSpec:
    dct = dict(mapping, **kwargs)
    if any(not _is_pytreespec_instance(child) for child in dct.values()):
        raise ValueError(f"Expected a dictionary of TreeSpecs, got: {dct!r}.")
    if any(child.none_is_leaf != none_is_leaf for child in dct.values()):
        raise ValueError(
            "All children PyTreeSpecs must have the same `none_is_leaf` value "
            f"as the parent; expected {none_is_leaf}, got: {dct!r}.",
        )
    if any(child.namespace not in (namespace, "") for child in dct.values()):
        raise ValueError(
            "All children PyTreeSpecs must have the same `namespace` value "
            f"as the parent; expected {namespace!r}, got: {dct!r}.",
        )

    (
        children,
        metadata,
        entries,
        unflatten_func,
    ) = optree.tree_flatten_one_level(  # type: ignore[assignment,var-annotated]
        dct,  # type: ignore[arg-type]
        none_is_leaf=none_is_leaf,
        namespace=namespace,
    )
    return PyTreeSpec(
        tuple(children),  # type: ignore[arg-type]
        dict,
        metadata,
        entries,
        unflatten_func,  # type: ignore[arg-type]
        none_is_leaf=none_is_leaf,
        namespace=namespace,
    )


@substitute_in_graph(  # type: ignore[arg-type]
    optree.tree_flatten,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def tree_flatten(
    tree: PyTree,
    /,
    is_leaf: Callable[[PyTree], bool] | None = None,
    *,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> tuple[list[Any], PyTreeSpec]:
    def helper(node: PyTree, leaves: list[Any]) -> PyTreeSpec:
        if tree_is_leaf(
            node,
            is_leaf=is_leaf,
            none_is_leaf=none_is_leaf,
            namespace=namespace,
        ):
            leaves.append(node)
            return PyTreeSpec(
                (),
                None,
                None,
                (),
                None,
                none_is_leaf=none_is_leaf,
                namespace=namespace,
            )

        (
            children,
            metadata,
            entries,
            unflatten_func,
        ) = optree.tree_flatten_one_level(
            node,
            is_leaf=is_leaf,
            none_is_leaf=none_is_leaf,
            namespace=namespace,
        )

        # Recursively flatten the children
        subspecs = tuple(helper(child, leaves) for child in children)
        return PyTreeSpec(
            subspecs,
            type(node),
            metadata,
            entries,
            unflatten_func,  # type: ignore[arg-type]
            none_is_leaf=none_is_leaf,
            namespace=namespace,
        )  # type: ignore[arg-type]

    leaves: list[Any] = []
    treespec = helper(tree, leaves)
    return leaves, treespec


@substitute_in_graph(  # type: ignore[arg-type]
    optree._C.flatten,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def _C_flatten(
    tree: PyTree,
    /,
    leaf_predicate: Callable[[PyTree], bool] | None = None,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> tuple[list[Any], PyTreeSpec]:
    return tree_flatten(  # type: ignore[return-value]
        tree,
        is_leaf=leaf_predicate,
        none_is_leaf=none_is_leaf,
        namespace=namespace,
    )


@substitute_in_graph(  # type: ignore[arg-type]
    optree.tree_flatten_with_path,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def tree_flatten_with_path(
    tree: PyTree,
    /,
    is_leaf: Callable[[PyTree], bool] | None = None,
    *,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> tuple[list[tuple[Any, ...]], list[Any], PyTreeSpec]:
    leaves, treespec = tree_flatten(
        tree,
        is_leaf=is_leaf,
        none_is_leaf=none_is_leaf,
        namespace=namespace,
    )
    return treespec.paths(), leaves, treespec  # type: ignore[return-value]


@substitute_in_graph(  # type: ignore[arg-type]
    optree._C.flatten_with_path,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def _C_flatten_with_path(
    tree: PyTree,
    /,
    leaf_predicate: Callable[[PyTree], bool] | None = None,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> tuple[list[tuple[Any, ...]], list[Any], PyTreeSpec]:
    return tree_flatten_with_path(  # type: ignore[return-value]
        tree,
        is_leaf=leaf_predicate,
        none_is_leaf=none_is_leaf,
        namespace=namespace,
    )


@substitute_in_graph(  # type: ignore[arg-type]
    optree.tree_structure,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def tree_structure(
    tree: PyTree,
    /,
    is_leaf: Callable[[PyTree], bool] | None = None,
    *,
    none_is_leaf: bool = False,
    namespace: str = "",
) -> PyTreeSpec:
    return tree_flatten(  # type: ignore[return-value]
        tree,
        is_leaf=is_leaf,
        none_is_leaf=none_is_leaf,
        namespace=namespace,
    )[1]


@substitute_in_graph(  # type: ignore[arg-type]
    optree.tree_unflatten,
    # We need to disable constant folding here because we want the function to reference the
    # PyTreeSpec class defined above, not the one in the C++ module.
    can_constant_fold_through=False,
)
def tree_unflatten(treespec: PyTreeSpec, leaves: Iterable[Any]) -> PyTree:
    if not _is_pytreespec_instance(treespec):
        raise TypeError(
            f"Expected `treespec` to be an instance of "
            f"PyTreeSpec but got item of type {type(treespec)}."
        )
    return treespec.unflatten(leaves)


_none_registration = optree.register_pytree_node.get(type(None))
assert _none_registration is not None


@substitute_in_graph(  # type: ignore[arg-type]
    _none_registration.unflatten_func,
    can_constant_fold_through=True,
    skip_signature_check=True,
)
def none_unflatten(_: None, children: Iterable[_T], /) -> None:
    if len(list(children)) != 0:
        raise ValueError("Expected no children.")
    return None


with optree.dict_insertion_ordered(False, namespace="torch"):
    _dict_registration = optree.register_pytree_node.get(dict)
    assert _dict_registration is not None


@substitute_in_graph(  # type: ignore[arg-type]
    _dict_registration.flatten_func,
    can_constant_fold_through=True,
    skip_signature_check=True,
)
def dict_flatten(
    dct: dict[_KT, _VT], /
) -> tuple[list[_VT], tuple[list[_KT], list[_KT]], tuple[_KT, ...]]:
    sorted_keys = optree.utils.total_order_sorted(dct)
    values = [dct[key] for key in sorted_keys]
    original_keys = list(dct)
    return values, (original_keys, sorted_keys), tuple(sorted_keys)


@substitute_in_graph(  # type: ignore[arg-type]
    _dict_registration.unflatten_func,
    can_constant_fold_through=True,
    skip_signature_check=True,
)
def dict_unflatten(
    metadata: tuple[list[_KT], list[_KT]],
    values: Iterable[_VT],
    /,
) -> dict[_KT, _VT]:
    original_keys, sorted_keys = metadata
    d = dict.fromkeys(original_keys)
    d.update(zip(sorted_keys, values, strict=True))
    return d  # type: ignore[return-value]