treex.compact_module

A decorator that enable the definition of functional Modules

Source code in treex/module.py

def compact_module(f) -> type:
    """
    A decorator that enable the definition of functional Modules
    """
    name = utils._get_name(f)

    @functools.wraps(f)
    @to.compact
    def __call__(self, *args, **kwargs):
        return f(*args, **kwargs)

    module_class = type(
        name,
        (Module,),
        dict(
            __call__=__call__,
        ),
    )

    return module_class