Quick StartΒΆ

Numba compiles Python code to LLVM IR which can be natively executed at runtime much faster than pure Python code. The first step to using Numba is becoming familiar with the jit decorator, which tells Numba which functions to compile:

from numba import jit

@jit
def sum(x, y):
    return x + y

The very basic example above is compiled for any compatible input types automatically when the sum function is called. The result is a new function with performance comparable to a compiled function written in C (assuming best case scenario; more on that later). To compile for specific input types, we can tell Numba what those input types are:

@jit('f8(f8,f8)')
def sum(x, y):
    return x + y

The string above passed to the jit decorator tells Numba the return type is an 8 byte float, and the two arguments passed in are also 8 byte floats. The string takes the form returntype(arg1type, arg2type, ...).

One of the main features of Numba is it’s support for NumPy arrays. The following example shows how a function can be compiled that takes a NumPy array of floats as an input:

@jit('f8(f8[:])')
def sum1d(array):
    sum = 0.0
    for i in range(array.shape[0]):
        sum += array[i]
    return sum

There are two main things to notice in the example above. The input argument is specified by the string f8[:], which means a 1d array of 8 byte floats. A 2d array would be specified as f8[:, :], a 3d array as f8[:, :, :], and so on. The other thing to take note of is the array indexing and shape method call, and the fact that we’re iterating over a NumPy array using Python. Normally this would be terribly slow and would be cause for writing a NumPy ufunc in C, but the performance of the code above is the same as NumPy’s sum method.

Numba can also infer the array type automatically like other elementary types:

@jit
def sum1d(array):
    ...

Numba’s elementary built in types in are summarized in the table below and can be found in the numba namespace.

Type Name Alias Result Type
boolean b1 uint8 (char)
bool_ b1 uint8 (char)
byte u1 unsigned char
uint8 u1 uint8 (char)
uint16 u2 uint16
uint32 u4 uint32
uint64 u8 uint64
char i1 signed char
int8 i1 int8 (char)
int16 i2 int16
int32 i4 int32
int64 i8 int64
float_ f4 float32
float32 f4 float32
double f8 float64
float64 f8 float64
complex64 c8 float complex
complex128 c16 double complex

Native platform-dependent types are also available under names such as int_, short, ulonglong, etc.

Function signatures can also be expressed with the type objects directly as opposed to using strings. For example:

from numba import jit, f8

@jit(f8(f8[:]))
def sum1d(array):
    ...

In the example above, the argument type object is passed in to the return type object’s constructor.

Numba attempts to compile everything down to LLVM IR, but in some cases this isn’t (yet) possible. If Numba can’t infer the type of a variable or doesn’t support a particular data type, it falls back to using Python objects. This is of course much slower. If you’re having performance issues and suspect Python objects are to blame, you can use the nopython flag to force Numba to abort if it can’t avoid using Python objects:

@jit(nopython=True)
def sum1d(array):
    ...

Another useful debugging tool is Numba’s new inspect_types method. This can be called for any Numba compiled function to get a listing of the Numba IR generated from the Python code as well as the inferred types of each variable:

>>> sum1d.inspect_types()
sum1d (array(float64, 1d, A),) -> float64
``-----------------------------------------------------------------------``
# --- LINE 5 ---

@jit('f8(f8[:])')

# --- LINE 6 ---

def sum1d(array):

    # --- LINE 7 ---
    # label 0
    #   $0.1 = const(<type 'float'>, 0.0)  :: float64
    #   sum = $0.1  :: float64

    sum = 0.0

    # --- LINE 8 ---
    #   jump 6
    # label 6
    #   $6.1 = global(range: <built-in function range>)  :: range
    #   $6.2 = getattr(attr=shape, value=array)  :: (int64 x 1)
    #   $6.3 = const(<type 'int'>, 0)  :: int32
    #   $6.4 = getitem(index=$6.3, target=$6.2)  :: int64
    #   $6.5 = call $6.1($6.4, )  :: (int64,) -> range_state64
    #   $6.6 = getiter(value=$6.5)  :: range_iter64
    #   jump 26
    # label 26
    #   $26.1 = iternext(value=$6.6)  :: int64
    #   $26.2 = itervalid(value=$6.6)  :: bool
    #   branch $26.2, 29, 50
    # label 29
    #   $29.1 = $26.1  :: int64
    #   i = $29.1  :: int64

    for i in range(array.shape[0]):

        # --- LINE 9 ---
        # label 49
        #   del $6.6
        #   $29.2 = getitem(index=i, target=array)  :: float64
        #   $29.3 = sum + $29.2  :: float64
        #   sum = $29.3  :: float64
        #   jump 26

        sum += array[i]

    # --- LINE 10 ---
    #   jump 50
    # label 50
    #   return sum

    return sum

For get a better feel of what Numba can do, see Examples.