Numba#

Numba is an just-in-time compiler for Python code.

Goal: automatically convert Python code to highly efficient compiled code (CPU and GPU).

Installation:

conda install numba

Numba is very simple to use.#

Decorate your Python function with jit decorator to automatically create an efficient, compiled version of the function:

from numba import jit
from numpy import arange

# jit decorator tells Numba to compile this function.
# The argument types will be inferred by Numba when function is called.


@jit(nopython=True)
def sum2d(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i, j]
    return result


N = 1000
a = arange(N**2).reshape(N, N)
print(sum2d(a))

499999500000.0

def sum2d(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i, j]
    return result


sum2d = jit(nopython=True)(sum2d)

Performance test#

Let us compare the performance to a standard Python implementation:

def sum2d_py(arr):
    M, N = arr.shape
    result = 0.0
    for i in range(M):
        for j in range(N):
            result += arr[i, j]
    return result

Benchmark of Python version:

%timeit sum2d_py(a)

123 ms ± 2.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Benchmark of Python+Numba version:

%timeit sum2d(a)

933 µs ± 1.44 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Numba tries to be clever by default#

  • Infers data types of the arguments at runtime.

  • You can call a @jit function from another @jit function. Numba might use even use inlining!

  • Automatically generate C and Python wrapper code of the function

  • Tries to compile high-performant non-Python CPU code - if it fails, fall back to Python version.

    Use @jit(nopython=True) to force Numba to use the non-Python mode.

You can, but do not have to, take more control, such as:

  • Provide function signature explicitely

  • Generate GPU code (CUDA), CPU parallel code

  • Call external C functions

  • Compilation in advance

  • Vectorization of functions (use the vectorize decorator)

See http://numba.pydata.org/numba-doc/latest

Numba pros and cons#

Uses Just in Time compilation (JIT) to speed up Python code

  • [+++] Simple to use, highly automated

  • [-] Needs to be installed as external dependency

  • [-] Is like magic, which makes debugging hard