Typing necessary for static Python-Numpy compilers like Pythran and Cython

[paugier]

Static Python-Numpy compilers like Pythran and Cython (and others) need type information.

For Transonic (a tools to use Python-Numpy compilers, see https://transonic.readthedocs.io), I had to implement a simple system to be able to give enough information in Python about the types supported by functions.

I give here an example of what is possible (more complex things are sometimes needed but this is the minimum, fused data-types and "fused" number of dimensions of the array):

import numpy as np

from transonic import Array, NDim, Type, boost

T = Type(int, np.complex128)
N = NDim(1, 3)

A = Array[T, N]
A1 = Array[np.float32, N + 1]

@boost
def compute(a: A, b: A, c: T, d: A1, e: str):
    ...

It would be great if array-api-typing could be good enough to be able to express such things in a quite simple way.

Note that this leads to the following possible Pythran signatures:

export compute(complex128[:, :, :], complex128[:, :, :], complex128, float32[:, :, :, :], str)
export compute(complex128[:], complex128[:], complex128, float32[:, :], str)
export compute(int[:, :, :], int[:, :, :], int, float32[:, :, :, :], str)
export compute(int[:], int[:], int, float32[:, :], str)

and Cython "signatures":

import cython

import numpy as np
cimport numpy as np

ctypedef fused __compute__Array_TypeIint_complex128I_NDimI1_3I:
   np.ndarray[np.complex128_t, ndim=1]
   np.ndarray[np.complex128_t, ndim=3]
   np.ndarray[np.int32_t, ndim=1]
   np.ndarray[np.int32_t, ndim=3]

ctypedef fused __compute__Array_float32_NDimI1_3Ip1:
   np.ndarray[np.float32_t, ndim=2]
   np.ndarray[np.float32_t, ndim=4]

ctypedef fused __compute__TypeIint_complex128I:
   cython.int
   np.complex128_t

cpdef compute(__compute__Array_TypeIint_complex128I_NDimI1_3I a, __compute__Array_TypeIint_complex128I_NDimI1_3I b, __compute__TypeIint_complex128I c, __compute__Array_float32_NDimI1_3Ip1 d, cython.str e)

Is there already enough in array-api-typing to be able to express such things? If not, would it be possible in future?

[jorenham]

This won't be possible to express in a way that's compatible with the Python typing spec, I'm afraid.

[paugier]

Is it at least possible to express that an argument is an array of a particular dtype and a particular dimension?

[jorenham]

Is it at least possible to express that an argument is an array of a particular dtype and a particular dimension?

This is possible in NumPy for dtypes already, and we're working hard on implementing "rank-typing", i.e. statically describing the no. of dimensions. It might sound easy, but in reality it's one of the most difficult things I've ever worked on in my career.

And if it's possible to do in NumPy, then it should also be possible to do in general. But getting there is be no means easy.

[lucascolley]

x-ref numpy/numtype#539