This notebook contains a demonstration of new features present in the 0.53.0 release of Numba. Whilst release notes are produced as part of the CHANGE_LOG, there's nothing like seeing code in action!
This release contains a few new features. In this notebook, the new CPU target features are demonstrated. The CUDA target also gained a lot of new features in 0.53.0 and @gmarkall has created a demo notebook especially for these!
Key internal changes:
@guvectorize
) (@guilhermeleobas).@njit(parallel=True)
) now supports Fortran ordered arrays (@DrTodd13 and @sklam).Intel also kindly sponsored research and development that lead to two new features for profiling the compiler:
Demonstrations of these compiler profiling features can be found in a different notebook.
This notebook will focus on:
import numba
assert numba.version_info.short >= (0, 53)
from numba.core import types
@numba.njit
def identity(x):
return x
identity()
is a generic function that accepts any type.
Let's define a 1-arity function type that takes and returns a intp
:
# Define a function type that takes an int and returns an int.
fn_sig = types.FunctionType(types.intp(types.intp))
fn_sig
FunctionType[int64(int64)]
fn_sig
is a subtype of identity()
. Therefore, we can use identity()
in any place that expects a fn_sig
.
Let's define a function that takes (fn_sig, intp)
:
@numba.njit((fn_sig, types.intp))
def invoke_callback(callback, arg):
return callback(arg)
# Disable compilation
invoke_callback.disable_compile()
invoke_callback.signatures
[(FunctionType[int64(int64)], int64)]
We can use identity
for the first argument.
invoke_callback(identity, 123)
123
The invoke_callback()
function can take any function that can be cast to the subtype fn_sig
. We will define two more functions that are compatible with fn_sig
.
@numba.njit
def cb_add_one(x):
return x + 1
@numba.njit
def cb_twice(x):
return x * 2
print('cb_add_one', invoke_callback(cb_add_one, 123))
print('cb_twice', invoke_callback(cb_twice, 123))
cb_add_one 124 cb_twice 246
No new signature needs to be compiled:
invoke_callback.signatures
[(FunctionType[int64(int64)], int64)]
Without declaring the expected function type, Numba would specialize to the given function:
@numba.njit
def invoke_callback_generic(callback, arg):
return callback(arg)
print(f"# of compiled version {len(invoke_callback_generic.signatures)}")
print("call identity() =", invoke_callback_generic(identity, 123))
print(f"# of compiled version {len(invoke_callback_generic.signatures)}")
print("call cb_add_one() =", invoke_callback_generic(cb_add_one, 123))
print(f"# of compiled version {len(invoke_callback_generic.signatures)}")
print("signatures", invoke_callback_generic.signatures)
# of compiled version 0 call identity() = 123 # of compiled version 1 call cb_add_one() = 124 # of compiled version 2 signatures [(type(CPUDispatcher(<function identity at 0x11dce0b80>)), int64), (type(CPUDispatcher(<function cb_add_one at 0x11f048310>)), int64)]
One other advantage to function subtyping is that first-class functions are not locked as previously required. We can request new subtypes from the functions.
First, let's look at the compiled versions of cb_add_one()
(there's just one):
cb_add_one.signatures
[(int64,)]
Now let's make a list that contains functions of the signature float64(float64)
and add our callback functions to it:
new_fn_sig = types.FunctionType(types.float64(types.float64))
lst = numba.typed.List.empty_list(new_fn_sig)
lst.append(cb_add_one)
lst.append(cb_twice)
A new float version of cb_add_one()
is compiled as requested during the .append()
:
cb_add_one.signatures
[(int64,), (float64,)]
We can then use the list of functions:
import numpy as np
@numba.njit
def many_callbacks(cblist, arg, incr):
return [cb(arg) + incr for cb in cblist]
print(many_callbacks(lst, 0.234, 10))
print(many_callbacks(lst, 0.234, np.arange(4)))
[11.234, 10.468] [array([1.234, 2.234, 3.234, 4.234]), array([0.468, 1.468, 2.468, 3.468])]
@guilhermeleobas added dynamic gufunc support, allowing gufuncs to be used without pre-defining their accepted function types.
Previously, gufuncs must be defined with a fix set of function types. For example:
@numba.guvectorize([(types.float64, types.float64[:]),
(types.complex128, types.complex128[:])], "()->()")
def static_twice(inp, out):
out[()] = inp * 2
inp = np.arange(5, dtype=np.float64)
print(f"{inp.dtype} version", static_twice(inp))
inp2 = inp + .5j
print(f"{inp2.dtype} version", static_twice(inp2))
float64 version [0. 2. 4. 6. 8.] complex128 version [0.+1.j 2.+1.j 4.+1.j 6.+1.j 8.+1.j]
With dynamic gufuncs, we can omit the function types. New compilation is triggered dynamically as needed. However, due to a limitation of type inference, the output argument must be specified.
# creates a dynamic gufunc by omitting the function types
@numba.guvectorize("()->()")
def dynamic_twice(inp, out):
out[()] = inp * 2
# Use the gufunc with different dtypes
out1 = np.zeros_like(inp)
dynamic_twice(inp, out1)
print(f"{inp.dtype} version", out1)
inp2 = inp + 0.5j
out2 = np.zeros_like(inp2)
dynamic_twice(inp2, out2)
print(f"{inp2.dtype} version", out2)
inp3 = inp.astype(np.intp)
out3 = np.zeros_like(inp3)
dynamic_twice(inp3, out3)
print(f"{inp3.dtype} version", out3)
float64 version [0. 2. 4. 6. 8.] complex128 version [0.+1.j 2.+1.j 4.+1.j 6.+1.j 8.+1.j] int64 version [0 2 4 6 8]