Commit 5cdc8604 by ziheng Committed by Tianqi Chen

[TAG] Fix signature of decorated function (#228)

* [TAG] Fix signature of decorated function

* Add dep
parent 72fcd4e6
......@@ -67,6 +67,7 @@ setuptools.setup(
description='A domain specific language(DSL) for tensor computations.',
install_requires=[
'numpy',
'decorator',
],
zip_safe=False,
packages=[
......
"""Tag class for TVM operators."""
from functools import wraps
from decorator import decorate
class TagScope(object):
"""Tag scope object to set tag for operators, working as context
......@@ -22,11 +22,10 @@ class TagScope(object):
TagScope.current = self._old_scope
def __call__(self, fdecl):
@wraps(fdecl)
def tagged_fdecl(*args, **kwargs):
def tagged_fdecl(func, *args, **kwargs):
with self:
return fdecl(*args, **kwargs)
return tagged_fdecl
return func(*args, **kwargs)
return decorate(fdecl, tagged_fdecl)
def tag_scope(tag):
......
......@@ -4,5 +4,5 @@ apt-get update && apt-get install -y python-pip python-dev python3-dev
# the version of the pip shipped with ubuntu may be too lower, install a recent version here
cd /tmp && wget https://bootstrap.pypa.io/get-pip.py && python3 get-pip.py && python2 get-pip.py
pip2 install nose pylint numpy nose-timer cython
pip3 install nose pylint numpy nose-timer cython
pip2 install nose pylint numpy nose-timer cython decorator
pip3 install nose pylint numpy nose-timer cython decorator
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment