o Ni$@sUdZddlmZddlZddlZddlZddlZddlZddlm Z m Z m Z m Z ddl mZmZddlmZddlmZmZmZmZmZmZmZmZddlmZddlmZmZm Z m!Z!m"Z"m#Z#dd l$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+dd l,m-Z-m.Z.m/Z/erddl0Z0dd l1m2Z2m3Z3m4Z4dd l$m5Z5m6Z6m7Z7m8Z8m9Z9ed e:e7dZ;gdZej?dddddde>Gdddee)Z@[>GdddZAeAZBdeCd<eZDdeCd<[AerddlEmFZFeFdZGedZHeFdZIed ZJGd!d"d"eeGeHeIeJfZKdd'd(ZLed)d*dd/d0ZMe dd)d*dd3d0ZM dd)d*dd6d0ZMeLeMed7dd?d@ZN[M[Le ddddAddEdFZOeddGdFZO ddddAddJdFZOddKdLZPejQddPdQZRddUdVZSddYdZZTdd]d^ZUddbdcZVddedfZWddjdkZXddldmZYddqdrZZddwdxZ[ddydzZ\dd{d|Z]dd~dZ^dddZ_dddZ`dddZadddZbdddZcdddZddddZedddZfdddZgdddZhdddZie:de@e:deTeUe!e'jjdeke@ekeVeWe"e'jldeme@emeXeYe"e'jndeoe@eoeZe[ee'jpde e@e efege e'jqde e@e e^e_ee'jrde e@e e`eaee'jsde e@e edeee"e'jtde(e@e(eheie#e'judi ZvdeCd<e@eoe\e]ee'jpdZwe@e ebecee'jsdZxdS)z&Registry for custom pytree node types.) annotationsN) OrderedDict defaultdictdeque namedtuple) itemgetter methodcaller)Lock) TYPE_CHECKINGAnyCallableClassVarGeneric NamedTupleTypeVaroverload) AutoEntry MappingEntryNamedTupleEntry PyTreeEntry SequenceEntryStructSequenceEntry)ChildrenMetaData PyTreeKindStructSequenceTis_namedtuple_classis_structseq_class)safe_ziptotal_order_sortedunzip2) Collection GeneratorIterable)KTVTCustomTreeNode FlattenFunc UnflattenFuncCustomTreeNodeType)bound)register_pytree_noderegister_pytree_node_classunregister_pytree_nodedict_insertion_ordered slotsT)initrepreqfrozenc@sbeZdZUdZded<ded<ded<ejdkrd ed <eZd ed <e j Z d ed<dZ ded<dS)PyTreeNodeRegistryEntryz>A dataclass that stores the information of a pytree node type.zbuiltins.type[Collection[T]]typeFlattenFunc[T] flatten_funcUnflattenFunc[T]unflatten_funcr0zdataclasses.KW_ONLY_zbuiltins.type[PyTreeEntry]path_entry_typerkindstr namespaceN) __name__ __module__ __qualname____doc____annotations__sys version_inforr?rCUSTOMr@rCrLrLC/home/ubuntu/.local/lib/python3.10/site-packages/optree/registry.pyr8Is   r8c@s$eZdZUdZded<d ddZdS) GlobalNamespacerLzClassVar[tuple[()]] __slots__returnrBcCdS)NzrLselfrLrLrM__repr__`zGlobalNamespace.__repr__N)rPrB)rDrErFrOrHrTrLrLrLrMrN]s rNrB__GLOBAL_NAMESPACEr __REGISTRY_LOCK) ParamSpec_P_T_GetP_GetTc@s eZdZdddZdd d ZdS)_CallableWithGetargs_P.argskwargs _P.kwargsrPrZcOtNNotImplementedErrorrSr^r`rLrLrM__call__rrUz_CallableWithGet.__call__ _GetP.args _GetP.kwargsr\cOrbrcrdrfrLrLrMgetvrUz_CallableWithGet.getN)r^r_r`rarPrZ)r^rhr`rirPr\)rDrErFrgrjrLrLrLrMr]qs r]rjCallable[_GetP, _GetT]rPDCallable[[Callable[_P, _T]], _CallableWithGet[_P, _T, _GetP, _GetT]]csdfdd }|S)NfuncCallable[_P, _T]rP&_CallableWithGet[_P, _T, _GetP, _GetT]cs |_|Srcrj)rmrprLrM decoratorsz_add_get..decorator)rmrnrProrL)rjrqrLrprM_add_getzsrrrA)rCclsr9rCPyTreeNodeRegistryEntry | NonecCdSrcrLrsrCrLrLrMpytree_node_registry_getrwNone#dict[type, PyTreeNodeRegistryEntry]cCrurcrLrvrLrLrMrwrx type | NoneDdict[type, PyTreeNodeRegistryEntry] | PyTreeNodeRegistryEntry | Nonecs>|turd}|dur|turt|std|dt|ts(td|d|dur^t|dhtfddt D}Wdn1sJwYt |r\t |t<t|t<|S|dkrot ||f}|duro|St |r|turzt S|turtSt |}|dur|St|rt tSt|rt tSdS)aLook up the pytree node registry. >>> register_pytree_node.get() # doctest: +IGNORE_WHITESPACE,ELLIPSIS { : PyTreeNodeRegistryEntry( type=, flatten_func=, unflatten_func=, path_entry_type=, kind=, namespace='' ), : PyTreeNodeRegistryEntry( type=, flatten_func=, unflatten_func=, path_entry_type=, kind=, namespace='' ), : PyTreeNodeRegistryEntry( type=, flatten_func=, unflatten_func=, path_entry_type=, kind=, namespace='' ), ... } >>> register_pytree_node.get(defaultdict) # doctest: +IGNORE_WHITESPACE,ELLIPSIS PyTreeNodeRegistryEntry( type=, flatten_func=, unflatten_func=, path_entry_type=, kind=, namespace='' ) >>> register_pytree_node.get(frozenset) # frozenset is considered as a leaf node None Args: cls (type or None, optional): The class of the pytree node to retrieve. If not provided, all the registered pytree nodes in the namespace are returned. namespace (str, optional): The namespace of the registry to retrieve. If not provided, the global namespace is used. Returns: If the ``cls`` is not provided, a dictionary of all the registered pytree nodes in the namespace is returned. If the ``cls`` is provided, the corresponding registry entry is returned if the ``cls`` is registered as a pytree node. Otherwise, :data:`None` is returned, i.e., the ``cls`` is represented as a leaf node. rANzExpected a class or None, got .$The namespace must be a string, got csi|] }|jvr|j|qSrL)rCr9).0handler namespacesrLrM s  z,pytree_node_registry_get..)rVrinspectisclass TypeError isinstancerB frozensetrW_NODETYPE_REGISTRYvalues_Cis_dict_insertion_ordered&_DICT_INSERTION_ORDERED_REGISTRY_ENTRYdict-_DEFAULTDICT_INSERTION_ORDERED_REGISTRY_ENTRYrrjrrr)rsrCregistryrrLrrMrwsN<         )r?r;r:r=r<type[Collection[T]]r?type[PyTreeEntry]cCst|s td|dt|rt|tstd|d|tur0t|ts0td|d|dkr8td|turA|}d}n||f}t t |||||t |||||dt |<Wd|S1sgwY|S) aExtend the set of types that are considered internal nodes in pytrees. See also :func:`register_pytree_node_class` and :func:`unregister_pytree_node`. The ``namespace`` argument is used to avoid collisions that occur when different libraries register the same Python type with different behaviors. It is recommended to add a unique prefix to the namespace to avoid conflicts with other libraries. Namespaces can also be used to specify the same class in different namespaces for different use cases. .. warning:: For safety reasons, a ``namespace`` must be specified while registering a custom type. It is used to isolate the behavior of flattening and unflattening a pytree node type. This is to prevent accidental collisions between different libraries that may register the same type. Args: cls (type): A Python type to treat as an internal pytree node. flatten_func (callable): A function to be used during flattening, taking an instance of ``cls`` and returning a triple or optionally a pair, with (1) an iterable for the children to be flattened recursively, and (2) some hashable metadata to be stored in the treespec and to be passed to the ``unflatten_func``, and (3) (optional) an iterable for the tree path entries to the corresponding children. If the entries are not provided or given by :data:`None`, then `range(len(children))` will be used. unflatten_func (callable): A function taking two arguments: the metadata that was returned by ``flatten_func`` and stored in the treespec, and the unflattened children. The function should return an instance of ``cls``. path_entry_type (type, optional): The type of the path entry to be used in the treespec. (default: :class:`AutoEntry`) namespace (str): A non-empty string that uniquely identifies the namespace of the type registry. This is used to isolate the registry from other modules that might register a different custom behavior for the same type. Returns: The same type as the input ``cls``. Raises: TypeError: If the input type is not a class. TypeError: If the path entry class is not a subclass of :class:`PyTreeEntry`. TypeError: If the namespace is not a string. ValueError: If the namespace is an empty string. ValueError: If the type is already registered in the registry. .. versionadded:: 0.12.0 The ``path_entry_type`` argument to specify the path entry type used in :meth:`PyTreeSpec.accessors` and :func:`tree_flatten_with_accessor`. If not provided, :class:`AutoEntry` will be used. Examples: >>> # Register a Python type with lambda functions >>> register_pytree_node( ... set, ... lambda s: (sorted(s), None, None), ... lambda _, children: set(children), ... namespace='set', ... ) >>> # Register a Python type into a namespace >>> import torch >>> register_pytree_node( ... torch.Tensor, ... flatten_func=lambda tensor: ( ... (tensor.cpu().detach().numpy(),), ... {'dtype': tensor.dtype, 'device': tensor.device, 'requires_grad': tensor.requires_grad}, ... ), ... unflatten_func=lambda metadata, children: torch.tensor(children[0], **metadata), ... namespace='torch2numpy', ... ) >>> # doctest: +SKIP >>> tree = {'weight': torch.ones(size=(1, 2)).cuda(), 'bias': torch.zeros(size=(2,))} >>> tree {'weight': tensor([[1., 1.]], device='cuda:0'), 'bias': tensor([0., 0.])} >>> # Flatten without specifying the namespace >>> tree_flatten(tree) # `torch.Tensor`s are leaf nodes ([tensor([0., 0.]), tensor([[1., 1.]], device='cuda:0')], PyTreeSpec({'bias': *, 'weight': *})) >>> # Flatten with the namespace >>> tree_flatten(tree, namespace='torch2numpy') ( [array([0., 0.], dtype=float32), array([[1., 1.]], dtype=float32)], PyTreeSpec( { 'bias': CustomTreeNode(Tensor[{'dtype': torch.float32, 'device': device(type='cpu'), 'requires_grad': False}], [*]), 'weight': CustomTreeNode(Tensor[{'dtype': torch.float32, 'device': device(type='cuda', index=0), 'requires_grad': False}], [*]) }, namespace='torch2numpy' ) ) >>> # Register the same type with a different namespace for different behaviors >>> def tensor2flatparam(tensor): ... return [torch.nn.Parameter(tensor.reshape(-1))], tensor.shape, None ... ... def flatparam2tensor(metadata, children): ... return children[0].reshape(metadata) ... ... register_pytree_node( ... torch.Tensor, ... flatten_func=tensor2flatparam, ... unflatten_func=flatparam2tensor, ... namespace='tensor2flatparam', ... ) >>> # Flatten with the new namespace >>> tree_flatten(tree, namespace='tensor2flatparam') ( [ Parameter containing: tensor([0., 0.], requires_grad=True), Parameter containing: tensor([1., 1.], device='cuda:0', requires_grad=True) ], PyTreeSpec( { 'bias': CustomTreeNode(Tensor[torch.Size([2])], [*]), 'weight': CustomTreeNode(Tensor[torch.Size([1, 2])], [*]) }, namespace='tensor2flatparam' ) ) Expected a class, got r}(Expected a subclass of PyTreeEntry, got r~rA(The namespace cannot be an empty string.r?rCN)rrr issubclassrrVrrB ValueErrorrWr register_noder8r)rsr;r=r?rCregistration_keyrLrLrMr,s@  r,r str | Nonetype[PyTreeEntry] | None2Callable[[CustomTreeNodeType], CustomTreeNodeType]cCrurcrLrsr?rCrLrLrMr-r-cCrurcrLrrLrLrMr-rCustomTreeNodeType | str | NoneGCustomTreeNodeType | Callable[[CustomTreeNodeType], CustomTreeNodeType]cstus ttrdurtddkrtdddur&tdtur6tts6tddkr>tddurLd%fd d }|StsYtd d durctdttrmt t sutdd t fdddDst fdddDstdt jd&dd}tt tjdjd'd!d"}|_|_ttd#jd$S)(aExtend the set of types that are considered internal nodes in pytrees. See also :func:`register_pytree_node` and :func:`unregister_pytree_node`. The ``namespace`` argument is used to avoid collisions that occur when different libraries register the same Python type with different behaviors. It is recommended to add a unique prefix to the namespace to avoid conflicts with other libraries. Namespaces can also be used to specify the same class in different namespaces for different use cases. .. warning:: For safety reasons, a ``namespace`` must be specified while registering a custom type. It is used to isolate the behavior of flattening and unflattening a pytree node type. This is to prevent accidental collisions between different libraries that may register the same type. Args: cls (type, optional): A Python type to treat as an internal pytree node. path_entry_type (type, optional): The type of the path entry to be used in the treespec. (default: :class:`AutoEntry`) namespace (str, optional): A non-empty string that uniquely identifies the namespace of the type registry. This is used to isolate the registry from other modules that might register a different custom behavior for the same type. Returns: The same type as the input ``cls`` if the argument presents. Otherwise, return a decorator function that registers the class as a pytree node. Raises: TypeError: If the path entry class is not a subclass of :class:`PyTreeEntry`. TypeError: If the namespace is not a string. TypeError: If the class does not define the required method pairs. ValueError: If the namespace is an empty string. ValueError: If the type is already registered in the registry. .. versionadded:: 0.12.0 The ``TREE_PATH_ENTRY_TYPE`` class variable to specify the path entry type used in :meth:`PyTreeSpec.accessors` and :func:`tree_flatten_with_accessor`. If not provided, :class:`AutoEntry` will be used. .. versionadded:: 0.18.0 Previously, this function looked for methods named ``tree_flatten`` and ``tree_unflatten`` for the given class. Since version 0.18.0, it prefers methods named ``__tree_flatten__`` and ``__tree_unflatten__`` instead. The old method names are still supported for backward compatibility, but it is recommended to use the new method names. The method resolution follows this priority: 1. If both ``__tree_flatten__`` and ``__tree_unflatten__`` are defined, use them directly. 2. If both ``tree_flatten`` and ``tree_unflatten`` are defined, wrap them as dunder methods. 3. If neither complete pair is available, raise a :exc:`TypeError` suggesting the new method names. This function is a thin wrapper around :func:`register_pytree_node`, and provides a class-oriented interface: .. code-block:: python @register_pytree_node_class(namespace='foo') class Special: TREE_PATH_ENTRY_TYPE = GetAttrEntry def __init__(self, x, y): self.x = x self.y = y def __tree_flatten__(self): return ((self.x, self.y), None, ('x', 'y')) @classmethod def __tree_unflatten__(cls, metadata, children): return cls(*children) @register_pytree_node_class('mylist') class MyList(UserList): TREE_PATH_ENTRY_TYPE = SequenceEntry def __tree_flatten__(self): return self.data, None, None @classmethod def __tree_unflatten__(cls, metadata, children): return cls(*children) # Legacy style (still supported but not recommended) @register_pytree_node_class(namespace='legacy') class LegacyStyleMyList(UserList): def tree_flatten(self): # Implementation automatically wrapped as __tree_flatten__ return self.data, None, None @classmethod def tree_unflatten(cls, metadata, children): # Implementation automatically wrapped as __tree_unflatten__ return cls(*children) Nz?Cannot specify `namespace` when the first argument is a string.rArz.decoratorrr}TREE_PATH_ENTRY_TYPErc3 |] }tt|dVqdSrccallablegetattrrmethodrrLrM Gs  z-register_pytree_node_class..)__tree_flatten____tree_unflatten__c3rrcrrrrLrMrLs ) tree_flattentree_unflattenzh must define both `__tree_flatten__` and `__tree_unflatten__` methods for registration as a pytree node.rSCustomTreeNode[T]Qtuple[Children[T], MetaData] | tuple[Children[T], MetaData, Iterable[Any] | None]cSs|Src)rrRrLrLrMrWz4register_pytree_node_class..__tree_flatten____func__type[CustomTreeNode[T]]metadatarchildren Children[T]cSs |||Src)r)rsrrrLrLrMr^s z6register_pytree_node_class..__tree_unflatten__rr)rsr*rPr*)rSrrPr)rsrrrrrrPr)rVrrBrrrrrrrrall functoolswrapsr classmethodrrrr,r)rsr?rCrqrrrL)rsrCr?rMr-sZb        cCst|s td|d|turt|tstd|d|dkr&td|tur/|}d}n||f}tt ||t |WdS1sKwYdS)aRemove a type from the pytree node registry. See also :func:`register_pytree_node` and :func:`register_pytree_node_class`. This function is the inverse operation of function :func:`register_pytree_node`. Args: cls (type): A Python type to remove from the pytree node registry. namespace (str): The namespace of the pytree node registry to remove the type from. Returns: The removed registry entry. Raises: TypeError: If the input type is not a class. TypeError: If the namespace is not a string. ValueError: If the namespace is an empty string. ValueError: If the type is a built-in type that cannot be unregistered. ValueError: If the type is not found in the registry. Examples: >>> # Register a Python type with lambda functions >>> register_pytree_node( ... set, ... lambda s: (sorted(s), None, None), ... lambda _, children: set(children), ... namespace='temp', ... ) >>> # Unregister the Python type >>> unregister_pytree_node(set, namespace='temp') rr}r~rArN) rrrrVrrBrrWrunregister_noderpop)rsrCrrLrLrMr.ws " $r.modeboolGenerator[None]ccs|turt|tstd|d|dkrtd|tur d}ttj|dd}tt ||Wdn1s>> tree = {'b': (2, [3, 4]), 'a': 1, 'c': None, 'd': 5} >>> tree_flatten(tree) # doctest: +IGNORE_WHITESPACE ( [1, 2, 3, 4, 5], PyTreeSpec({'a': *, 'b': (*, [*, *]), 'c': None, 'd': *}) ) >>> with dict_insertion_ordered(True, namespace='some-namespace'): # doctest: +IGNORE_WHITESPACE ... tree_flatten(tree, namespace='some-namespace') ( [2, 3, 4, 1, 5], PyTreeSpec({'b': (*, [*, *]), 'a': *, 'c': None, 'd': *}, namespace='some-namespace') ) .. warning:: The dictionary sorting mode is a global setting and is **not thread-safe**. It is recommended to use this context manager in a single-threaded environment. Args: mode (bool): The dictionary sorting mode to set. namespace (str): The namespace to set the dictionary sorting mode for. r~r}rArF)inherit_global_namespaceN) rVrrBrrrWrrset_dict_insertion_orderedr)rrCprevrLrLrMr/s$(r/itemsIterable[tuple[KT, VT]]list[tuple[KT, VT]]cCst|tddS)Nr)key)r r)rrLrLrM _sorted_itemssrr>tuple[tuple[()], None]cCrQ)N)rLNrL)r>rLrLrM _none_flattenrUrr Iterable[Any]cCs$t}tt|||urtddS)NzExpected no children.)objectnextiterr)r>rsentinelrLrLrM_none_unflattensrtup tuple[T, ...]tuple[tuple[T, ...], None]cC|dfSrcrLrrLrLrM_tuple_flattenr Iterable[T]cCt|Src)tupler>rrLrLrM_tuple_unflattenrrlstlist[T]tuple[list[T], None]cCrrcrL)rrLrLrM _list_flattenrrcCrrc)listrrLrLrM_list_unflattenrrdct dict[KT, VT]/tuple[tuple[VT, ...], list[KT], tuple[KT, ...]]cCs"tt|\}}|t||fSrc)r!rrrrkeysrrLrLrM _dict_flattensrrlist[KT]r Iterable[VT]cCtt||SrcrrrrrLrLrM_dict_unflattenrcCt|\}}|t||fSrcr!rrrrLrLrM_dict_insertion_ordered_flattenrcCrrcrrrLrLrM!_dict_insertion_ordered_unflattenrrOrderedDict[KT, VT]cCrrcrrrLrLrM_ordereddict_flattenrrcCrrc)rrrrLrLrM_ordereddict_unflattenrrdefaultdict[KT, VT]Otuple[tuple[VT, ...], tuple[Callable[[], VT] | None, list[KT]], tuple[KT, ...]]cCt|\}}}||j|f|fSrc)rdefault_factoryrr dict_metadataentriesrLrLrM_defaultdict_flatten"rr!tuple[Callable[[], VT], list[KT]]cC|\}}t|t||Src)rrrrrrrLrLrM_defaultdict_unflatten.rcCrrc)rrrrLrLrM&_defaultdict_insertion_ordered_flatten7rrcCrrc)rrrrLrLrM(_defaultdict_insertion_ordered_unflattenCrrdeqdeque[T]tuple[deque[T], int | None]cCs ||jfSrcmaxlen)rrLrLrM_deque_flattenLs r r  int | NonecCs t||dS)Nr )r)r rrLrLrM_deque_unflattenP r NamedTuple[T])tuple[tuple[T, ...], type[NamedTuple[T]]]cC |t|fSrcr9rrLrLrM_namedtuple_flattenTrrtype[NamedTuple[T]]cCs||SrcrLrsrrLrLrM_namedtuple_unflattenYrrseqStructSequence[T]-tuple[tuple[T, ...], type[StructSequence[T]]]cCrrcr)rrLrLrM_structseq_flatten]rrtype[StructSequence[T]]cCs||SrcrLrrLrLrM_structseq_unflattenarr)r?r@z6dict[type | tuple[str, type], PyTreeNodeRegistryEntry]rrL)rjrkrPrl)rsr9rCrBrPrtrc)rsryrCrBrPrz)rsr{rCrBrPr|) r;r:r=r<rsrr?rrCrBrPr)rsrr?rrCrrPr)rsr*r?rrCrBrPr*)rsrr?rrCrrPr)rsr9rCrBrPr8)rrrCrBrPr)rrrPr)r>ryrPr)r>ryrrrPry)rrrPr)r>ryrrrPr)rrrPr)r>ryrrrPr)rrrPr)rrrrrPr)rrrPr)rrrrrPr)rrrPr)rrrrrPr)rrrPr )r r rrrPr)rrrPr)rsrrrrPr)rrrPr)rsrrrrPr)yrG __future__r contextlib dataclassesrrrI collectionsrrrroperatorrr threadingr typingr r r r rrrr optree._Croptree.accessorsrrrrrr optree.typingrrrrrrr optree.utilsrr r!builtinscollections.abcr"r#r$r%r&r'r(r)r9r*__all__rJSLOTS dataclassr8rNrVrHrWtyping_extensionsrXrYrZr[r\r]rrrwr,r-r.contextmanagerr/rrrrrrrrrrrrrrrrrr rrrrrNONErTUPLErLISTrDICT NAMEDTUPLE ORDEREDDICT DEFAULTDICTDEQUESTRUCTSEQUENCErrrrLrLrLrMsX  (  $    k&   45  .                  C