{"id":13688834,"url":"https://github.com/metaopt/optree","last_synced_at":"2025-05-01T20:30:56.019Z","repository":{"id":61791591,"uuid":"540486972","full_name":"metaopt/optree","owner":"metaopt","description":"OpTree: Optimized PyTree Utilities","archived":false,"fork":false,"pushed_at":"2025-05-01T11:55:07.000Z","size":1194,"stargazers_count":178,"open_issues_count":5,"forks_count":9,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-05-01T12:38:59.236Z","etag":null,"topics":["optree","pytree","tree","tree-flatten","tree-map","tree-operation","tree-structure"],"latest_commit_sha":null,"homepage":"https://optree.readthedocs.io","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/metaopt.png","metadata":{"files":{"readme":"README.md","changelog":"CHANGELOG.md","contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","threat_model":null,"audit":null,"citation":null,"codeowners":".github/CODEOWNERS","security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2022-09-23T14:46:39.000Z","updated_at":"2025-05-01T11:55:10.000Z","dependencies_parsed_at":"2023-02-19T19:40:18.089Z","dependency_job_id":"2456ac04-43e3-441a-a238-51af344d81c2","html_url":"https://github.com/metaopt/optree","commit_stats":{"total_commits":157,"total_committers":6,"mean_commits":"26.166666666666668","dds":0.07006369426751591,"last_synced_commit":"60707129e9cf37b015f9fee300a2f3611ff852cc"},"previous_names":[],"tags_count":23,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/metaopt%2Foptree","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/metaopt%2Foptree/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/metaopt%2Foptree/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/metaopt%2Foptree/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/metaopt","download_url":"https://codeload.github.com/metaopt/optree/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":251878841,"owners_count":21658680,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["optree","pytree","tree","tree-flatten","tree-map","tree-operation","tree-structure"],"created_at":"2024-08-02T15:01:24.266Z","updated_at":"2025-05-01T20:30:56.011Z","avatar_url":"https://github.com/metaopt.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"\u003c!-- markdownlint-disable html --\u003e\n\n# OpTree\n\n![Python 3.9+](https://img.shields.io/badge/Python-3.9%2B-brightgreen)\n[![PyPI](https://img.shields.io/pypi/v/optree?logo=pypi)](https://pypi.org/project/optree)\n![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/metaopt/optree/build.yml?label=build\u0026logo=github)\n![GitHub Workflow Status](https://img.shields.io/github/actions/workflow/status/metaopt/optree/tests.yml?label=tests\u0026logo=github)\n[![Codecov](https://img.shields.io/codecov/c/github/metaopt/optree/main?logo=codecov)](https://codecov.io/gh/metaopt/optree)\n[![Documentation Status](https://img.shields.io/readthedocs/optree?logo=readthedocs)](https://optree.readthedocs.io)\n[![Downloads](https://static.pepy.tech/personalized-badge/optree?period=total\u0026left_color=grey\u0026right_color=blue\u0026left_text=downloads)](https://pepy.tech/project/optree)\n[![GitHub Repo Stars](https://img.shields.io/github/stars/metaopt/optree?color=brightgreen\u0026logo=github)](https://github.com/metaopt/optree/stargazers)\n\nOptimized PyTree Utilities.\n\n--------------------------------------------------------------------------------\n\n### Table of Contents  \u003c!-- omit in toc --\u003e \u003c!-- markdownlint-disable heading-increment --\u003e\n\n- [Installation](#installation)\n- [PyTrees](#pytrees)\n  - [Tree Nodes and Leaves](#tree-nodes-and-leaves)\n    - [Built-in PyTree Node Types](#built-in-pytree-node-types)\n    - [Registering a Container-like Custom Type as Non-leaf Nodes](#registering-a-container-like-custom-type-as-non-leaf-nodes)\n    - [Notes about the PyTree Type Registry](#notes-about-the-pytree-type-registry)\n  - [`None` is Non-leaf Node vs. `None` is Leaf](#none-is-non-leaf-node-vs-none-is-leaf)\n  - [Key Ordering for Dictionaries](#key-ordering-for-dictionaries)\n- [Changelog](#changelog)\n- [License](#license)\n\n--------------------------------------------------------------------------------\n\n## Installation\n\nInstall from PyPI ([![PyPI](https://img.shields.io/pypi/v/optree?logo=pypi)](https://pypi.org/project/optree) / ![Status](https://img.shields.io/pypi/status/optree)):\n\n```bash\npip3 install --upgrade optree\n```\n\nInstall from conda-forge ([![conda-forge](https://img.shields.io/conda/v/conda-forge/optree?logo=condaforge)](https://anaconda.org/conda-forge/optree)):\n\n```bash\nconda install conda-forge::optree\n```\n\nInstall the latest version from GitHub:\n\n```bash\npip3 install git+https://github.com/metaopt/optree.git#egg=optree\n```\n\nOr, clone this repo and install manually:\n\n```bash\ngit clone --depth=1 https://github.com/metaopt/optree.git\ncd optree\n\npip3 install .\n```\n\nThe following options are available while building the Python C extension from the source:\n\n```bash\nexport CMAKE_COMMAND=\"/path/to/custom/cmake\"\nexport CMAKE_BUILD_TYPE=\"Debug\"\nexport CMAKE_CXX_STANDARD=\"20\"  # C++17 is tested on Linux/macOS (C++20 is required on Windows)\nexport OPTREE_CXX_WERROR=\"OFF\"\nexport _GLIBCXX_USE_CXX11_ABI=\"1\"\nexport pybind11_DIR=\"/path/to/custom/pybind11\"\npip3 install .\n```\n\nCompiling from the source requires Python 3.9+, a compiler (`gcc` / `clang` / `icc` / `cl.exe`) that supports C++20 and a `cmake` installation.\n\n--------------------------------------------------------------------------------\n\n## PyTrees\n\nA PyTree is a recursive structure that can be an arbitrarily nested Python container (e.g., `tuple`, `list`, `dict`, `OrderedDict`, `NamedTuple`, etc.) or an opaque Python object.\nThe key concepts of tree operations are tree flattening and its inverse (tree unflattening).\nAdditional tree operations can be performed based on these two basic functions (e.g., `tree_map = tree_unflatten ∘ map ∘ tree_flatten`).\n\nTree flattening is traversing the entire tree in a left-to-right depth-first manner and returning the leaves of the tree in a deterministic order.\n\n```python\n\u003e\u003e\u003e tree = {'b': (2, [3, 4]), 'a': 1, 'c': 5, 'd': 6}\n\u003e\u003e\u003e optree.tree_flatten(tree)\n([1, 2, 3, 4, 5, 6], PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': *, 'd': *}))\n\u003e\u003e\u003e optree.tree_flatten(1)\n([1], PyTreeSpec(*))\n\u003e\u003e\u003e optree.tree_flatten(None)\n([], PyTreeSpec(None))\n\u003e\u003e\u003e optree.tree_map(lambda x: x**2, tree)\n{'b': (4, [9, 16]), 'a': 1, 'c': 25, 'd': 36}\n```\n\nThis usually implies that the equal pytrees return equal lists of leaves and the same tree structure.\nSee also section [Key Ordering for Dictionaries](#key-ordering-for-dictionaries).\n\n```python\n\u003e\u003e\u003e {'a': [1, 2], 'b': [3]} == {'b': [3], 'a': [1, 2]}\nTrue\n\u003e\u003e\u003e optree.tree_leaves({'a': [1, 2], 'b': [3]}) == optree.tree_leaves({'b': [3], 'a': [1, 2]})\nTrue\n\u003e\u003e\u003e optree.tree_structure({'a': [1, 2], 'b': [3]}) == optree.tree_structure({'b': [3], 'a': [1, 2]})\nTrue\n\u003e\u003e\u003e optree.tree_map(lambda x: x**2, {'a': [1, 2], 'b': [3]})\n{'a': [1, 4], 'b': [9]}\n\u003e\u003e\u003e optree.tree_map(lambda x: x**2, {'b': [3], 'a': [1, 2]})\n{'b': [9], 'a': [1, 4]}\n```\n\n\u003e [!TIP]\n\u003e\n\u003e Since OpTree v0.14.1, a new namespace `optree.pytree` is introduced as aliases for `optree.tree_*` functions. The following examples are equivalent to the above:\n\u003e\n\u003e ```python\n\u003e \u003e\u003e\u003e import optree.pytree as pt\n\u003e \u003e\u003e\u003e tree = {'b': (2, [3, 4]), 'a': 1, 'c': 5, 'd': 6}\n\u003e \u003e\u003e\u003e pt.flatten(tree)\n\u003e ([1, 2, 3, 4, 5, 6], PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': *, 'd': *}))\n\u003e \u003e\u003e\u003e pt.flatten(1)\n\u003e ([1], PyTreeSpec(*))\n\u003e \u003e\u003e\u003e pt.flatten(None)\n\u003e ([], PyTreeSpec(None))\n\u003e \u003e\u003e\u003e pt.map(lambda x: x**2, tree)\n\u003e {'b': (4, [9, 16]), 'a': 1, 'c': 25, 'd': 36}\n\u003e \u003e\u003e\u003e pt.map(lambda x: x**2, {'a': [1, 2], 'b': [3]})\n\u003e {'a': [1, 4], 'b': [9]}\n\u003e \u003e\u003e\u003e pt.map(lambda x: x**2, {'b': [3], 'a': [1, 2]})\n\u003e {'b': [9], 'a': [1, 4]}\n\u003e ```\n\n### Tree Nodes and Leaves\n\nA tree is a collection of non-leaf nodes and leaf nodes, where the leaf nodes have no children to flatten.\n`optree.tree_flatten(...)` will flatten the tree and return a list of leaf nodes while the non-leaf nodes will store in the tree specification.\n\n#### Built-in PyTree Node Types\n\nOpTree out-of-box supports the following Python container types in the registry:\n\n- [`tuple`](https://docs.python.org/3/library/stdtypes.html#tuple)\n- [`list`](https://docs.python.org/3/library/stdtypes.html#list)\n- [`dict`](https://docs.python.org/3/library/stdtypes.html#dict)\n- [`collections.namedtuple`](https://docs.python.org/3/library/collections.html#collections.namedtuple) and its subclasses\n- [`collections.OrderedDict`](https://docs.python.org/3/library/collections.html#collections.OrderedDict)\n- [`collections.defaultdict`](https://docs.python.org/3/library/collections.html#collections.defaultdict)\n- [`collections.deque`](https://docs.python.org/3/library/collections.html#collections.deque)\n- [`PyStructSequence`](https://docs.python.org/3/c-api/tuple.html#struct-sequence-objects) types created by C API [`PyStructSequence_NewType`](https://docs.python.org/3/c-api/tuple.html#c.PyStructSequence_NewType)\n\nwhich are considered non-leaf nodes in the tree.\nPython objects that the type is not registered will be treated as leaf nodes.\nThe registry lookup uses the `is` operator to determine whether the type is matched.\nSo subclasses will need to explicitly register in the registry, otherwise, an object of that type will be considered a leaf.\nThe [`NoneType`](https://docs.python.org/3/library/constants.html#None) is a special case discussed in section [`None` is non-leaf Node vs. `None` is Leaf](#none-is-non-leaf-node-vs-none-is-leaf).\n\n#### Registering a Container-like Custom Type as Non-leaf Nodes\n\nA container-like Python type can be registered in the type registry with a pair of functions that specify:\n\n- `flatten_func(container) -\u003e (children, metadata, entries)`: convert an instance of the container type to a `(children, metadata, entries)` triple, where `children` is an iterable of subtrees and `entries` is an iterable of path entries of the container (e.g., indices or keys).\n- `unflatten_func(metadata, children) -\u003e container`: convert such a pair back to an instance of the container type.\n\nThe `metadata` is some necessary data apart from the children to reconstruct the container, e.g., the keys of the dictionary (the children are values).\n\nThe `entries` can be omitted (only returns a pair) or is optional to implement (returns `None`). If so, use `range(len(children))` (i.e., flat indices) as path entries of the current node. The signature for the flatten function can be one of the following:\n\n- `flatten_func(container) -\u003e (children, metadata, entries)`\n- `flatten_func(container) -\u003e (children, metadata, None)`\n- `flatten_func(container) -\u003e (children, metadata)`\n\nThe following examples show how to register custom types and utilize them for `tree_flatten` and `tree_map`. Please refer to section [Notes about the PyTree Type Registry](#notes-about-the-pytree-type-registry) for more information.\n\n```python\n# Registry a Python type with lambda functions\noptree.register_pytree_node(\n    set,\n    # (set) -\u003e (children, metadata, None)\n    lambda s: (sorted(s), None, None),\n    # (metadata, children) -\u003e (set)\n    lambda _, children: set(children),\n    namespace='set',\n)\n\n# Register a Python type into a namespace\nimport torch\n\nclass Torch2NumpyEntry(optree.PyTreeEntry):\n    def __call__(self, obj):\n        assert self.entry == 0\n        return obj.cpu().detach().numpy()\n\n    def codify(self, node=''):\n        assert self.entry == 0\n        return f'{node}.cpu().detach().numpy()'\n\noptree.register_pytree_node(\n    torch.Tensor,\n    # (tensor) -\u003e (children, metadata)\n    flatten_func=lambda tensor: (\n        (tensor.cpu().detach().numpy(),),\n        {'dtype': tensor.dtype, 'device': tensor.device, 'requires_grad': tensor.requires_grad},\n    ),\n    # (metadata, children) -\u003e tensor\n    unflatten_func=lambda metadata, children: torch.tensor(children[0], **metadata),\n    path_entry_type=Torch2NumpyEntry,\n    namespace='torch2numpy',\n)\n```\n\n```python\n\u003e\u003e\u003e tree = {'weight': torch.ones(size=(1, 2)).cuda(), 'bias': torch.zeros(size=(2,))}\n\u003e\u003e\u003e tree\n{'weight': tensor([[1., 1.]], device='cuda:0'), 'bias': tensor([0., 0.])}\n\n# Flatten without specifying the namespace\n\u003e\u003e\u003e optree.tree_flatten(tree)  # `torch.Tensor`s are leaf nodes\n([tensor([0., 0.]), tensor([[1., 1.]], device='cuda:0')], PyTreeSpec({'bias': *, 'weight': *}))\n\n# Flatten with the namespace\n\u003e\u003e\u003e leaves, treespec = optree.tree_flatten(tree, namespace='torch2numpy')\n\u003e\u003e\u003e leaves, treespec\n(\n    [array([0., 0.], dtype=float32), array([[1., 1.]], dtype=float32)],\n    PyTreeSpec(\n        {\n            'bias': CustomTreeNode(Tensor[{'dtype': torch.float32, 'device': device(type='cpu'), 'requires_grad': False}], [*]),\n            'weight': CustomTreeNode(Tensor[{'dtype': torch.float32, 'device': device(type='cuda', index=0), 'requires_grad': False}], [*])\n        },\n        namespace='torch2numpy'\n    )\n)\n\n# `entries` are not defined and use `range(len(children))`\n\u003e\u003e\u003e optree.tree_paths(tree, namespace='torch2numpy')\n[('bias', 0), ('weight', 0)]\n\n# Custom path entry type defines the pytree access behavior\n\u003e\u003e\u003e optree.tree_accessors(tree, namespace='torch2numpy')\n[\n    PyTreeAccessor(*['bias'].cpu().detach().numpy(), (MappingEntry(key='bias', type=\u003cclass 'dict'\u003e), Torch2NumpyEntry(entry=0, type=\u003cclass 'torch.Tensor'\u003e))),\n    PyTreeAccessor(*['weight'].cpu().detach().numpy(), (MappingEntry(key='weight', type=\u003cclass 'dict'\u003e), Torch2NumpyEntry(entry=0, type=\u003cclass 'torch.Tensor'\u003e)))\n]\n\n# Unflatten back to a copy of the original object\n\u003e\u003e\u003e optree.tree_unflatten(treespec, leaves)\n{'weight': tensor([[1., 1.]], device='cuda:0'), 'bias': tensor([0., 0.])}\n```\n\nUsers can also extend the pytree registry by decorating the custom class and defining an instance method `tree_flatten` and a class method `tree_unflatten`.\n\n```python\nfrom collections import UserDict\n\n@optree.register_pytree_node_class(namespace='mydict')\nclass MyDict(UserDict):\n    TREE_PATH_ENTRY_TYPE = optree.MappingEntry  # used by accessor APIs\n\n    def tree_flatten(self):  # -\u003e (children, metadata, entries)\n        reversed_keys = sorted(self.keys(), reverse=True)\n        return (\n            [self[key] for key in reversed_keys],  # children\n            reversed_keys,  # metadata\n            reversed_keys,  # entries\n        )\n\n    @classmethod\n    def tree_unflatten(cls, metadata, children):\n        return cls(zip(metadata, children))\n```\n\n```python\n\u003e\u003e\u003e tree = MyDict(b=4, a=(2, 3), c=MyDict({'d': 5, 'f': 6}))\n\n# Flatten without specifying the namespace\n\u003e\u003e\u003e optree.tree_flatten_with_path(tree)  # `MyDict`s are leaf nodes\n(\n    [()],\n    [MyDict(b=4, a=(2, 3), c=MyDict({'d': 5, 'f': 6}))],\n    PyTreeSpec(*)\n)\n\n# Flatten with the namespace\n\u003e\u003e\u003e optree.tree_flatten_with_path(tree, namespace='mydict')\n(\n    [('c', 'f'), ('c', 'd'), ('b',), ('a', 0), ('a', 1)],\n    [6, 5, 4, 2, 3],\n    PyTreeSpec(\n        CustomTreeNode(MyDict[['c', 'b', 'a']], [CustomTreeNode(MyDict[['f', 'd']], [*, *]), *, (*, *)]),\n        namespace='mydict'\n    )\n)\n\u003e\u003e\u003e optree.tree_flatten_with_accessor(tree, namespace='mydict')\n(\n    [\n        PyTreeAccessor(*['c']['f'], (MappingEntry(key='c', type=\u003cclass 'MyDict'\u003e), MappingEntry(key='f', type=\u003cclass 'MyDict'\u003e))),\n        PyTreeAccessor(*['c']['d'], (MappingEntry(key='c', type=\u003cclass 'MyDict'\u003e), MappingEntry(key='d', type=\u003cclass 'MyDict'\u003e))),\n        PyTreeAccessor(*['b'], (MappingEntry(key='b', type=\u003cclass 'MyDict'\u003e),)),\n        PyTreeAccessor(*['a'][0], (MappingEntry(key='a', type=\u003cclass 'MyDict'\u003e), SequenceEntry(index=0, type=\u003cclass 'tuple'\u003e))),\n        PyTreeAccessor(*['a'][1], (MappingEntry(key='a', type=\u003cclass 'MyDict'\u003e), SequenceEntry(index=1, type=\u003cclass 'tuple'\u003e)))\n    ],\n    [6, 5, 4, 2, 3],\n    PyTreeSpec(\n        CustomTreeNode(MyDict[['c', 'b', 'a']], [CustomTreeNode(MyDict[['f', 'd']], [*, *]), *, (*, *)]),\n        namespace='mydict'\n    )\n)\n```\n\n#### Notes about the PyTree Type Registry\n\nThere are several key attributes of the pytree type registry:\n\n1. **The type registry is per-interpreter-dependent.** This means registering a custom type in the registry affects all modules that use OpTree.\n\n\u003e [!WARNING]\n\u003e For safety reasons, a `namespace` must be specified while registering a custom type. It is\n\u003e used to isolate the behavior of flattening and unflattening a pytree node type. This is to\n\u003e prevent accidental collisions between different libraries that may register the same type.\n\n2. **The elements in the type registry are immutable.** Users can neither register the same type twice in the same namespace (i.e., update the type registry), nor remove a type from the type registry. To update the behavior of an already registered type, simply register it again with another `namespace`.\n\n3. **Users cannot modify the behavior of already registered built-in types** listed in [Built-in PyTree Node Types](#built-in-pytree-node-types), such as key order sorting for `dict` and `collections.defaultdict`.\n\n4. **Inherited subclasses are not implicitly registered.** The registry lookup uses `type(obj) is registered_type` rather than `isinstance(obj, registered_type)`. Users need to register the subclasses explicitly. To register all subclasses, it is easy to implement with [`metaclass`](https://docs.python.org/3/reference/datamodel.html#metaclasses) or [`__init_subclass__`](https://docs.python.org/3/reference/datamodel.html#customizing-class-creation), for example:\n\n    ```python\n    from collections import UserDict\n\n    @optree.register_pytree_node_class(namespace='mydict')\n    class MyDict(UserDict):\n        TREE_PATH_ENTRY_TYPE = optree.MappingEntry  # used by accessor APIs\n\n        def __init_subclass__(cls):  # define this in the base class\n            super().__init_subclass__()\n            # Register a subclass to namespace 'mydict'\n            optree.register_pytree_node_class(cls, namespace='mydict')\n\n        def tree_flatten(self):  # -\u003e (children, metadata, entries)\n            reversed_keys = sorted(self.keys(), reverse=True)\n            return (\n                [self[key] for key in reversed_keys],  # children\n                reversed_keys,  # metadata\n                reversed_keys,  # entries\n            )\n\n        @classmethod\n        def tree_unflatten(cls, metadata, children):\n            return cls(zip(metadata, children))\n\n    # Subclasses will be automatically registered in namespace 'mydict'\n    class MyAnotherDict(MyDict):\n        pass\n    ```\n\n    ```python\n    \u003e\u003e\u003e tree = MyDict(b=4, a=(2, 3), c=MyAnotherDict({'d': 5, 'f': 6}))\n    \u003e\u003e\u003e optree.tree_flatten_with_path(tree, namespace='mydict')\n    (\n        [('c', 'f'), ('c', 'd'), ('b',), ('a', 0), ('a', 1)],\n        [6, 5, 4, 2, 3],\n        PyTreeSpec(\n            CustomTreeNode(MyDict[['c', 'b', 'a']], [CustomTreeNode(MyAnotherDict[['f', 'd']], [*, *]), *, (*, *)]),\n            namespace='mydict'\n        )\n    )\n    \u003e\u003e\u003e optree.tree_accessors(tree, namespace='mydict')\n    [\n        PyTreeAccessor(*['c']['f'], (MappingEntry(key='c', type=\u003cclass 'MyDict'\u003e), MappingEntry(key='f', type=\u003cclass 'MyAnotherDict'\u003e))),\n        PyTreeAccessor(*['c']['d'], (MappingEntry(key='c', type=\u003cclass 'MyDict'\u003e), MappingEntry(key='d', type=\u003cclass 'MyAnotherDict'\u003e))),\n        PyTreeAccessor(*['b'], (MappingEntry(key='b', type=\u003cclass 'MyDict'\u003e),)),\n        PyTreeAccessor(*['a'][0], (MappingEntry(key='a', type=\u003cclass 'MyDict'\u003e), SequenceEntry(index=0, type=\u003cclass 'tuple'\u003e))),\n        PyTreeAccessor(*['a'][1], (MappingEntry(key='a', type=\u003cclass 'MyDict'\u003e), SequenceEntry(index=1, type=\u003cclass 'tuple'\u003e)))\n    ]\n    ```\n\n5. **Be careful about the potential infinite recursion of the custom flatten function.** The returned `children` from the custom flatten function are considered subtrees. They will be further flattened recursively. The `children` can have the same type as the current node. Users must design their termination condition carefully.\n\n    ```python\n    import numpy as np\n    import torch\n\n    optree.register_pytree_node(\n        np.ndarray,\n        # Children are nest lists of Python objects\n        lambda array: (np.atleast_1d(array).tolist(), array.ndim == 0),\n        lambda scalar, rows: np.asarray(rows) if not scalar else np.asarray(rows[0]),\n        namespace='numpy1',\n    )\n\n    optree.register_pytree_node(\n        np.ndarray,\n        # Children are Python objects\n        lambda array: (\n            list(array.ravel()),  # list(1DArray[T]) -\u003e List[T]\n            dict(shape=array.shape, dtype=array.dtype)\n        ),\n        lambda metadata, children: np.asarray(children, dtype=metadata['dtype']).reshape(metadata['shape']),\n        namespace='numpy2',\n    )\n\n    optree.register_pytree_node(\n        np.ndarray,\n        # Returns a list of `np.ndarray`s without termination condition\n        lambda array: ([array.ravel()], array.dtype),\n        lambda shape, children: children[0].reshape(shape),\n        namespace='numpy3',\n    )\n\n    optree.register_pytree_node(\n        torch.Tensor,\n        # Children are nest lists of Python objects\n        lambda tensor: (torch.atleast_1d(tensor).tolist(), tensor.ndim == 0),\n        lambda scalar, rows: torch.tensor(rows) if not scalar else torch.tensor(rows[0])),\n        namespace='torch1',\n    )\n\n    optree.register_pytree_node(\n        torch.Tensor,\n        # Returns a list of `torch.Tensor`s without termination condition\n        lambda tensor: (\n            list(tensor.view(-1)),  # list(1DTensor[T]) -\u003e List[0DTensor[T]] (STILL TENSORS!)\n            tensor.shape\n        ),\n        lambda shape, children: torch.stack(children).reshape(shape),\n        namespace='torch2',\n    )\n    ```\n\n    ```python\n    \u003e\u003e\u003e optree.tree_flatten(np.arange(9).reshape(3, 3), namespace='numpy1')\n    (\n        [0, 1, 2, 3, 4, 5, 6, 7, 8],\n        PyTreeSpec(\n            CustomTreeNode(ndarray[False], [[*, *, *], [*, *, *], [*, *, *]]),\n            namespace='numpy1'\n        )\n    )\n    # Implicitly casts `float`s to `np.float64`\n    \u003e\u003e\u003e optree.tree_map(lambda x: x + 1.5, np.arange(9).reshape(3, 3), namespace='numpy1')\n    array([[1.5, 2.5, 3.5],\n           [4.5, 5.5, 6.5],\n           [7.5, 8.5, 9.5]])\n\n    \u003e\u003e\u003e optree.tree_flatten(np.arange(9).reshape(3, 3), namespace='numpy2')\n    (\n        [0, 1, 2, 3, 4, 5, 6, 7, 8],\n        PyTreeSpec(\n            CustomTreeNode(ndarray[{'shape': (3, 3), 'dtype': dtype('int64')}], [*, *, *, *, *, *, *, *, *]),\n            namespace='numpy2'\n        )\n    )\n    # Explicitly casts `float`s to `np.int64`\n    \u003e\u003e\u003e optree.tree_map(lambda x: x + 1.5, np.arange(9).reshape(3, 3), namespace='numpy2')\n    array([[1, 2, 3],\n           [4, 5, 6],\n           [7, 8, 9]])\n\n    # Children are also `np.ndarray`s, recurse without termination condition.\n    \u003e\u003e\u003e optree.tree_flatten(np.arange(9).reshape(3, 3), namespace='numpy3')\n    Traceback (most recent call last):\n        ...\n    RecursionError: Maximum recursion depth exceeded during flattening the tree.\n\n    \u003e\u003e\u003e optree.tree_flatten(torch.arange(9).reshape(3, 3), namespace='torch1')\n    (\n        [0, 1, 2, 3, 4, 5, 6, 7, 8],\n        PyTreeSpec(\n            CustomTreeNode(Tensor[False], [[*, *, *], [*, *, *], [*, *, *]]),\n            namespace='torch1'\n        )\n    )\n    # Implicitly casts `float`s to `torch.float32`\n    \u003e\u003e\u003e optree.tree_map(lambda x: x + 1.5, torch.arange(9).reshape(3, 3), namespace='torch1')\n    tensor([[1.5000, 2.5000, 3.5000],\n            [4.5000, 5.5000, 6.5000],\n            [7.5000, 8.5000, 9.5000]])\n\n    # Children are also `torch.Tensor`s, recurse without termination condition.\n    \u003e\u003e\u003e optree.tree_flatten(torch.arange(9).reshape(3, 3), namespace='torch2')\n    Traceback (most recent call last):\n        ...\n    RecursionError: Maximum recursion depth exceeded during flattening the tree.\n    ```\n\n### `None` is Non-leaf Node vs. `None` is Leaf\n\nThe [`None`](https://docs.python.org/3/library/constants.html#None) object is a special object in the Python language.\nIt serves some of the same purposes as `null` (a pointer does not point to anything) in other programming languages, which denotes a variable is empty or marks default parameters.\nHowever, the `None` object is a singleton object rather than a pointer.\nIt may also serve as a sentinel value.\nIn addition, if a function has returned without any return value or the return statement is omitted, the function will also implicitly return the `None` object.\n\nBy default, the `None` object is considered a non-leaf node in the tree with arity 0, i.e., _**a non-leaf node that has no children**_.\nThis is like the behavior of an empty tuple.\nWhile flattening a tree, it will remain in the tree structure definitions rather than in the leaves list.\n\n```python\n\u003e\u003e\u003e tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5}\n\u003e\u003e\u003e optree.tree_flatten(tree)\n([1, 2, 3, 4, 5], PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': None, 'd': *}))\n\u003e\u003e\u003e optree.tree_flatten(tree, none_is_leaf=True)\n([1, 2, 3, 4, None, 5], PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': *, 'd': *}, NoneIsLeaf))\n\u003e\u003e\u003e optree.tree_flatten(1)\n([1], PyTreeSpec(*))\n\u003e\u003e\u003e optree.tree_flatten(None)\n([], PyTreeSpec(None))\n\u003e\u003e\u003e optree.tree_flatten(None, none_is_leaf=True)\n([None], PyTreeSpec(*, NoneIsLeaf))\n```\n\nOpTree provides a keyword argument `none_is_leaf` to determine whether to consider the `None` object as a leaf, like other opaque objects.\nIf `none_is_leaf=True`, the `None` object will be placed in the leaves list.\nOtherwise, the `None` object will remain in the tree specification (structure).\n\n```python\n\u003e\u003e\u003e import torch\n\n\u003e\u003e\u003e linear = torch.nn.Linear(in_features=3, out_features=2, bias=False)\n\u003e\u003e\u003e linear._parameters  # a container has None\nOrderedDict({\n    'weight': Parameter containing:\n              tensor([[-0.6677,  0.5209,  0.3295],\n                      [-0.4876, -0.3142,  0.1785]], requires_grad=True),\n    'bias': None\n})\n\n\u003e\u003e\u003e optree.tree_map(torch.zeros_like, linear._parameters)\nOrderedDict({\n    'weight': tensor([[0., 0., 0.],\n                      [0., 0., 0.]]),\n    'bias': None\n})\n\n\u003e\u003e\u003e optree.tree_map(torch.zeros_like, linear._parameters, none_is_leaf=True)\nTraceback (most recent call last):\n    ...\nTypeError: zeros_like(): argument 'input' (position 1) must be Tensor, not NoneType\n\n\u003e\u003e\u003e optree.tree_map(lambda t: torch.zeros_like(t) if t is not None else 0, linear._parameters, none_is_leaf=True)\nOrderedDict({\n    'weight': tensor([[0., 0., 0.],\n                      [0., 0., 0.]]),\n    'bias': 0\n})\n```\n\n### Key Ordering for Dictionaries\n\nThe built-in Python dictionary (i.e., [`builtins.dict`](https://docs.python.org/3/library/stdtypes.html#dict)) is an unordered mapping that holds the keys and values.\nThe leaves of a dictionary are the values. Although since Python 3.6, the built-in dictionary is insertion ordered ([PEP 468](https://peps.python.org/pep-0468)).\nThe dictionary equality operator (`==`) does not check for key ordering.\nTo ensure [referential transparency](https://en.wikipedia.org/wiki/Referential_transparency) that \"equal `dict`\" implies \"equal ordering of leaves\", the order of values of the dictionary is sorted by the keys.\nThis behavior is also applied to [`collections.defaultdict`](https://docs.python.org/3/library/collections.html#collections.defaultdict).\n\n```python\n\u003e\u003e\u003e optree.tree_flatten({'a': [1, 2], 'b': [3]})\n([1, 2, 3], PyTreeSpec({'a': [*, *], 'b': [*]}))\n\u003e\u003e\u003e optree.tree_flatten({'b': [3], 'a': [1, 2]})\n([1, 2, 3], PyTreeSpec({'a': [*, *], 'b': [*]}))\n```\n\nIf users want to keep the values in the insertion order in pytree traversal, they should use [`collections.OrderedDict`](https://docs.python.org/3/library/collections.html#collections.OrderedDict), which will take the order of keys under consideration:\n\n```python\n\u003e\u003e\u003e OrderedDict([('a', [1, 2]), ('b', [3])]) == OrderedDict([('b', [3]), ('a', [1, 2])])\nFalse\n\u003e\u003e\u003e optree.tree_flatten(OrderedDict([('a', [1, 2]), ('b', [3])]))\n([1, 2, 3], PyTreeSpec(OrderedDict({'a': [*, *], 'b': [*]})))\n\u003e\u003e\u003e optree.tree_flatten(OrderedDict([('b', [3]), ('a', [1, 2])]))\n([3, 1, 2], PyTreeSpec(OrderedDict({'b': [*], 'a': [*, *]})))\n```\n\n**Since OpTree v0.9.0, the key order of the reconstructed output dictionaries from `tree_unflatten` is guaranteed to be consistent with the key order of the input dictionaries in `tree_flatten`.**\n\n```python\n\u003e\u003e\u003e leaves, treespec = optree.tree_flatten({'b': [3], 'a': [1, 2]})\n\u003e\u003e\u003e leaves, treespec\n([1, 2, 3], PyTreeSpec({'a': [*, *], 'b': [*]}))\n\u003e\u003e\u003e optree.tree_unflatten(treespec, leaves)\n{'b': [3], 'a': [1, 2]}\n\u003e\u003e\u003e optree.tree_map(lambda x: x, {'b': [3], 'a': [1, 2]})\n{'b': [3], 'a': [1, 2]}\n\u003e\u003e\u003e optree.tree_map(lambda x: x + 1, {'b': [3], 'a': [1, 2]})\n{'b': [4], 'a': [2, 3]}\n```\n\nThis property is also preserved during serialization/deserialization.\n\n```python\n\u003e\u003e\u003e leaves, treespec = optree.tree_flatten({'b': [3], 'a': [1, 2]})\n\u003e\u003e\u003e leaves, treespec\n([1, 2, 3], PyTreeSpec({'a': [*, *], 'b': [*]}))\n\u003e\u003e\u003e restored_treespec = pickle.loads(pickle.dumps(treespec))\n\u003e\u003e\u003e optree.tree_unflatten(treespec, leaves)\n{'b': [3], 'a': [1, 2]}\n\u003e\u003e\u003e optree.tree_unflatten(restored_treespec, leaves)\n{'b': [3], 'a': [1, 2]}\n```\n\n\u003e [!NOTE]\n\u003e Note that there are no restrictions on the `dict` to require the keys to be comparable (sortable).\n\u003e There can be multiple types of keys in the dictionary.\n\u003e The keys are sorted in ascending order by `key=lambda k: k` first if capable otherwise fallback to `key=lambda k: (f'{k.__class__.__module__}.{k.__class__.__qualname__}', k)`. This handles most cases.\n\u003e\n\u003e ```python\n\u003e \u003e\u003e\u003e sorted({1: 2, 1.5: 1}.keys())\n\u003e [1, 1.5]\n\u003e \u003e\u003e\u003e sorted({'a': 3, 1: 2, 1.5: 1}.keys())\n\u003e Traceback (most recent call last):\n\u003e     ...\n\u003e TypeError: '\u003c' not supported between instances of 'int' and 'str'\n\u003e \u003e\u003e\u003e sorted({'a': 3, 1: 2, 1.5: 1}.keys(), key=lambda k: (f'{k.__class__.__module__}.{k.__class__.__qualname__}', k))\n\u003e [1.5, 1, 'a']\n\u003e ```\n\n--------------------------------------------------------------------------------\n\n## Changelog\n\nSee [CHANGELOG.md](https://github.com/metaopt/optree/blob/HEAD/CHANGELOG.md).\n\n--------------------------------------------------------------------------------\n\n## License\n\nOpTree is released under the Apache License 2.0.\n\nOpTree is heavily based on JAX's implementation of the PyTree utility, with deep refactoring and several improvements.\nThe original licenses can be found at [JAX's Apache License 2.0](https://github.com/google/jax/blob/HEAD/LICENSE) and [Tensorflow's Apache License 2.0](https://github.com/tensorflow/tensorflow/blob/HEAD/LICENSE).\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmetaopt%2Foptree","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fmetaopt%2Foptree","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fmetaopt%2Foptree/lists"}