[Numpy-discussion] Type annotation for Numpy arrays, accelerators and numpy.typing

PIERRE AUGIER pierre.augier at univ-grenoble-alpes.fr
Tue Feb 16 04:14:18 EST 2021


Hi,

When Numpy 1.20 was released, I discovered numpy.typing and its documentation https://numpy.org/doc/stable/reference/typing.html

I know that it is very new but I'm a bit lost. A good API to describe Array type would be useful not only for type checkers but also for Python accelerators using ndarrays (in particular Pythran, Numba, Cython, Transonic).

For Transonic, I'd like to be able to use internally numpy.typing to have a better implementation of what we need in transonic.typing (in particular compatible with type checkers like MyPy).

However, it seems that I can't do anything with what I see today in numpy.typing.

For Python-Numpy accelerators, we need to be able to define precise array types to limit the compilation time and give useful hints for optimizations (ndim, partial or full shape). We also need fused types.

What can be done with Transonic is described in these pages: https://transonic.readthedocs.io/en/latest/examples/type_hints.html and https://transonic.readthedocs.io/en/latest/generated/transonic.typing.html

I think it would be good to be able to do things like that with numpy.typing. It may be already possible but I can't find how in the doc.

I can give few examples here. First very simple:

from transonic import Array

Af3d = Array[float, "3d"]

# Note that this can also be written without Array just as
Af3d = "float[:,:,:]"

# same thing but only contiguous C ordered
Af3d = Array[float, "3d", "C"]

Note: being able to limit the compilation just for C-aligned arrays is very important since it can drastically decrease the compilation time/memory and that some numerical kernels are anyway written to be efficient only with C (or Fortran) ordered arrays.

# 2d color image
A_im = Array[np.int16, "[:,:,3]"]

Now, fused types. This example is taken from a real life case (https://foss.heptapod.net/fluiddyn/fluidsim/-/blob/branch/default/fluidsim/base/time_stepping/pseudo_spect.py) so it's really useful in practice.

from transonic import Type, NDim, Array, Union

N = NDim(2, 3, 4)
A = Array[np.complex128, N, "C"]
Am1 = Array[np.complex128, N - 1, "C"]

N123 = NDim(1, 2, 3)
A123c = Array[np.complex128, N123, "C"]
A123f = Array[np.float64, N123, "C"]

T = Type(np.float64, np.complex128)
A1 = Array[T, N, "C"]
A2 = Array[T, N - 1, "C"]
ArrayDiss = Union[A1, A2]

To summarize, type annotations are and will also be used for Python-Numpy accelerators. It would be good to also consider this application when designing numpy.typing.

Cheers,
Pierre


More information about the NumPy-Discussion mailing list