Merge pull request #88 from opendilab/dev/dict

dev(hansbug): add register custom dicts

docs/source/api_doc/tree/tree.rst

@@ -220,3 +220,11 @@ loads
 
 .. autofunction:: loads
 
+
+register_dict_type
+----------------------------
+
+.. autofunction:: register_dict_type
+
+
+

pytest.ini

@@ -1,5 +1,5 @@
 [pytest]
-timeout = 20
+timeout = 60
 markers =
     unittest
     benchmark

requirements-doc.txt

@@ -1,7 +1,7 @@
 Jinja2~=3.0.0
 sphinx~=3.2.0
 sphinx_rtd_theme~=0.4.3
-enum_tools
+enum_tools~=0.9.0
 sphinx-toolbox
 plantumlcli>=0.0.4
 packaging

requirements-test-extra.txt

@@ -1,2 +1,2 @@
 jax[cpu]>=0.3.25; platform_system != 'Windows'
-torch>=1.1.0; python_version < '3.12'
+torch>=1.1.0,<2.1.0; python_version < '3.12'

test/testings/__init__.py

@@ -0,0 +1 @@
+from .mapping import CustomMapping

test/testings/mapping.py

@@ -0,0 +1,15 @@
+import collections.abc
+
+
+class CustomMapping(collections.abc.Mapping):
+    def __init__(self, **kwargs):
+        self._kwargs = kwargs
+
+    def __getitem__(self, __key):
+        return self._kwargs[__key]
+
+    def __len__(self):
+        return len(self._kwargs)
+
+    def __iter__(self):
+        yield from self._kwargs

test/tree/general/base.py

@@ -1,3 +1,4 @@
+import collections.abc
 import unittest
 from functools import reduce
 from operator import __mul__
@@ -7,8 +8,11 @@
 import pytest
 from hbutils.testing import cmdv
 
+from treevalue import register_dict_type
 from treevalue.tree import func_treelize, TreeValue, raw, mapping, delayed, FastTreeValue
 from ..tree.base import get_treevalue_test
+from ...testings import CustomMapping
+
 
 
 def get_fasttreevalue_test(treevalue_class: Type[FastTreeValue]):
@@ -813,4 +817,32 @@ def test_unpack(self):
             assert y == pytest.approx(7.7)
             assert z is None
 
+        def test_init_with_custom_mapping_type(self):
+            origin_t = CustomMapping(a=1, b=2, c={'x': 15, 'y': CustomMapping(z=100)})
+            t = treevalue_class(origin_t)
+            assert t == treevalue_class({'a': 1, 'b': 2, 'c': {'x': 15, 'y': {'z': 100}}})
+
+        def test_init_with_custom_type(self):
+            class _CustomMapping:
+                def __init__(self, **kwargs):
+                    self._kwargs = kwargs
+
+                def __getitem__(self, __key):
+                    return self._kwargs[__key]
+
+                def __len__(self):
+                    return len(self._kwargs)
+
+                def __iter__(self):
+                    yield from self._kwargs
+
+                def iter_items(self):
+                    yield from self._kwargs.items()
+
+            register_dict_type(_CustomMapping, _CustomMapping.iter_items)
+
+            origin_t = _CustomMapping(a=1, b=2, c={'x': 15, 'y': _CustomMapping(z=100)})
+            t = treevalue_class(origin_t)
+            assert t == treevalue_class({'a': 1, 'b': 2, 'c': {'x': 15, 'y': {'z': 100}}})
+
     return _TestClass

test/tree/integration/test_torch.py

@@ -161,3 +161,20 @@ def forward(self, x):
             'c': torch.Size([14]),
             'd': torch.Size([2, 5, 3]),
         })
+
+    @skipUnless(torch is not None and vpip('torch') and OS.linux and vpython < '3.11', 'torch required')
+    def test_moduledict(self):
+        with torch.no_grad():
+            md = torch.nn.ModuleDict({
+                'a': torch.nn.Linear(3, 5),
+                'b': torch.nn.Linear(3, 6),
+            })
+            t = FastTreeValue(md)
+
+            input_ = torch.randn(2, 3)
+            output_ = t(input_)
+
+            assert output_.shape == FastTreeValue({
+                'a': (2, 5),
+                'b': (2, 6),
+            })

test/tree/tree/base.py

@@ -7,9 +7,10 @@
 from hbutils.testing import OS, cmdv
 
 from test.tree.tree.test_constraint import GreaterThanConstraint
-from treevalue import raw, TreeValue, delayed, ValidationError
+from treevalue import raw, TreeValue, delayed, ValidationError, register_dict_type
 from treevalue.tree.common import create_storage
 from treevalue.tree.tree.constraint import cleaf
+from ...testings import CustomMapping
 
 try:
     _ = reversed({}.keys())
@@ -787,4 +788,32 @@ def test_repr_jpeg(self):
             assert min_size <= len(_repr_jpeg_) <= max_size, \
                 f'Size within [{min_size!r}, {max_size!r}] required, but {len(_repr_jpeg_)!r} found.'
 
