diff --git a/python/tvm/_ffi/function.py b/python/tvm/_ffi/function.py index 526d972f6d282740dbc77a0281a824a9c85ad9cb..90a0c18dec5d49634485904bfc34e3d4fdbe7a8b 100644 --- a/python/tvm/_ffi/function.py +++ b/python/tvm/_ffi/function.py @@ -280,15 +280,23 @@ def _get_api(f): return flocal(*args) return my_api_func -def _init_api(namespace): +def _init_api(namespace, target_module_name=None): """Initialize api for a given module name - mod : str - The name of the module. + namespace : str + The namespace of the source registry + + target_module_name : str + The target module name if different from namespace """ - assert namespace.startswith("tvm.") - prefix = namespace[4:] - _init_api_prefix(namespace, prefix) + target_module_name = ( + target_module_name if target_module_name else namespace) + if namespace.startswith("tvm."): + prefix = namespace[4:] + _init_api_prefix(target_module_name, namespace[4:]) + else: + _init_api_prefix(target_module_name, namespace) + def _init_api_prefix(module_name, prefix): module = sys.modules[module_name] diff --git a/topi/tests/python/test_topi_conv2d.py b/topi/tests/python/test_topi_conv2d.py index 3e6978752b29c7ba2a7028cf9a601fc91b460c58..e7ea956eea78ecdf7bd0d48ef644e744106f156b 100644 --- a/topi/tests/python/test_topi_conv2d.py +++ b/topi/tests/python/test_topi_conv2d.py @@ -3,6 +3,7 @@ import os import numpy as np import tvm import topi +import topi.testing from tvm.contrib.pickle_memoize import memoize from topi.util import get_const_tuple diff --git a/topi/tests/python/test_topi_conv2d_hwcn.py b/topi/tests/python/test_topi_conv2d_hwcn.py index 7761a90a0eed35a18b6403ac139d97feda050da8..84962a0b478902179913187a7330e270fe47fd9c 100644 --- a/topi/tests/python/test_topi_conv2d_hwcn.py +++ b/topi/tests/python/test_topi_conv2d_hwcn.py @@ -3,6 +3,7 @@ import os import numpy as np import tvm import topi +import topi.testing from tvm.contrib.pickle_memoize import memoize from topi.util import get_const_tuple diff --git a/topi/tests/python/test_topi_conv2d_nhwc.py b/topi/tests/python/test_topi_conv2d_nhwc.py index 7fc5b841908f1c32711418ca768494c6bfc7eff3..40aa3e550a86121b02b0e0bf334026481ae60038 100644 --- a/topi/tests/python/test_topi_conv2d_nhwc.py +++ b/topi/tests/python/test_topi_conv2d_nhwc.py @@ -3,6 +3,7 @@ import os import numpy as np import tvm import topi +import topi.testing from tvm.contrib.pickle_memoize import memoize from topi.util import get_const_tuple diff --git a/topi/tests/python/test_topi_conv2d_transpose_nchw.py b/topi/tests/python/test_topi_conv2d_transpose_nchw.py index 712f64ca2c87e0de19461addb7da4b7fe28c8e6e..dcdddca9107be9192e12a4ca4c0340bb884c5216 100644 --- a/topi/tests/python/test_topi_conv2d_transpose_nchw.py +++ b/topi/tests/python/test_topi_conv2d_transpose_nchw.py @@ -2,6 +2,7 @@ import numpy as np import tvm import topi +import topi.testing from tvm.contrib.pickle_memoize import memoize from topi.util import get_const_tuple diff --git a/topi/tests/python/test_topi_dense.py b/topi/tests/python/test_topi_dense.py index cbae2055a3cd6041a29315ef87ef146070a7be3d..779e1af1f35e860771dc3fe48d56c79966596680 100644 --- a/topi/tests/python/test_topi_dense.py +++ b/topi/tests/python/test_topi_dense.py @@ -2,6 +2,7 @@ import numpy as np import tvm import topi +import topi.testing from topi.util import get_const_tuple from tvm.contrib.pickle_memoize import memoize diff --git a/topi/tests/python/test_topi_depthwise_conv2d.py b/topi/tests/python/test_topi_depthwise_conv2d.py index 8ecfdd063a940ec45f4db6cf58a4a9e77263ba0d..5a12d7abb20894963dd295a0847750b50f03daf2 100644 --- a/topi/tests/python/test_topi_depthwise_conv2d.py +++ b/topi/tests/python/test_topi_depthwise_conv2d.py @@ -1,5 +1,6 @@ import tvm import topi +import topi.testing import numpy as np from scipy import signal from topi.util import get_const_tuple diff --git a/topi/tests/python/test_topi_depthwise_conv2d_back_input.py b/topi/tests/python/test_topi_depthwise_conv2d_back_input.py index 4280e92bdcb913495745ce53380cb978f3ee420d..663364c5345beb040d2ceecaeff293f736859390 100644 --- a/topi/tests/python/test_topi_depthwise_conv2d_back_input.py +++ b/topi/tests/python/test_topi_depthwise_conv2d_back_input.py @@ -5,6 +5,7 @@ from tvm.contrib.pickle_memoize import memoize from scipy import signal from topi.util import get_const_tuple from topi.nn.util import get_pad_tuple +import topi.testing from topi.cuda.depthwise_conv2d import schedule_depthwise_conv2d_backward_input_nhwc diff --git a/topi/tests/python/test_topi_depthwise_conv2d_back_weight.py b/topi/tests/python/test_topi_depthwise_conv2d_back_weight.py index 5aa65d2c4039963b43141c8dd2f3b36f9a24fbc2..25d08e1b188287c0bbe2a16f2aea88f4579c6584 100644 --- a/topi/tests/python/test_topi_depthwise_conv2d_back_weight.py +++ b/topi/tests/python/test_topi_depthwise_conv2d_back_weight.py @@ -1,5 +1,6 @@ import tvm import topi +import topi.testing import numpy as np from tvm.contrib.pickle_memoize import memoize from scipy import signal diff --git a/topi/tests/python/test_topi_math.py b/topi/tests/python/test_topi_math.py index 2c31d0fea62c94f33d1d88aa3fcd958aa4646e8a..dc3c015d4f25ee86d6ce07b808aac933fe695ae4 100644 --- a/topi/tests/python/test_topi_math.py +++ b/topi/tests/python/test_topi_math.py @@ -1,6 +1,7 @@ import numpy as np import tvm import topi +import topi.testing from topi import util