Typing support for shapes

See how contracts package are trying to provide support for something similar to shapes. They are extending annotations in a different way than just standard typing and maybe something like that could be also done. So instead of providing specific extension (PEP) for typing to allow thing like shapes, it might be maybe more useful to determine a syntax for general constraints on types and use that, in addition to standard types through typing.

This will take quite some work here and in mypy. At a discussion held several weeks ago, the agreement was on the plan in the README. So while this will happen, adding support for checking shapes will not happen soon. There is much to decide, such as how to define syntax related to declaring array shape, and how an actual type checking implementation should work.

@shoyer

@shoyer: I'm not sure where this should be noted, but the framing of numpy/numpy-stubs#6 and python/typing#516 seems to indicate that you're favoring putting the full shape in the type specification: that is, for a two-dimensional array of size (M,N), the proposal seems to be that the type annotation would convey both the value of M and N.

I think we should consider a more modest version in which only the number of dimensions is tracked via the type, but not the shape itself. That is, instead of ndarray[Shape[100, 200], Dtype] for an array of shape=(100,200), we go for ndarray[2, Dtype], to indicate only the fact that the array has ndim=2.

Tracking only the rank (as opposed to the full shape) of the array through the type system is, for what it's worth, how all of the C++ tensor libraries work to the best of my knowledge. It's more consistent with the rest of the PEP484 container types as well, in the sense that the type List[int] doesn't track the number of elements.

@rmcgibbo

@rmcgibbo Yes, this needs justification.

My reasoning for favoring specifying shapes in types comes down to:

Rank and shape are not mutually exclusive: if we do shape, we get rank as well. We could even have an explicit Rank type, which we tell type checkers should be treated as equivalent to anonymous dimensions, e.g., ndarray[Rank[2]] could be the same as ndarray[Shape[:, :]].
For code directly using NumPy, rank checks would not be anywhere near as valuable as shape checks. The problem is that nearly NumPy operation is valid on inputs with an arbitrary number of dimensions (even between arguments, with broadcasting). Giving up on shape checks would means we couldn't type ufuncs in a useful way, or catch the common errors with broadcasting. This would basically mean reserving shape/rank checks for the (many) applications and libraries that only use on particular ranks (e.g., vectors and matrices).

With regards to C++ libraries, Eigen at least does allow for either runtime and statically defined shapes. That said, I certainly would not bother with static/compile-time shapes on their own, because it is indeed rare that you know the exact shape of arrays before running a program. But with generics, we have the possibility of doing much more sophisticated pattern matching. This isn't feasible in C++, where you would need to use templates for every possible dimension size. In contrast, we can probably do something more sophisticated in Python to allow for typing shapes without enumerating all of them (similar to what you see in Haskell).

Shape is clearly more expressive than rank. So in that sense, I agree with both of your bulleted points above (not mutually exclusive, and static shape checks, if they worked, would provide more value to users). There are two factors on the other side that come to mind.

In almost any real code, it will be impossible for mypy to statically infer the shapes. For example, take the following:

M = 5
N = 10
arr = np.zeros((M, N))
reveal_type(arr)

Mypy tracks the type of M and N, but not their values. It can infer the rank of arr, but not the full shape. Want to do something more complicated?

def compute_length() -> int:
    return <something>
def compute_width() > int:
    return <something>
arr = np.zeros((compute_length(), compute_width()))
reveal_type(arr)

Not going to work.

If the shape is to be tracked explicitly, for example, to the best of my understanding there is no way to express even the simplest numpy function signatures with the current generation of PEP484/mypy. For example, take

