Looking for a good way to use @jit even numba is not installed

Hello,

I would like to write python code that can run either with numba as well as without numba installed.
Is the follwing approach ok.?

At import section I define a dummy jit decorator in case numba can’t be imported:

try:
  from numba import jit
except:
  def jit(func=None, *args, **kwargs):
    if not(callable(func)):
      return jit
    else:
      return func

Is there a reason not to? Did I forget anything?

As an alternative to the try / except block, in this case, consider the importlib.util.find_spec.

Thanks for the hint.
Can I still leave the @jit decorator in the code using the importlib.util.find_spec?

My topic is how to leave the numba code unchanged (even the decorators) with minimum effort.
In most cases I prefer not to import extra libraries if possible.

I tried to cover 4 different decorators:

  1. @jit with no arguments
  2. @jit(*args) with arbitrary positional arguments
  3. @jit(**kwargs) with arbitrary keyword arguments
  4. @jit(*args, **kwargs) with arbitrary mixture of positional and keyword arguments

Is your hinted way shorter or more stable?

I guess it depends, to an extent, on your use-case(s).

One option is to tuck the find_spec into the decorator itself:

def try_njit(*args, **kwargs):
    if find_spec("numba") is not None:
        from numba import njit
        if args and callable(args[0]):
            return njit(args[0])
        return njit(*args, **kwargs)
    if args and callable(args[0]):
        return args[0]
    return lambda _: _

Another is something like this:

import importlib
from importlib.util import find_spec


def _void_njit(*args, **kws):
    if args and callable(args[0]):
        return args[0]
    return lambda _: _


try_njit = find_spec("numba") and importlib.import_module("numba").njit or _void_njit

or along the lines thereof.

Nothing wrong per se with your try / except approach, but I’d probably catch ModuleNotFoundError rather than the generic Exception.

Hope it helps.

1 Like

Thanks a lot. I will try a little.