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( ...@@ -67,6 +67,7 @@ setuptools.setup(
description='A domain specific language(DSL) for tensor computations.', description='A domain specific language(DSL) for tensor computations.',
install_requires=[ install_requires=[
'numpy', 'numpy',
'decorator',
], ],
zip_safe=False, zip_safe=False,
packages=[ packages=[
......
"""Tag class for TVM operators.""" """Tag class for TVM operators."""
from functools import wraps from decorator import decorate
class TagScope(object): class TagScope(object):
"""Tag scope object to set tag for operators, working as context """Tag scope object to set tag for operators, working as context
...@@ -22,11 +22,10 @@ class TagScope(object): ...@@ -22,11 +22,10 @@ class TagScope(object):
TagScope.current = self._old_scope TagScope.current = self._old_scope
def __call__(self, fdecl): def __call__(self, fdecl):
@wraps(fdecl) def tagged_fdecl(func, *args, **kwargs):
def tagged_fdecl(*args, **kwargs):
with self: with self:
return fdecl(*args, **kwargs) return func(*args, **kwargs)
return tagged_fdecl return decorate(fdecl, tagged_fdecl)
def tag_scope(tag): def tag_scope(tag):
......
...@@ -4,5 +4,5 @@ apt-get update && apt-get install -y python-pip python-dev python3-dev ...@@ -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 # 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 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 pip2 install nose pylint numpy nose-timer cython decorator
pip3 install nose pylint numpy nose-timer cython 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