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:

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 =, 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): = numpy.array(data) = cal
        self.topk = topk

    def __call__(self, query: list):
        query = numpy.array(query)
        if == 'L2':
            dists = calEuclideanDistance(query,  
        elif == 'IP':
            dists = calInnerDistance(query,
        dists = numpy.array(dists)
        dis, ids = get_topk(dists, self.topk)
        return dis, ids

And I will import the module in my code:
import importlib
from pathlib import Path

path = Path('./')
modname = 'search'
spec = importlib.util.spec_from_file_location(modname, path)
module = importlib.util.module_from_spec(spec)
op = getattr(module, 'SimilaritySearch')

obj = op(data=[[1,2,3],[4,5,6]], cal='L2')
res = obj(query=[[4,5,6], [7,8,9]])

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 but want to JIT the functions in the search when called by the imported module?