# requires a type-level addition function to express the return type.
def concatenate_1d_arrays(arr1 : ndarray[D, Shape[M]], arr2 : ndarray[D, Shape[N]] -> ?

Take another look at my doc on shape typing. You're right that functions like concatenate are hard, but I would claim that even mere matching of integer sizes (without arithmetic) along with dimension identity would be enough to be valuable:

You could catch broadcasting errors when using ufuncs.
You could catch errors with regards to which dimension is removed by aggregation or indexing
You could verify that the shapes of non-concatenated axes of arrays in concatenate match, and are propagated to the result, e.g., stack_matrices(x: ndarray[Shape[:, N], y: ndarray[Shape[:, N]]) -> ndarray[Shape[:, N]].

In almost any real code, it will be impossible for mypy to statically infer the shapes. For example, take the following:

M = 5
N = 10
arr = np.zeros((M, N))
reveal_type(arr)

I'm not certain that mypy will be unable to infer shapes in this case. If typing adds support for literals (which is almost certainly necessary regardless for NumPy), then I imagine it will also support literals in variables, in the same way that it supports type aliases.

My suspicions are that shape tracking is (1) going to add a small marginal value, since it will be rarely be possible to statically infer them (e.g. my two examples above) in real world code, whereas it will often be possible to statically infer rank in real world code; and (2) going to add a ton of complexity. That suggests to me it's not the best place to start for the "1.0".

If there's a way to design the syntax so that it is initially rank-only, build a viable package around that, and add shape information afterwards in a second stage, then that would be my preference.

Maybe for v1.0 we could have ndarray[2, Dtype] be an alias for ndarray[Shape[Any, Any], Dtype], ignoring the length of the shape but keeping the ndim in mind. This would probably be a good-ish idea to keep forward compatibility.

Edit: It might also be a good idea to have a warning if the shape contains anything other than Any (or similar), for now.

Yes, I agree that it could make sense to add shape support incrementally. We'll probably still need upstream work in typing even to make that work well, though (e.g., to support integer values in types).

It seems like there is some interesting work being done on variadic generics:
https://mail.python.org/archives/list/typing-sig@python.org/thread/SQVTQYWIOI4TIO7NNBTFFWFMSMS2TA4J/

Depending on how this progresses from here it is possible we could end up with either
ndarray[Dtype, [i, j]] (or ndarray[[i, j], Dtype]) without the need for a dedicated Shape object.

It seems like we now have a (draft) pep for variadic generics,
i.e. generics with not one but an arbitrary number of variables (like typing.Tuple): PEP 646.

I have the impression, from the discussion above, that only "an array whose shape is a typecheck-time constant" and "an array whose rank is a typecheck-time constant but all of whose actual dimensions can vary at runtime" are currently being considered for shape typing.

Yea that's about right. Statically annotating the rank of an array is the best we can do within the limitations of Python type system. Currently, it's not possible to statically describe the size of individual axes, neither as constant, nor symbolically. And, unfortunately, we won't be able to use Literal in shape types, as explained in numpy/numtype#578.

However, in my experience, being able to annotate the rank of an array can already be very helpful.

And although it might seem like a very easy thing to implement, I can guarantee you that it's quite the opposite. For example, the broadcasting rules require something akin to a max(rank1, rank2, ...) function. But that doesn't exist, so we have to emulate it, and I'm sure you can imagine that's not an easy thing to do, to say the least.

I can definitely see the difficulties. It's just unfortunate, since I think this means type checking won't be able to help me with the most common numpy-related bugs in the kind of code I tend to write.

I wonder if a runtime type checking solution might be worth pursuing? An example of this sort of approach is the jaxtyping package, which does runtime shape & dtype checks for JAX arrays based on type annotations.

As far as I'm aware, jaxtyping is purely static. So you won't get IDE support, and will still have to run the code to see if you did something wrong. But if that suits your use-case, then I'm not gonna stop you. But personally, I'm not a big fan of their string-based syntax, and don't see the need for officially supporting something like jaxtyping in NumPy.

will still have to run the code to see if you did something wrong

Exactly - this is why I emphasized runtime type checking.

don't see the need for officially supporting something like jaxtyping in NumPy

I agree that supporting something like jaxtyping in NumPy is unlikely. But I bring it up because it points to the possiblity of interested parties striking out on their own to support runtime shape checks (similar to how jaxtyping is a community package not supported by the core JAX team).

Sorry for self promo, following this issue bc i maintain numpydantic for runtime shape validation. Doesn't have to be used with pydantic, supports shape specs, isinstance checks, etc.

Sorry for self promo, following this issue bc i maintain numpydantic for runtime shape validation. Doesn't have to be used with pydantic, supports shape specs, isinstance checks, etc.

Does numpydantic support using array types from numpy.typing, so that it can also be used for static type-checking?

Does numpydantic support using array types from numpy.typing, so that it can also be used for static type-checking?

Unfortunately, Pydantic isn't able to support third-party types to be used as is without either subclassing the type or defining Annotated metadata. We have plans for a new API that would allow this though.

but numpydantic does (and can be used outside of pydantic models) - that's sorta it's whole deal :)

from pydantic import BaseModel
from numpydantic import NDArray, Shape
import numpy as np

class MyModel(BaseModel):
    array: NDArray[Shape["1, 2, 3"], int]

# works
MyModel(array=np.zeros((1,2,3), dtype=int))
# fails validation bc shape mismatch
MyModel(array=np.zeros((4,5,6), dtype=int))

numpydantic does the best it can with static type checking given its purpose as a mostly runtime tool that is an interface to lots of different array backends. It also inherited nptyping's syntax which linters don't like (this will be fixed in ~a month or so). static shape checking is pretty hard (and likely to be impossible to do fully correctly, i would be worried about the potential for the false sense of security from operations that mutate shape without things that would be obvious to a static checker plugin like slicing or using reshape). I plan to write a mypy plugin that does best effort but that's tbd/help wanted.

but it does do duck typing for things that can be coerced into arrays, and it supports isinstance checks, so you could do integration using mypy's type narrowing/use typing.TypeGuard or TypeIs, and i would be happy to add whatever is needed for whatever compatibility with numpy.typing you want :)

from typing import Any, TypeGuard
from numpydantic import NDArray, Shape
import numpy as np

# type narrowing
some_array = np.zeros((1,2,3))
if isinstance(some_array, NDArray[Shape["1,2,3"], Any]):
    pass

# type guard
MyShapeArray = NDArray[Shape["1, 2, 3"], Any]
def is_my_shape(array: np.ndarray) -> TypeGuard[MyShapeArray]:
    return isinstance(array, MyShapeArray)

# use as a callable for a cheap assert - array will be passed through if shape/dtype validate
# otherwise, raise validation error
array = MyShapeArray(array)

It looks to me like numpydantic.NDArray isn't compatible with numpy.typing.NDArray. And looking at the implementation of it, I can't imagine that static type-checkers will accept assignments like a: numpydantic.NDArray[Shape["1"], Any] = np.empty((1,)), or am I missing something?

Well, I suppose I'm to blame for the red herring, but we're deviating from the subject here. I'll raise an issue regarding static typing support, but I'll mark your comment as off-topic.

For anyone interested; I'm tracking the shape-typing progress in https://gist.github.com/jorenham/36c78602a7ca866ddb2a917de48ccdd0 ¹.

Note that this applies to the current main branch + soon-to-be-merged PRs; not the latest NumPy release.

it's a gist instead of a tracker issue to avoid attracting hordes of slop PRs ↩

The upcoming NumPy 2.6 release will support shape-typing :)

