numba_jit_decorator_numpy_function_compilation_benchmark.py

1from numba import jit
2import numpy as np
3import time
4
5@jit(nopython=True) # Set "nopython" mode for best performance, equivalent to @njit
6def go_fast(a): # Function is compiled to machine code when called the first time
7    trace = 0.0
8    for i in range(a.shape[0]):
9        trace += np.tanh(a[i, i])
10    return a + trace
11
12x = np.arange(10000).reshape(100, 100)
13
14# DO NOT REPORT TIMING DURING FIRST EXECUTION AS IT INCLUDES COMPILATION TIME
15go_fast(x)
16
17# NOW THE FUNCTION IS COMPILED, RE-RUN IT TO SEE THE SPEED
18start = time.time()
19go_fast(x)
20end = time.time()
21print("Elapsed (with compilation) = %s" % (end - start))
22
23# NOW THE FUNCTION IS COMPILED, RE-RUN IT TO SEE THE SPEED
24start = time.time()
25go_fast(x)
26end = time.time()
27print("Elapsed (after compilation) = %s" % (end - start))