+        def test_init_with_custom_mapping_type(self):
+            origin_t = CustomMapping(a=1, b=2, c={'x': 15, 'y': CustomMapping(z=100)})
+            t = treevalue_class(origin_t)
+            assert t == treevalue_class({'a': 1, 'b': 2, 'c': {'x': 15, 'y': {'z': 100}}})
+
+        def test_init_with_custom_type(self):
+            class _CustomMapping:
+                def __init__(self, **kwargs):
+                    self._kwargs = kwargs
+
+                def __getitem__(self, __key):
+                    return self._kwargs[__key]
+
+                def __len__(self):
+                    return len(self._kwargs)
+
+                def __iter__(self):
+                    yield from self._kwargs
+
+                def iter_items(self):
+                    yield from self._kwargs.items()
+
+            register_dict_type(_CustomMapping, _CustomMapping.iter_items)
+
+            origin_t = _CustomMapping(a=1, b=2, c={'x': 15, 'y': _CustomMapping(z=100)})
+            t = treevalue_class(origin_t)
+            assert t == treevalue_class({'a': 1, 'b': 2, 'c': {'x': 15, 'y': {'z': 100}}})
+
     return _TestClass

treevalue/tree/integration/torch.py

@@ -1,6 +1,8 @@
 import warnings
 from functools import wraps
 
+from ..tree import register_dict_type
+
 try:
     import torch
     from torch.utils._pytree import _register_pytree_node
@@ -20,3 +22,10 @@ def register_for_torch(cls):
 
     register_for_torch(TreeValue)
     register_for_torch(FastTreeValue)
+
+try:
+    from torch.nn import ModuleDict
+except (ModuleNotFoundError, ImportError):
+    pass
+else:
+    register_dict_type(ModuleDict, ModuleDict.items)

treevalue/tree/tree/__init__.py

@@ -5,4 +5,4 @@
 from .io import loads, load, dumps, dump
 from .service import jsonify, clone, typetrans, walk
 from .structural import subside, union, rise
-from .tree import TreeValue, delayed, ValidationError
+from .tree import TreeValue, delayed, ValidationError, register_dict_type

treevalue/tree/tree/tree.pxd

@@ -14,6 +14,7 @@ cdef class _SimplifiedConstraintProxy:
     cdef readonly Constraint cons
 
 cdef Constraint _c_get_constraint(object cons)
+cpdef register_dict_type(object type_, object f_items)
 
 cdef class ValidationError(Exception):
     cdef readonly TreeValue _object

treevalue/tree/tree/tree.pyx

@@ -31,24 +31,59 @@ cdef inline object _item_unwrap(object v):
         return v
 
 _GET_NO_DEFAULT = SingletonMark('get_no_default')
+_KNOWN_DICT_TYPES = {
+    Mapping: Mapping.items,
+}
 
-cdef inline TreeStorage _dict_unpack(dict d):
+@cython.binding(True)
+cpdef inline register_dict_type(object type_, object f_items):
+    """
+    Overview:
+        Register dict-like type for TreeValue.
+
+    :param type_: Type to register.
+    :param f_items: Function to get items, the format should be like ``dict.items``.
+
+    .. note::
+        If torch detected, the ``torch.nn.ModuleDict`` is registered by default.
+
+    """
+    if isinstance(object, type):
+        _KNOWN_DICT_TYPES[type_] = f_items
+    else:
+        raise TypeError(f'Not a type - {type_!r}.')
+
+_DEFAULT_STORAGE = create_storage({})
+
+cdef inline TreeStorage _generic_dict_unpack(object d):
     cdef str k
     cdef object v
     cdef dict result = {}
 
-    for k, v in d.items():
+    cdef object d_items = None
+    if isinstance(d, dict):
+        d_items = d.items()
+    else:
+        for d_type, df_items in _KNOWN_DICT_TYPES.items():
+            if isinstance(d, d_type):
+                d_items = df_items(d)
+                break
+        if d_items is None:
+            raise TypeError(f'Unknown dict type - {d!r}.')
+
+    for k, v in d_items:
         if isinstance(v, dict):
-            result[k] = _dict_unpack(v)
-        elif isinstance(v, TreeValue):
+            result[k] = _generic_dict_unpack(v)
+        if isinstance(v, TreeValue):
             result[k] = v._detach()
         else:
-            result[k] = v
+            try:
+                result[k] = _generic_dict_unpack(v)
+            except TypeError:
+                result[k] = v
 
     return create_storage(result)
 
-_DEFAULT_STORAGE = create_storage({})
-
 cdef class _SimplifiedConstraintProxy:
     def __cinit__(self, Constraint cons):
         self.cons = cons
@@ -131,13 +166,14 @@ cdef class TreeValue:
                 self._child_constraints = data._child_constraints
             else:
                 self.constraint = _c_get_constraint(constraint)
-        elif isinstance(data, dict):
-            self._st = _dict_unpack(data)
-            self.constraint = _c_get_constraint(constraint)
         else:
-            raise TypeError(
-                "Unknown initialization type for tree value - {type}.".format(
-                    type=repr(type(data).__name__)))
+            try:
+                self._st = _generic_dict_unpack(data)
+                self.constraint = _c_get_constraint(constraint)
+            except TypeError:
+                raise TypeError(
+                    "Unknown initialization type for tree value - {type}.".format(
+                        type=repr(type(data).__name__)))
 
     def __getnewargs_ex__(self):  # for __cinit__, when pickle.loads
         return ({},), {}