Supported NumPy features
One objective of Numba is having a seamless integration with NumPy. NumPy arrays provide an efficient storage method for homogeneous sets of data. NumPy dtypes provide type information useful when compiling, and the regular, structured storage of potentially large amounts of data in memory provides an ideal memory layout for code generation. Numba excels at generating code that executes on top of NumPy arrays.
NumPy support in Numba comes in many forms:
Numba understands calls to NumPy ufuncs and is able to generate equivalent native code for many of them.
NumPy arrays are directly supported in Numba. Access to NumPy arrays is very efficient, as indexing is lowered to direct memory accesses when possible.
Numba is able to generate ufuncs and gufuncs. This means that it is possible to implement ufuncs and gufuncs within Python, getting speeds comparable to that of ufuncs/gufuncs implemented in C extension modules using the NumPy C API.
The following sections focus on the NumPy features supported in nopython mode, unless otherwise stated.
Scalar types
Numba supports the following NumPy scalar types:
Integers: all integers of either signedness, and any width up to 64 bits
Booleans
Real numbers: single-precision (32-bit) and double-precision (64-bit) reals
Complex numbers: single-precision (2x32-bit) and double-precision (2x64-bit) complex numbers
Datetimes and timestamps: of any unit
Character sequences (but no operations are available on them)
Structured scalars: structured scalars made of any of the types above and arrays of the types above
The following scalar types and features are not supported:
Arbitrary Python objects
Half-precision and extended-precision real and complex numbers
Nested structured scalars the fields of structured scalars may not contain other structured scalars
The operations supported on NumPy scalars are almost the same as on the
equivalent built-in types such as int or float. You can use a type’s
constructor to convert from a different type or width. In addition you can use
the view(np.<dtype>) method to bitcast all int and float types
within the same width. However, you must define the scalar using a NumPy
constructor within a jitted function. For example, the following will work:
>>> import numpy as np
>>> from numba import njit
>>> @njit
... def bitcast():
... i = np.int64(-1)
... print(i.view(np.uint64))
...
>>> bitcast()
18446744073709551615
Whereas the following will not work:
>>> import numpy as np
>>> from numba import njit
>>> @njit
... def bitcast(i):
... print(i.view(np.uint64))
...
>>> bitcast(np.int64(-1))
---------------------------------------------------------------------------
TypingError Traceback (most recent call last)
...
TypingError: Failed in nopython mode pipeline (step: ensure IR is legal prior to lowering)
'view' can only be called on NumPy dtypes, try wrapping the variable with 'np.<dtype>()'
File "<ipython-input-3-fc40aaab84c4>", line 3:
def bitcast(i):
print(i.view(np.uint64))
Structured scalars support attribute getting and setting, as well as member lookup using constant strings. Strings stored in a local or global tuple are considered constant strings and can be used for member lookup.
import numpy as np
from numba import njit
arr = np.array([(1, 2)], dtype=[('a1', 'f8'), ('a2', 'f8')])
fields_gl = ('a1', 'a2')
@njit
def get_field_sum(rec):
fields_lc = ('a1', 'a2')
field_name1 = fields_lc[0]
field_name2 = fields_gl[1]
return rec[field_name1] + rec[field_name2]
get_field_sum(arr[0]) # returns 3
It is also possible to use local or global tuples together with literal_unroll:
import numpy as np
from numba import njit, literal_unroll
arr = np.array([(1, 2)], dtype=[('a1', 'f8'), ('a2', 'f8')])
fields_gl = ('a1', 'a2')
@njit
def get_field_sum(rec):
out = 0
for f in literal_unroll(fields_gl):
out += rec[f]
return out
get_field_sum(arr[0]) # returns 3
Record subtyping
Warning
This is an experimental feature.
Numba allows width subtyping of structured scalars.
For example, dtype([('a', 'f8'), ('b', 'i8')]) will be considered a subtype of dtype([('a', 'f8')], because
the second is a strict subset of the first, i.e. field a is of the same type and is in the same position in both
types. The subtyping relationship will matter in cases where compilation for a certain input is not allowed, but the
input is a subtype of another, allowed type.
import numpy as np
from numba import njit, typeof
from numba.core import types
record1 = np.array([1], dtype=[('a', 'f8')])[0]
record2 = np.array([(2,3)], dtype=[('a', 'f8'), ('b', 'f8')])[0]
@njit(types.float64(typeof(record1)))
def foo(rec):
return rec['a']
foo(record1)
foo(record2)
Without subtyping the last line would fail. With subtyping, no new compilation will be triggered, but the
compiled function for record1 will be used for record2.
See also
NumPy scalars reference.
Array types
NumPy arrays of any of the scalar types above are supported, regardless of the shape or layout.
Array access
Arrays support normal iteration. Full basic indexing and slicing is supported. A subset of advanced indexing is also supported: only one advanced index is allowed, and it has to be a one-dimensional array (it can be combined with an arbitrary number of basic indices as well).
See also
NumPy indexing reference.
Structured array access
Numba presently supports accessing fields of individual elements in structured arrays by attribute as well as by getting and setting. This goes slightly beyond the NumPy API, which only allows accessing fields by getting and setting. For example:
from numba import njit
import numpy as np
record_type = np.dtype([("ival", np.int32), ("fval", np.float64)], align=True)
def f(rec):
value = 2.5
rec[0].ival = int(value)
rec[0].fval = value
return rec
arr = np.ones(1, dtype=record_type)
cfunc = njit(f)
# Works
print(cfunc(arr))
# Does not work
print(f(arr))
The above code results in the output:
[(2, 2.5)]
Traceback (most recent call last):
File "repro.py", line 22, in <module>
print(f(arr))
File "repro.py", line 9, in f
rec[0].ival = int(value)
AttributeError: 'numpy.void' object has no attribute 'ival'
The Numba-compiled version of the function executes, but the pure Python version raises an error because of the unsupported use of attribute access.
Note
This behavior will eventually be deprecated and removed.
Attributes
The following attributes of NumPy arrays are supported:
The flags object
The object returned by the flags attribute supports
the contiguous, c_contiguous and f_contiguous attributes.
The flat object
The object returned by the flat attribute supports
iteration and indexing, but be careful: indexing is very slow on
non-C-contiguous arrays.
The real and imag attributes
NumPy supports these attributes regardless of the dtype but Numba chooses to
limit their support to avoid potential user error. For numeric dtypes,
Numba follows NumPy’s behavior. The real attribute
returns a view of the real part of the complex array and it behaves as an identity
function for other numeric dtypes. The imag attribute
returns a view of the imaginary part of the complex array and it returns a zero
array with the same shape and dtype for other numeric dtypes. For non-numeric
dtypes, including all structured/record dtypes, using these attributes will
result in a compile-time (TypingError) error. This behavior differs from
NumPy’s but it is chosen to avoid the potential confusion with field names that
overlap these attributes.
Calculation
The following methods of NumPy arrays are supported in their basic form (without any optional arguments):
The corresponding top-level NumPy functions (such as numpy.prod())
are similarly supported.
Other methods
The following methods of NumPy arrays are supported:
argmax()(axiskeyword argument supported).argmin()(axiskeyword argument supported).argsort()(kindkey word argument supported for values'quicksort'and'mergesort')astype()(only the 1-argument form)copy()(without arguments)dot()(only the 1-argument form)flatten()(no order argument; ‘C’ order only)item()(without arguments)itemset()(only the 1-argument form)ptp()(without arguments)ravel()(no order argument; ‘C’ order only)repeat()(no axis argument)reshape()(only the 1-argument form)sort()(without arguments)sum()(with or without theaxisand/ordtypearguments.)axisonly supportsintegervalues.If the
axisargument is a compile-time constant, all valid values are supported. An out-of-range value will result in aLoweringErrorat compile-time.If the
axisargument is not a compile-time constant, only values from 0 to 3 are supported. An out-of-range value will result in a runtime exception.All numeric
dtypesare supported in thedtypeparameter.timedeltaarrays can be used as input arrays buttimedeltais not supported asdtypeparameter.When a
dtypeis given, it determines the type of the internal accumulator. When it is not, the selection is made automatically based on the input array’sdtype, mostly following the same rules as NumPy. However, on 64-bit Windows, Numba uses a 64-bit accumulator for integer inputs (int64forint32inputs anduint64foruint32inputs), while NumPy would use a 32-bit accumulator in those cases.
view()(only the 1-argument form)
Where applicable, the corresponding top-level NumPy functions (such as
numpy.argmax()) are similarly supported.
Warning
Sorting may be slightly slower than NumPy’s implementation.
Functions
Linear algebra
Basic linear algebra is supported on 1-D and 2-D contiguous arrays of floating-point and complex numbers:
numpy.kron()(‘C’ and ‘F’ order only)numpy.trace()(only the first argument).On Python 3.5 and above, the matrix multiplication operator from PEP 465 (i.e.
a @ bwhereaandbare 1-D or 2-D arrays).numpy.linalg.cond()(only non string values inp).numpy.linalg.eig()(only running with data that does not cause a domain change is supported e.g. real input -> real output, complex input -> complex output).numpy.linalg.eigh()(only the first argument).numpy.linalg.eigvals()(only running with data that does not cause a domain change is supported e.g. real input -> real output, complex input -> complex output).numpy.linalg.eigvalsh()(only the first argument).numpy.linalg.norm()(only the 2 first arguments and only non string values inord).numpy.linalg.qr()(only the first argument).numpy.linalg.svd()(only the 2 first arguments).
Note
The implementation of these functions needs SciPy to be installed.
Reductions
The following reduction functions are supported:
numpy.diff()(only the 2 first arguments)numpy.median()(only the first argument)numpy.nancumprod()(only the first argument)numpy.nancumsum()(only the first argument)numpy.nanmax()(only the first argument)numpy.nanmean()(only the first argument)numpy.nanmedian()(only the first argument)numpy.nanmin()(only the first argument)numpy.nanpercentile()(only the 2 first arguments, complex dtypes unsupported)numpy.nanquantile()(only the 2 first arguments, complex dtypes unsupported)numpy.nanprod()(only the first argument)numpy.nanstd()(only the first argument)numpy.nansum()(only the first argument)numpy.nanvar()(only the first argument)numpy.percentile()(only the 2 first arguments, complex dtypes unsupported)numpy.quantile()(only the 2 first arguments, complex dtypes unsupported)
Other functions
The following top-level functions are supported:
numpy.argsort()(kindkey word argument supported for values'quicksort'and'mergesort')numpy.array()(only the 2 first arguments)numpy.asarray()(only the 2 first arguments)numpy.asarray_chkfinite()(only the 2 first arguments)numpy.asfortranarray()(only the first argument)numpy.broadcast_to()(only the 2 first arguments)numpy.broadcast_arrays()(only the first argument)numpy.convolve()(only the 2 first arguments)numpy.copy()(only the first argument)numpy.corrcoef()(only the 3 first arguments, requires SciPy)numpy.correlate()(only the 2 first arguments)numpy.count_nonzero()(axis only supports scalar values)numpy.cov()(only the 5 first arguments)numpy.cross()(only the 2 first arguments; at least one of the input arrays should haveshape[-1] == 3)If
shape[-1] == 2for both inputs, please replace yournumpy.cross()call withnumba.np.extensions.cross2d().
numpy.delete()(only the 2 first arguments)numpy.dtype()(only the first argument)numpy.empty()(only the 2 first arguments)numpy.empty_like()(only the 2 first arguments)numpy.flatten()(no order argument; ‘C’ order only)numpy.flip()(no axis argument)numpy.frombuffer()(only the 2 first arguments)numpy.full()(only the 3 first arguments)numpy.full_like()(only the 3 first arguments)numpy.histogram()(only the 3 first arguments)numpy.interp()(only the 3 first arguments)numpy.intersect1d()(only first 2 arguments, ar1 and ar2)numpy.linspace()(only the 3-argument form)numpy.logspace()(only the 3 first arguments)numpy.nditer(only the first argument)numpy.ones()(only the 2 first arguments)numpy.ones_like()(only the 2 first arguments)numpy.partition()(only the 2 first arguments)numpy.ptp()(only the first argument)numpy.ravel()(no order argument; ‘C’ order only)numpy.repeat()(no axis argument)numpy.reshape()(no order argument; ‘C’ order only)numpy.roll()(only the 2 first arguments; second argumentshiftmust be an integer)numpy.rot90()(only the 2 first arguments)numpy.round_()numpy.searchsorted()(only the 3 first arguments)numpy.select()(only using homogeneous lists or tuples for the first two arguments, condlist and choicelist). Additionally, these two arguments can only contain arrays (unlike NumPy that also accepts tuples).numpy.sort()(no optional arguments, quicksort accepts multi-dimensional array and sorts its last axis).numpy.take()(only the 2 first arguments)numpy.take_along_axis()(the axis argument must be a literal value)numpy.trapz()(only the 3 first arguments)numpy.tri()(only the 3 first arguments; third argumentkmust be an integer)numpy.tril()(second argumentkmust be an integer)numpy.tril_indices()(all arguments must be integer)numpy.tril_indices_from()(second argumentkmust be an integer)numpy.triu()(second argumentkmust be an integer)numpy.triu_indices()(all arguments must be integer)numpy.triu_indices_from()(second argumentkmust be an integer)numpy.unique()(only the first argument)numpy.zeros()(only the 2 first arguments)numpy.zeros_like()(only the 2 first arguments)
The following constructors are supported, both with a numeric input (to construct a scalar) or a sequence (to construct an array):
numpy.complex64numpy.complex128numpy.float32numpy.float64numpy.int8numpy.int16numpy.int32numpy.int64numpy.intpnumpy.uint8numpy.uint16numpy.uint32numpy.uint64numpy.uintp
The following machine parameter classes are supported, with all purely numerical attributes:
numpy.finfo(macharattribute not supported)numpy.MachAr(with no arguments to the constructor)
Literal arrays
Neither Python nor Numba has actual array literals, but you can construct
arbitrary arrays by calling numpy.array() on a nested tuple:
a = numpy.array(((a, b, c), (d, e, f)))
(nested lists are not yet supported by Numba)
Modules
random
Generator Objects
Numba supports numpy.random.Generator() objects. As of version 0.56, users can pass
individual NumPy Generator objects into Numba functions and use their
methods inside the functions. The same algorithms are used as NumPy for
random number generation hence maintaining parity between the random
number generated using NumPy and Numba under identical arguments
(also the same documentation notes as NumPy Generator methods apply).
The current Numba support for Generator is not thread-safe, hence we
do not recommend using Generator methods in methods with parallel
execution logic.
Note
NumPy’s Generator objects rely on BitGenerator to manage state
and generate the random bits, which are then transformed into random
values from useful distributions. Numba will unbox the Generator objects
and will maintain a reference to the underlying BitGenerator objects using NumPy’s
ctypes interface bindings. Hence Generator objects can cross the JIT boundary
and their functions be used within Numba-Jit code. Note that since only references
to BitGenerator objects are maintained, any change to the state of a particular
Generator object outside Numba code would affect the state of Generator
inside the Numba code.
x = np.random.default_rng(1)
y = np.random.default_rng(1)
size = 10
@numba.njit
def do_stuff(gen):
return gen.random(size=int(size / 2))
original = x.random(size=size)
# [0.51182162 0.9504637 0.14415961 0.94864945 0.31183145
# 0.42332645 0.82770259 0.40919914 0.54959369 0.02755911]
numba_func_res = do_stuff(y)
# [0.51182162 0.9504637 0.14415961 0.94864945 0.31183145]
after_numba = y.random(size=int(size / 2))
# [0.42332645 0.82770259 0.40919914 0.54959369 0.02755911]
The following Generator methods are supported:
numpy.random.Generator().random()
RandomState and legacy Random number generation
Numba supports top-level functions from the numpy.random module, but does not allow you to create individual RandomState instances. The same algorithms are used as for the standard random module (and therefore the same notes apply), but with an independent internal state: seeding or drawing numbers from one generator won’t affect the other.
The following functions are supported.
Initialization
numpy.random.seed(): with an integer argument only
Warning
Calling numpy.random.seed() from interpreted code (including from object mode
code) will seed the NumPy random generator, not the Numba random generator.
To seed the Numba random generator, see the example below.
from numba import njit
import numpy as np
@njit
def seed(a):
np.random.seed(a)
@njit
def rand():
return np.random.rand()
# Incorrect seeding
np.random.seed(1234)
print(rand())
np.random.seed(1234)
print(rand())
# Correct seeding
seed(1234)
print(rand())
seed(1234)
print(rand())
Simple random data
Permutations
numpy.random.choice(): the optional p argument (probabilities array) is not supportednumpy.random.shuffle(): the sequence argument must be a one-dimension NumPy array or buffer-providing object (such as abytearrayorarray.array)
Distributions
The following functions support all arguments.
Note
Calling numpy.random.seed() from non-Numba code (or from
object mode code) will seed the NumPy random generator, not the
Numba random generator.
Note
Since version 0.28.0, the generator is thread-safe and fork-safe. Each thread and each process will produce independent streams of random numbers.
stride_tricks
The following function from the numpy.lib.stride_tricks module
is supported:
as_strided()(the strides argument is mandatory, the subok argument is not supported)
Standard ufuncs
One objective of Numba is having all the standard ufuncs in NumPy understood by Numba. When a supported ufunc is found when compiling a function, Numba maps the ufunc to equivalent native code. This allows the use of those ufuncs in Numba code that gets compiled in nopython mode.
Limitations
Right now, only a selection of the standard ufuncs work in nopython mode. Following is a list of the different standard ufuncs that Numba is aware of, sorted in the same way as in the NumPy documentation.
Math operations
UFUNC |
MODE |
|
|---|---|---|
name |
object mode |
nopython mode |
add |
Yes |
Yes |
subtract |
Yes |
Yes |
multiply |
Yes |
Yes |
divide |
Yes |
Yes |
logaddexp |
Yes |
Yes |
logaddexp2 |
Yes |
Yes |
true_divide |
Yes |
Yes |
floor_divide |
Yes |
Yes |
negative |
Yes |
Yes |
power |
Yes |
Yes |
float_power |
Yes |
Yes |
remainder |
Yes |
Yes |
mod |
Yes |
Yes |
fmod |
Yes |
Yes |
divmod (*) |
Yes |
Yes |
abs |
Yes |
Yes |
absolute |
Yes |
Yes |
fabs |
Yes |
Yes |
rint |
Yes |
Yes |
sign |
Yes |
Yes |
conj |
Yes |
Yes |
exp |
Yes |
Yes |
exp2 |
Yes |
Yes |
log |
Yes |
Yes |
log2 |
Yes |
Yes |
log10 |
Yes |
Yes |
expm1 |
Yes |
Yes |
log1p |
Yes |
Yes |
sqrt |
Yes |
Yes |
square |
Yes |
Yes |
cbrt |
Yes |
Yes |
reciprocal |
Yes |
Yes |
conjugate |
Yes |
Yes |
gcd |
Yes |
Yes |
lcm |
Yes |
Yes |
(*) not supported on timedelta types
Trigonometric functions
UFUNC |
MODE |
|
|---|---|---|
name |
object mode |
nopython mode |
sin |
Yes |
Yes |
cos |
Yes |
Yes |
tan |
Yes |
Yes |
arcsin |
Yes |
Yes |
arccos |
Yes |
Yes |
arctan |
Yes |
Yes |
arctan2 |
Yes |
Yes |
hypot |
Yes |
Yes |
sinh |
Yes |
Yes |
cosh |
Yes |
Yes |
tanh |
Yes |
Yes |
arcsinh |
Yes |
Yes |
arccosh |
Yes |
Yes |
arctanh |
Yes |
Yes |
deg2rad |
Yes |
Yes |
rad2deg |
Yes |
Yes |
degrees |
Yes |
Yes |
radians |
Yes |
Yes |
Bit-twiddling functions
UFUNC |
MODE |
|
|---|---|---|
name |
object mode |
nopython mode |
bitwise_and |
Yes |
Yes |
bitwise_or |
Yes |
Yes |
bitwise_xor |
Yes |
Yes |
bitwise_not |
Yes |
Yes |
invert |
Yes |
Yes |
left_shift |
Yes |
Yes |
right_shift |
Yes |
Yes |
Comparison functions
UFUNC |
MODE |
|
|---|---|---|
name |
object mode |
nopython mode |
greater |
Yes |
Yes |
greater_equal |
Yes |
Yes |
less |
Yes |
Yes |
less_equal |
Yes |
Yes |
not_equal |
Yes |
Yes |
equal |
Yes |
Yes |
logical_and |
Yes |
Yes |
logical_or |
Yes |
Yes |
logical_xor |
Yes |
Yes |
logical_not |
Yes |
Yes |
maximum |
Yes |
Yes |
minimum |
Yes |
Yes |
fmax |
Yes |
Yes |
fmin |
Yes |
Yes |
Floating functions
UFUNC |
MODE |
|
|---|---|---|
name |
object mode |
nopython mode |
isfinite |
Yes |
Yes |
isinf |
Yes |
Yes |
isnan |
Yes |
Yes |
signbit |
Yes |
Yes |
copysign |
Yes |
Yes |
nextafter |
Yes |
Yes |
modf |
Yes |
No |
ldexp |
Yes (*) |
Yes |
frexp |
Yes |
No |
floor |
Yes |
Yes |
ceil |
Yes |
Yes |
trunc |
Yes |
Yes |
spacing |
Yes |
Yes |
(*) not supported on windows 32 bit
Datetime functions
UFUNC |
MODE |
|
|---|---|---|
name |
object mode |
nopython mode |
isnat |
Yes |
Yes |