Function _check_obj_attr

imperative/python/megengine/traced_module/utils.py:121–145

(obj)

def _check_obj_attr(obj):
    # check if all the attributes of a obj is serializable
    from .pytree import tree_flatten
    from .pytree import SUPPORTED_LEAF_CLS, SUPPORTED_LEAF_TYPE, TreeDef
    from .expr import Expr
    from .traced_module import TracedModule, InternalGraph, NameSpace

    def _check_leaf_type(leaf):
        leaf_type = leaf if isinstance(leaf, type) else type(leaf)
        traced_module_types = [Expr, TreeDef, TracedModule, InternalGraph, NameSpace]
        return (
            issubclass(leaf_type, tuple(SUPPORTED_LEAF_CLS + traced_module_types))
            or leaf_type in SUPPORTED_LEAF_TYPE
        )

    for _, v in obj.items():
        leafs, _ = tree_flatten(v, is_leaf=lambda _: True)
        for leaf in leafs:
            assert _check_leaf_type(leaf), (
                "Type {} is not supported in TracedModule serialization by default. "
                "If you want to save this object to file, please call tm.register_supported_type({}) "
                "before saving.".format(
                    leaf if isinstance(leaf, type) else type(leaf), type(leaf).__name__
                )
            )


def _check_builtin_module_attr(mod):