Typing for multi-dimensional arrays

[shoyer]

I'd like to open a discussion about typing for multi-dimensional arrays in general, and more specifically for NumPy. We have already been discussing this over in the NumPy issue tracker (numpy/numpy#7370) and recently opened a new repository to start writing type stubs (https://github.com/numpy/numpy_stubs).

To help guide discussion, I wrote a document outlining ideas for array shape typing.

To summarize:

• We would like to be able to type-check both data types (e.g., float64) and shapes (e.g., a 3x4 array) for multi-dimensional arrays.
• There are many uses cases where support for checks using dimension identity would be valuable, e.g., to indicate that a function transforms an array with shape (N, M) to shape (N,) for arbitrary integers N and M. These dimension variables look very similar to TypeVar, if TypeVar supported integers as types.
• A notion of "zero or more additional dimensions" would also be quite valuable, and is a core part of the type for many NumPy operations (generalized ufuncs). This might be naturally written with Ellipsis, e.g., (...., N) for an array with a last dimension of length N and any number of proceeding dimensions. There are particular rules (broadcasting) that should be enforced for matching multiple arguments with variable numbers of dimensions.

This will likely require some new typing features (as well as type-checker support). Notably:

• Support for literal values (Check that literal strings/int/float belong to /is excluded from a set/range of values #478), so we can type check operations like array.sum(axis=0).
• Variadic generics (Allow variadic generics #193), we can write types like NDArray[N] and NDArray[N, M].
• Some sort of support for dimension identity in shapes (e.g., integer types, or DimensionVar as described in my doc).
• Standard syntax for writing array dtype/shape annotations: what should these look like?

[ilevkivskyi]

It looks like the proposal of integer generics is also relevant here python/mypy#3345 (it looks almost identical to what you call DimensionVar).

In general, I am very supportive of this project (I have heard many times that static typing would be very helpful for data science, numerics and related fields, but current support in mypy and PEP 484 is very limited). The main obstacle however is the size of this project (it may require its own PEP). I will read your document (thanks for writing it), but already now it seems to me that it may make sense to start from features that will be useful in general (i.e. also outside of numeric stack) such as literal types and variadic generics.

Also tagging @JukkaL here just in case.

[shoyer]

The main obstacle however is the size of this project (it may require its own PEP).

Yes, I expect a PEP will be necessary, especially if we want to standardize base types for typing multi-dimensional arrays in the typing module.

it seems to me that it may make sense to start from features that will be useful in general (i.e. also outside of numeric stack) such as literal types and variadic generics.

Indeed, this is probably the best place where the broader typing community can help.

[shoyer]

I've opened a sub-issue for discussing syntax for array typing: #516

[ilevkivskyi]

Some update on the issue:

Our (mypy core team) previous schedule for working on this was Q4 2018. However, we decided that some type system features (such as literal types and variadic generics) needed to efficiently support NumPy will be also useful in general, so we decided to implement the general support for such features first. Literal types are almost already there, and variadic generics are going to be added in coming months. After that we will start working on dedicated NumPy support (around Q2), sorry for a delay.

[ilevkivskyi]

Sorry, I forgot to post notes from the latest Python typing meetup on numeric stack typing here. Here they are

[vsiles]

Are you specifically looking at numpy, or at the machine learning echosystem with numpy/pytorch/... ?
I found today that they are quite heterogeneous:

>>> x = torch.zeros([4], dtype=torch.int8)
>>> y = torch.zeros([4], dtype=torch.float32)
>>> torch.add(x, y)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: expected type torch.FloatTensor but got torch.CharTensor
>>> xx = numpy.array([4], dtype=numpy.int8)
>>> yy = numpy.array([4], dtype=numpy.float32)
>>> xx + yy
array([8.], dtype=float32)

Pytorch doesn't seems to do auto cast when types are different whereas Numpy is doing some upcast (see https://stackoverflow.com/questions/56022497/numpy-pytorch-dtype-conversion-compatibility/56022918?noredirect=1#comment98689941_56022918)

[ilevkivskyi]

Are you specifically looking at numpy, or at the machine learning echosystem with numpy/pytorch/... ?

At all of them. Dimensionality/shape will be an additional abstraction orthogonal to container type and element type.

[vsiles]

Sorry I wasn't clear, I wanted to ask for the numerical stack part specifically. Do we have a current target in numpy / pytorch / tensorflow that would focus most of the effort are are people looking to their favorite flavor (which seems incompatible with each other)