numba-benchmarks.py

@jit(nopython=True)
def g0(x, y):
    s = 0
    for i in range(len(x)):
        s += x[i] * y[i]
    return s

@jit(nopython=True)
def g1(x, y):
    s = 0
    for xi, yi in zip(x.flat, y.flat):
        s += xi * yi
    return s

@jit(nopython=True)
def g2(x, y):
    s = 0
    for i in np.ndindex(*x.shape):
        s += x[i] * y[i]
    return s

@jit(nopython=True)
def g3(x, y):
    s = 0
    for i, xi in np.ndenumerate(x):
        s += xi * y[i]
    return s

x = np.random.randn(1000, 10)
y = np.random.randn(1000, 10)
x0 = x.ravel()
y0 = y.ravel()

# compile each function
g0(x0, y0)
g1(x0, y0)
g2(x0, y0)
g3(x0, y0)
g1(x, y)
g2(x, y)
g3(x, y)

print('1d arguments')
%timeit g0(x0, y0)
%timeit g1(x0, y0)
%timeit g2(x0, y0)
%timeit g3(x0, y0)

print('2d arguments')
%timeit g1(x, y)
%timeit g2(x, y)
%timeit g3(x, y)

# 1d arguments
# 100000 loops, best of 3: 10.7 µs per loop
# 10000 loops, best of 3: 18.8 µs per loop
# 10000 loops, best of 3: 22.5 µs per loop
# 100000 loops, best of 3: 14.4 µs per loop
# 2d arguments
# 100000 loops, best of 3: 17.5 µs per loop
# 10000 loops, best of 3: 29.2 µs per loop
# 10000 loops, best of 3: 21.8 µs per loop