annotate_demo.py

import numpy as np
from numba import jit


@jit
def simple(x, out):
    for i in range(x.shape[0]):
        out[i] = x[i]**2 + 1


@jit
def loop_lifting(x):
    out = np.empty_like(x)

    for i in range(x.shape[0]):
        out[i] = x[i]**2 + 1

    return out


def math_function_that_i_forgot_to_compile(y):
    return (1.0 - y) / (1.0 + y)


@jit
def accidental_object_mode(x, out):
    for i in range(x.shape[0]):
        out[i] = math_function_that_i_forgot_to_compile(x[i])


x = np.ones(1000, dtype=np.float32)
out = np.empty_like(x)

simple(x, out)
loop_lifting(x)
accidental_object_mode(x, out)