Numba - Jit

This page can be downloaded as interactive jupyter notebook

Jit is a module from the Numba package and is a abbreviation for Just In Time compiler. The module makes it possible to compile Python methods to fast machine code (which can give a significant time improvement). The usage is very simple. The first step is to import the module:

from numba import jit
import numpy as np

The second step is to add the @jit decorator in front of every method that should be compiled.

Speed test:

The following cell contains three methods. All of them take a integer n and sum up all numbers from 0 to n-1. The first function is a native Python implementation with a single for-in-loop. The second function contains the same code, but this time with the @jit decorator. The third one is a Numpy implementation, that does not use any explicit loops, but instead the methods arange() and sum(). At last the iPython feature timeit is used, to run each method 100000 times and compare the runtime.

def sum_up(n):
    s = 0
    for i in range(n):
        s += i
    return s

@jit
def sum_up_jit(n):
    s = 0
    for i in range(n):
        s += i
    return s
    
def sum_up_numpy(n):
    v = np.arange(n)
    return np.sum(v)

n = 1000

print(sum_up(n),sum_up_jit(n),sum_up_numpy(n)) # check if all methods have same result and compile jit

t_reg = %timeit -o -q -n 100000 -r 1 sum_up(n)
t_jit = %timeit -o -q -n 100000 -r 1 sum_up_jit(n)
t_num = %timeit -o -q -n 100000 -r 1 sum_up_numpy(n)

print('jit was {0:.1f} times faster than regular python'.format(t_reg.average/t_jit.average))
print('numpy was {0:.1f} times faster than regular python'.format(t_reg.average/t_num.average))

499500 499500 499500
jit was 286.6 times faster than regular python
numpy was 9.4 times faster than regular python

Limitations

If jit gives such a significant time improvement, why not always use it? The reason is, that jit does not support all Python features and specially not many features of other libraries. Luckily it supports most of Numpys features albeit some of them only partially. The supportet features are listed at

• http://numba.pydata.org/numba-doc/dev/reference/pysupported.html (for Python features)
• http://numba.pydata.org/numba-doc/dev/reference/numpysupported.html (for Numpy features)

Whenever you try to compile a jit-decorated function that contains an unsupported feature, the compiler will fall back to Python, which can result in a half compiled method that is maybe even slower than the regular method without using jit. This can be prevented, by explicitly telling the compiler not to fall back to Python (nopython = True). Instead of falling back to Python, the compiler will then raise an error. But if the code will be compiled without an error, you can be sure that it has been fully compiled. Note that any function will only be compiled, if it is called at some point in the script.

Unsupported Python features

Although we get an error message whenever we try to use an unsupported feature, these errors are hard to debug because we won’t get the line number. Therefore we will have a look at some unsupported features:

1. Dictionaries: currently not supported
2. Strings: also currently not supported
3. Nested memory-managed objects (mmos): mmos are data types with mutable size like lists or sets. Nested mmos are e.g. lists of lists, sets of lists etc.
4. Calling uncompiled functions: not possible in nopython-mode
5. Lists, sets and arrays with varying data type: will not raise an error, but implicitly be casted to the same type

# @jit(nopython = True)
def a_python_function():
    return 

@jit(nopython = True)
def jit_test():
    # D = {1:2,2:3}           # 1. dictionary
    # s = 'a string'          # 2. strings (e.g. when returned)
    # LL = [[1,2],[3,4]]      # 3. nested memory-managed object 
    # r = a_python_function() # 4. calling a native Python method
    V = {5, True, 18.0}       # 5. set with int, bool and float -> will be casted to float
    return V
    
v = jit_test()
print(v)

{1.0, 18.0, 5.0}

The enumeration above is not complete. Usually it’s best practice to first implement a native Python / Numpy function and then try to compile it (add the decorator). If it does not work, try to get some informations from the compiler error message or comment some statements of the function.