Search code examples
pythoncachingimportnumbajit

How to import a cached numba function without a Python definition


Consider the following function in numba, which just serves as an example:

import numba as nb
import numpy as np

@nb.njit('float64(float64[::1])', cache=True)
def total (x):
    ''' Sum the elements of an array. '''

    total = 0
    for i in range(x.shape[0]):
        total += x[i]

    return total

x = np.arange(100,dtype=np.float64)
print(total(x))

Since I have specified the cache=True option, two files are created in the __pycache__ folder, one .nbc file and one .nbi file. I assume that these files contain all (compiled) information about the function. Let's say I delete the Python file that defines the function (i..e, the above code).

Can I still use compiled/cached functions? In other words, can I import them without having the original .py file that defined them?


Solution

  • @Michael Szczesny's comment about ahead of time compilation is exactly what I wanted. To make it work, I have adapted the code as follows:

    from numba.pycc import CC
    
    cc = CC('my_module')
    
    @cc.export('total', 'float64(float64[::1])')
    def total (x):
        ''' Sum the elements of an array. '''
    
        total = 0
        for i in range(x.shape[0]):
            total += x[i]
    
        return total
    
    if __name__ == "__main__":
        cc.compile()
    

    After running this file, a binary module file (.pyd) is saved in the directory and you can use it as follows:

    import numpy as np
    from my_module import total
    
    x = np.arange(100,dtype=np.float64)
    print(total(x))