python / python/typing

Typing for multi-dimensional arrays

Open
#513 21 comments 74 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

topic: feature
Dominant language
Python
Stars
1.8k
Forks
302
Avg merge
23h
Merged PRs (30d)
8

Description

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 (https://github.com/numpy/numpy/issues/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 (https://github.com/python/typing/issues/478), so we can type check operations like array.sum(axis=0).
  • Variadic generics (https://github.com/python/typing/issues/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?

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the issue, the linked NumPy issue and numpy_stubs repository, then review the linked document on array shape typing. Compare the proposals around literal values, variadic generics, dimension identity and annotation syntax; the issue does not define a concrete implementation target or a test that would establish completion.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
tooling
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.