Wrap @jit decorator

[Damian Barabonkov]

Hi all,

Is there a good way to wrap the @jit decorator? What I am attempting is to have @jit(parallel=True) dispatch a parallelized function normally, but when the function is being called already in a parallelized environment (such as joblib on top), to dispatch the function as if `parallel=False` to avoid parallelization within parallelization (which is detrimental to performance).

What I have so far as a proof of concept:

def jit(**decor_kwargs):
    """Modify the numba jit decorator for better parallel performance."""

    def modified_decorator(func):
        parallel_fn = numba_jit(**decor_kwargs)(func)

        def wrapper_func(*args, **kwargs):
            # TODO: Logic to dispatch parallel or non-parallel function
            return parallel_fn(*args, **kwargs)

        return wrapper_func

    return modified_decorator

But even this benign code issues warnings such as

- "Compilation is falling back to object mode WITH looplifting enabled because Function "..." failed type inference due to: non-precise type pyobject".

- "NumbaWarning: Function "..." was compiled in object mode without forceobj=True."

Thanks for your help!