How to jit compile with import module?

Numba is a perfect project to speed up my code, thank you guys to develop it :slight_smile:

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?