Numba is a perfect project to speed up my code, thank you guys to develop it
I have a question “How to compile all functions when importing the module?”. For example, there is a module that computes vector distances:
# search.py
#search.py
import numpy
def calEuclideanDistance(query, data):
dists = []
for vec in data:
dist = numpy.sqrt(numpy.sum(numpy.square(query - vec)))
dists.append(round(dist, 3))
return dists
def calInnerDistance(query, data):
dists = []
for vec in data:
dist = numpy.dot(vec, query)
dists.append(round(dist, 3))
return dists
def get_topk(dists, topk):
sorted_id = numpy.argsort(dists)
return sorted(dists)[0:topk], sorted_id[0:topk]
class SimilaritySearch():
def __init__(self, data: list, cal: str = 'L2', topk: int = 5):
self.data = numpy.array(data)
self.cal = cal
self.topk = topk
def __call__(self, query: list):
query = numpy.array(query)
if self.cal == 'L2':
dists = calEuclideanDistance(query, self.data)
elif self.cal == 'IP':
dists = calInnerDistance(query, self.data)
dists = numpy.array(dists)
dis, ids = get_topk(dists, self.topk)
return dis, ids
And I will import the module in my code:
#test.py
import importlib
from pathlib import Path
path = Path('./search.py')
modname = 'search'
spec = importlib.util.spec_from_file_location(modname, path)
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
op = getattr(module, 'SimilaritySearch')
obj = op(data=[[1,2,3],[4,5,6]], cal='L2')
res = obj(query=[[4,5,6], [7,8,9]])
print(res)
I see that jit_module
will automatically replace functions declared in a module with their jit
-wrapped equivalents, so does it support jitted when importing a module in another file? What should I do when I don’t want to modify or can’t modify search.py but want to JIT the functions in the search when called by the imported module?