This notebook contains a demonstration of new features present in the 0.51.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 notebook contains new CPU target features, the CUDA target also gained a lot of new features in 0.51.0 and so has it's own demo notebook!
Key internal changes:
Intel also kindly sponsored research and development that lead to some exciting new features:
StructRef
(@sklam).Demonstrations of new features/changes include:
First, import the necessary from Numba and NumPy...
from numba import jit, njit, config, __version__, errors, literal_unroll, types
from numba.extending import overload
import numba
import numpy as np
assert tuple(int(x) for x in __version__.split('.')[:2]) >= (0, 51)
Numba has supported heterogeneous immutable containers (e.g. tuples!) and homogeneous mutable containers (lists and dictionaries) for some time, Numba 0.51 adds support for additional types of immutable heterogeneous containers. Practically these take the form of "lists of mixed type items" and "string key dictionaries mapping to any type of value", these are only supported by direct definition in @jit
decorated functions (i.e. can't pass them in from Python). Motivating these by example:
@njit
def mixed_type_list():
# a list of type [literal intp, unicode string, NumPy 1d array of float64]
x = [1, 'a', np.zeros(5)]
# getitem works for constant indexes (a literal value known at compile time)
print('getitem', x[1]) # 1 is constant, this prints 'a'
print('len', len(x)) # non-mutating call on the list is ok
# iteration requires `literal_unroll` as the type of the induction variable
# depends on the iteration, but works for constant values as before
y = [100, 'apple', 200, 'orange']
for i in literal_unroll(y):
print(i)
mixed_type_list()
Heterogeneously typed lists are immutable, attempted mutation is a compilation error...
@njit
def mixed_type_list_error():
# a list of type [literal intp, unicode string, NumPy 1d array of float64]
x = [1, 'a', np.zeros(5)]
x.append(2j) # illegal mutation
try:
mixed_type_list_error()
except errors.TypingError as e:
# CANNOT MUTATE A LITERAL LIST!
print("Cannot mutate a literal list!")
assert "Cannot mutate a literal list" in str(e)
Heterogeneously typed lists also carry their type information, including literal values, such that it's possible to dispatch based on their value types
def bar(x):
pass
@overload(bar)
def ol_bar(x):
# If the string "NOP" is in the list then return a no-operation function
# else capture the types as strings and return that!
# Note that heterogeneous lists use `.literal_value` to hold
# the types of the item whereas e.g. a tuple uses `.types`, this is because
# heterogeneous lists inherit from `types.Literal`.
# Look for NOP, do nothing!
if any([getattr(lv, 'literal_value', None) == "NOP" for lv in x.literal_value]):
return lambda x: None
# Capture the type strings
type_str = ', '.join([str(lv) for lv in x.literal_value])
def impl(x):
return "Item types: " + type_str
return impl
@njit
def mixed_type_list():
# a list of type [literal intp, unicode string, NumPy 1d array of float64]
x = [1, 'a', np.zeros(5)]
print("type strings:", bar(x)) # prints the type strings
# a list with the magic "NOP" string
x = [1, 'a', np.zeros(5), "NOP"]
print("NOP does nothing...!", bar(x))
mixed_type_list()
Following on from immutable heterogeneous lists, immutable heterogeneous string key dictionaries are also now supported. For example:
@njit
def mixed_value_type_str_key_dict():
# str -> mixed types, including array and typed dictionary!
a = {'a': 1, 'b': 'string', 'c': np.arange(5), 'd': {10:20, 30:40}}
print('getitem', a['d']) # getitem works
[print("key", k) for k in a.keys()] # keys() works
[print("value", v) for v in literal_unroll(a.values())] # as does values()
print('len', len(a)) # non-mutating call on the dictionary is ok
print("contains ", 'a' in a, 'z' in a) # and contains as it's read only
# it's slightly contrived, but .items() also works
for item in literal_unroll(a.items()):
k, v = item
print(k, "->", v)
mixed_value_type_str_key_dict()
and a more advanced example might be to use a dictionary to provide configuration
@njit
def my_function(data, config):
tmp = data / np.linalg.norm(data, ord=config['normalize'])
iv = config['initial_value']
for i in tmp:
iv += i
return iv
@njit
def config_example(data):
# pass a dictionary as configuration
config_a = {'normalize': None, 'initial_value': 5}
result_a = my_function(data, config_a)
print(result_a)
config_b = {'normalize': np.inf, 'initial_value': 10j}
result_b = my_function(data, config_b)
print(result_b)
config_example(np.arange(10.))
NOTE: this is an advanced feature predominantly for use by library authors. It permits dispatching on values recorded from the definition site of the container.
Locally defined homogeneous lists and string key dictionaries can now do initial value capturing (and type capturing in the case of dictionaries), this requires the use of literally
to force literal value dispatch. These types now have an .initial_value
attribute which contains any information about the values at the definition site, as directly discovered from the bytecode. This is best demonstrated by example:
def demo_iv(x):
pass
@overload(demo_iv)
def ol_demo_iv(x):
# if the initial_value is not present, request literal value dispatch
if x.initial_value is None:
return lambda x: literally(x)
else: # initial_value is present on the type
print("type of x: {}. Initial value {}".format(x, x.initial_value))
return lambda x: ...
@njit
def initial_value_capturing():
l = [1, 2, 3, 4] # initial value [1, 2, 3, 4]
l.append(5) # not part of the initial value
demo_iv(l)
initial_value_capturing()
the same works for dictionaries
@njit
def dict_initial_value_capturing():
d = {'a': 10, 'b': 20, 'c': 30} # initial value {'a': 10, 'b': 20, 'c': 30}
d['d'] = 40 # not part of the initial value
demo_iv(d)
dict_initial_value_capturing()
given this information is evidently available at compile time it's naturally possible to dispatch specialisations based on this information.
Numerous improvements were made to on-disk function caching support in 0.51, to ensure the best performance on Python < 3.8 make sure the pickle5
module is installed!
A long requested piece of functionality was added in 0.51, that of being able to cache functions that contain object mode blocks. For example, this is now cacheable:
import time
from numba import objmode
n = 100
@njit(cache=True) # request caching!
def foo(): # this is a nopython mode function
x = y = 0
for i in range(n):
x += np.sqrt(np.cos(n) ** 2 + np.sin(n) ** 2)
# but this block jumps into object mode j is defined in object mode,
# so we need to tell `nopython` mode its type so it can be used
# outside this block in nopython mode
with objmode(j='int64'):
time.sleep(0.05)
j = i + 10 # j is defined in object mode
y += j
return x, y
print(foo()) # worked with no warnings!
As a result of fixing caching of object mode blocks, it's now also possible to cache functions defined in closures:
# the specialiser, close over a jitted function argument,
# the inner function is compiled and cached!
def make_function(specialise_on_this_function):
@njit(cache=True)
def specialised(x):
return specialise_on_this_function(x)
return specialised
@njit(cache=True)
def f(x):
print("f(x)", x)
@njit(cache=True)
def g(x):
print("g(x)", x)
# these both cache miss as it had to compile it, but no complaints about doing the caching!
special_f = make_function(f)
special_f(10)
print(special_f.stats)
special_g = make_function(g)
special_g(20)
print(special_g.stats)
A very common question from users is:
What can I use as a mutable structure that's also pass-by-reference?
the answer is the new StructRef
type (warning: this is experimental!), documentation is here.
from numba.experimental import structref
# Define a StructRef.
# `structref.register` associates the type with the default data model.
# This will also install getters and setters to the fields of
# the StructRef.
@structref.register
class FruitType(types.StructRef):
def preprocess_fields(self, fields):
# This method is called by the type constructor for additional
# preprocessing on the fields.
# Here, we don't want the struct to take Literal types.
return tuple((name, types.unliteral(typ)) for name, typ in fields)
# Define a Python type that can be used as a proxy to the StructRef
# allocated inside Numba. Users can construct the StructRef via
# the constructor for this type in python code and jit-code.
class Fruit(structref.StructRefProxy):
def __new__(cls, kind, amount):
# Overriding the __new__ method is optional, doing so
# allows Python code to use keyword arguments,
# or add other customized behavior.
# The default __new__ takes `*args`.
# IMPORTANT: Users should not override __init__.
return structref.StructRefProxy.__new__(cls, kind, amount)
# By default, the proxy type does not reflect the attributes or
# methods to the Python side. It is up to users to define
# these. (This may be automated in the future.)
@property
def kind(self):
# To access a field, we can define a function that simply
# return the field in jit-code. This is to permit access
# to the data in the jit representation of the structure.
# The definition is shown later.
return Fruit_get_kind(self)
@property
def amount(self):
# The definition of is shown later.
return Fruit_get_amount(self)
@njit
def Fruit_get_kind(self):
# In jit-code, the StructRef's attribute is exposed via
# structref.register
return self.kind
@njit
def Fruit_get_amount(self):
return self.amount
# This associates the proxy with FruitType for the given set of
# fields. Notice how we are not contraining the type of each field.
# Field types remain generic.
structref.define_proxy(Fruit, FruitType, ["kind", "amount"])
from numba.core.extending import overload_method
# Use @overload_method to add a method for "eat"
@overload_method(FruitType, "eat")
def ol_eat(self, this_many):
def impl(self, this_many):
if self.amount >= this_many:
self.amount -= this_many
else:
raise ValueError("Insufficient quantity")
return impl
Use the above, and also demonstrate the new str(int)
support (implemented by @guilhermeleobas, with thanks!)...
@njit
def demo_struct_mutation():
fruit = Fruit("apple", 5)
print("Have " + str(fruit.amount) + "s " + fruit.kind + ".\n\nGoing to eat 3...")
fruit.eat(3)
print("Now have ", str(fruit.amount) + "s " + fruit.kind + ".\n\nGoing to eat 4 more...")
try:
fruit.eat(4)
except:
print("Ran out of " + fruit.kind + "s!")
return fruit
python_struct = demo_struct_mutation()
print("Object returned to Python: kind={}, amount={}".format(python_struct.kind,
python_struct.amount))
Finally, and with many thanks to contributions from the community, this release contains support for:
setitem
with literal string on a record array (by @luk-f-a).np.ndarray
construction from literal value (by @guilhermeleobas).np.positive
ufunc support (by @niteya-shah).minlength
kwarg support to np.bincount
(by @AndrewEckart).np.divmod
ufunc support (by @eric-wieser).a demonstration of these features...
# Define a record array for use in the demo
rec_array = np.array([1, 2], dtype=np.dtype([('e', np.int32), ('f', np.float64)], align=True))
@njit
def new_numpy_features(rec):
print("original record", rec)
print("setitem with literal string on record array")
for f in literal_unroll(('e', 'f')):
rec[0][f] = 10 * ord(f)
print("record updated", rec)
print("np.ndarray from literal", np.asarray("abc"), np.asarray(123))
print("np.positive(np.arange(10))",np.positive(np.arange(10)))
print("np.bincount with minlength", np.bincount(np.array([0, 1, 2, 1, 3, 2, 4]),
minlength=10))
print("np.divmod, multi-output ufunc!", np.divmod(np.arange(10), 2 ))
new_numpy_features(rec_array)