Search code examples
pythonbenchmarkingsignaturejitnumba

numba.jit give function and signature as inputs


My objective:

I have a function and I want to see how it performs for a list of numba options/parameters (nopython, no gil, parallel, etc. and Signature!)

so I want to do something like:

def foo(a):
   return a*2


for signature in list:
   foo_jit = numba.jit(foo, signature)
   
   print("speed for jitted func given signature")

Sadly this raises an error, because @numba.jit() as decorator takes a signature for the first input, and numba.jit() as a function takes a function as input. I've tried assigning the signature before jitting the function, but you can't jit a function twice.


Solution

  • To test such approach, I would give a try to defining a function within a loop:

    for signature in list:
    
        @jit(signature)
        def foo(a):
           return a*2
    
        start = time.time()
        foo(whatever_data_you_have)
        end = time.time()
        print("Elapsed (with compilation) = %s" % (end - start))
    

    You can of course read in the guide that first time is with compilation, thus if you want time without compilation, add one more time the testing block, this it will load from cache and wont count compilation in elapsed time.