Cache njit function at import or install

I have a package that contains many public njit functions. Of course, these functions are only compiled just-in-time as needed but I was wondering if there was a way to compile these functions either at installation OR at import? Is there a standard approach for this? I know that there is a cache=True option but this still requires the function to be called at least one time by the user. I was thinking that I could sneak this into by calling the functions there but any suggestions would be greatly appreciated.

1 Like

@seanlaw I would like to know that as well. I think the only way right now to do that is to use AOT compilation

I was thinking: if calling those njit functions from a background thread will it work for the first time, in order to do not freeze the app at startup?

@rpopovici Yes, I’ve read that too but AOT sounds like a worse option as it is not optimized for the local machine and the code is pretty ugly from a maintenance perspective. I think performing the JIT compilation and placing the function call inside if is another option. This means that import time may be high as it will include compile time but, for a data science workflow (i.e., a Jupyter notebook where import happens once), this might be fine but the solution needs to be easily portable as it is not just for me and my machine.

I recently became aware of the approach used by Qibo to force compilation at install time. Each decorator has a signature and cache=True, e.g.

@njit("int64(int64, int32[:])",
def multicontrol_index(g, qubits):


Then after installation in, the library is imported:

class CustomInstall(install):
    def run(self):
        from qibo import K

This forces the compilation and caching of all functions.

Note that this will break binary distribution when you do python bdist, so Qibo uploads sdists to PyPI instead, creating them with python sdist.

(Many thanks to Stefano Carrazza for pointing me to this approach).

1 Like