diff --git a/python/tvm/autotvm/task/code_hash.py b/python/tvm/autotvm/task/code_hash.py index 3331fc13c719..2bd053da7244 100644 --- a/python/tvm/autotvm/task/code_hash.py +++ b/python/tvm/autotvm/task/code_hash.py @@ -19,6 +19,7 @@ code hashing is used to check the consistence of schedule code and the parameters loaded from log """ +import functools import inspect import zlib @@ -35,6 +36,7 @@ def attach_code_hash(s): """ def decorator(func): + @functools.wraps(func) def wrapper(*args, **kwargs): func(*args, **kwargs) raw_hash = zlib.crc32("".join(inspect.getsourcelines(func)[0]).encode()) @@ -56,6 +58,7 @@ def attach_code_hash_to_arg(arg_idx=1): """ def decorator(func): + @functools.wraps(func) def wrapper(*args, **kwargs): func(*args, **kwargs) assert isinstance(args[arg_idx], schedule.Schedule) diff --git a/python/tvm/autotvm/task/task.py b/python/tvm/autotvm/task/task.py index 668832b8a86c..1f5827d7e9d0 100644 --- a/python/tvm/autotvm/task/task.py +++ b/python/tvm/autotvm/task/task.py @@ -21,6 +21,8 @@ func is a state-less function, or a string that registers the standard task. """ +import functools + import numpy as np from tvm import runtime @@ -411,6 +413,7 @@ def matmul(N, L, M, dtype): """ def _decorate(f): + @functools.wraps(f) def wrapper(*args, **kwargs): assert not kwargs, "Do not support kwargs in template function call" workload = args_to_workload(args, task_name) diff --git a/python/tvm/autotvm/task/topi_integration.py b/python/tvm/autotvm/task/topi_integration.py index 2558c7669ac9..32d8674640ed 100644 --- a/python/tvm/autotvm/task/topi_integration.py +++ b/python/tvm/autotvm/task/topi_integration.py @@ -26,6 +26,8 @@ See tvm/topi/python/topi/arm_cpu/depthwise_conv2d.py for example usage. """ +import functools + import tvm.te._ffi_api from tvm.target import Target from tvm.te import tensor @@ -149,6 +151,7 @@ def register_topi_compute(task_name, func=None): """ def _decorate(topi_compute): + @functools.wraps(topi_compute) @_register_task_compute(task_name) def wrapper(*args, **kwargs): """wrapper function for topi compute""" @@ -224,6 +227,7 @@ def register_topi_schedule(task_name, func=None): """ def _decorate(topi_schedule): + @functools.wraps(topi_schedule) @_register_task_schedule(task_name) def wrapper(outs, *args, **kwargs): """wrapper function for topi schedule"""