Navigation Menu

Repository navigation

Typing support for shapes #16544

33 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Typing support for shapes #16544

Description

Activity

emmatyping commented on Dec 7, 2017

rmcgibbo commented on Dec 12, 2017

rmcgibbo commented on Dec 12, 2017

shoyer commented on Dec 12, 2017

rmcgibbo commented on Dec 12, 2017

shoyer commented on Dec 12, 2017

rmcgibbo commented on Dec 12, 2017

hameerabbasi commented on Feb 19, 2018

shoyer commented on Feb 19, 2018

BvB93 commented on Oct 7, 2020

BvB93 commented on Jan 27, 2021

33 remaining items

jorenham commented on Aug 13, 2025

zackw commented on Aug 13, 2025

jakevdp commented on Aug 13, 2025

jorenham commented on Aug 13, 2025

jakevdp commented on Aug 13, 2025

sneakers-the-rat commented on Aug 13, 2025

jorenham commented on Aug 13, 2025

Viicos commented on Aug 13, 2025

sneakers-the-rat commented on Aug 13, 2025

jorenham commented on Aug 14, 2025

sneakers-the-rat commented on Aug 14, 2025

jorenham commented on Aug 14, 2025

jorenham commented on Aug 5, 2026

Footnotes

jorenham commented on Sep 30, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions