Lazily importing jitted functions without importing numba

Following a short discussion on Gitter, I’m documenting an approach to importing a numba jitted function without A) importing numba, and B) compiling the function on each call.

The motivation is to reduce import times of libraries dependent on numba by delaying importing numba until necessary. On my 2020 windows computer, import numba takes 400 ms. Things add up.

Normally one can use e.g. lazyasd, which is available on pip and conda, but when jit is called (e.g. as a decorator), numba is loaded.

This is the solution:

import importlib
from lazyasd import lazyobject

# This is the lazy import of numba, and works great
@lazyobject
def numba():
    return importlib.import_module('numba')

# this function acts like _add(N), but doesn't import numba with the jit call
@lazyobject
def add():
    @numba.jit(nopython=True)
    def _add(N):
        "Docstring goes here"
        n = 0
        for i in range(N):
            n = n + 1
        return n
    return _add

Here, lazyobject expects the function to return a generic Python object which will be assigned the name of the outer function (thanks @Hannes). The function has the same performance (after being called once) as non-lazily version does (~140ns on my computer).

It is called by add(100).

Hope this is useful to someone.

1 Like

I just had another idea that I would like to share, but I cannot guarantee that there won’t be sand in the gears :stuck_out_tongue:

If this is something you need to do repeatedly in your code base I imagine that one may end up with a fair bit of boiler plate code, so I was wondering if one could make a custom decorator

from lazyasd import LazyObject

def lazyjit(f):
     return LazyObject(lambda: numba.njit(f), f.__globals__, f.__name__)


@lazyjit
def test(x=5):
    return x

Based on some extremely rudimentary play testing this seems to work as intended, but I tend to get scared when messing around with __globals__

1 Like