How do I use @jit wrapper interactively?

I want to quickly test Numba on various functions in IPython session. Going through my modules and decorating everything with @jit seem like a rather tedious task. Is there a way to use it with a function without modifying the original?

That would also help measure wrapped and non-wrapped functions in the same notebook.

Hi @abitrolly ,

the jit decorator is a completely normal decorator. As such it is just syntactic sugar around a function that modifies other functions.

def f(x):
    return x

# The above is basically the same as
def f(x):
    return x
f = decorator(f)

You can therefore do this

from numba import jit

def f(x):
    return x

# This is the plain replacement of @jit
f_jit = jit(f) 

#This is used when you want to pass extra options as in @jit(nopython=True)
f_jit_options = jit(nopython=True)(f) 

The reason that the latter works, is that jit is actually a little clever and figures out if the first argument passed is a function or not. If not it assumes you pass options and creates a specialised decorator on the fly, which is then applied to f. Otherwise the default options are used.

1 Like

jit(nopython=True)(f) is just what I need. Thanks. :smiley: