How to specify numpy arrays with numba?

Apologies if this has been posted elsewhere.
I am seeking some clarity as to how to proper specify numpy arrays so that numba can recognize them.
More specifically:

  • What is the correct way to specify them as part of a numba class?
  • In combination with Numpy array typing?

There are various versions floating around but I have yet to find some clarity on this.

As part of the specs I did the following:

import numpy as np
import numba as nb

('arr', nb.types.npytypes.Array(nb.types.f4, 2, 'C'))

in order to specify a 2D numpy array. Within the function definition (as a parameter) I used:

 def __init__(arr: np.array= np.zeros((1,1)),...):

Hey @MLLO ,
Numpy typing does not support array shape annotations, yet (see Numpy-Typing support for shapes Issue #16544). You can still use the concept of “Variadic Generics” (PEP 646) to annotate numpy array shapes if you like. I am not sure if my example is correct but an annotation could look similar like this:

from typing import TypeAlias, Tuple, Any
import numpy as np
Arr_Float64_2D: TypeAlias = np.ndarray[Tuple[Any, Any], np.float64]

If you want to annotate the constructor of your class than use “np.ndarray” instead of “np.array”.
Nevertheless, you will not need it for your jitclass to work. The jitclass infers the type of the argument from the spec variable that you pass to the decorator.
There are multiple ways to specify the array. Everything that works should be fine, I assume.

import numpy as np
import numba as nb
import numba.types as nbt

#@nb.experimental.jitclass([('arr', nbt.Array(nbt.f8, 2, 'A'))])
@nb.experimental.jitclass([('arr', nbt.f8[:, :])])
class Foo:
    # arr: nbt.f8[:, :]
    def __init__(self, arr: np.ndarray[np.float64]):
        self.arr = arr

arr = np.ones((2,2), dtype=np.float64)
foo = Foo(arr)
print(foo.arr)
# [[1. 1.]
#  [1. 1.]]

Excellent @Oyibo !

What you provided works. Many thanks!

I am still interested in finding where I can find more guidance for this because I do not think that the documentation is very clear on this. While I found out that my way of specifying the 2D array was incorrect, there are still different ways one can use to specify the same thing. E.g.

@nb.experimental.jitclass([('arr', nbt.f8[:, :])])

versus

@nb.experimental.jitclass([('arr', nbt.Array(nb.f8, 2, C))])

Is the way for specifying variables the same, regardless of whether one is doing so for a function or a class.

There are multiple ways to define arrays (1D, 2D, 3D… It is well described in the documentation.
https://numba.readthedocs.io/en/stable/reference/types.html#arrays
The way to specify variables in a function is similar to a jitclass but not the same.
Numba functions use lazy compilation, meaning you don’t need to provide detailed specifications for data types. You can simply apply the @jit decorator to optimize the function.
Numba jitclasses require more information about the data types of its fields, often provided in a specification (spec). Most of the time you need to specify the types of fields explicitly when defining the class, which helps Numba generate optimized code for the class.
There seems to be a mayor disadvantage using jitclasses.
Numba jitclasses cannot be compiled ahead of time or cached easily. While Numba functions can be compiled ahead of time (AOT) and cached for later use, the same isn’t available for jitclasses.
This means that when using jitclasses, the compilation happens at runtime, which might introduce some overhead compared to pre-compiled and cached functions. The lack of AOT support and caching for jitclasses can be a limitation when you need to optimize startup time or reuse the same compiled code.

1 Like

Thank you @Oyibo for your comments. They are very helpful. It is giving me some pause on whether I want to continue the effort of translating the code I had (originally in OO) into numba classes.

I will definitely try to make use of some of the classes I created but perhaps will switch to more functional programming to complete the job,

Thanks again!
M