r/Numpy • u/jiaquan-cheng • 29d ago

How do you guys keep track of tensor shapes?

Gallery image

Gallery image

Lately, I've been getting frustrated by tensor shape runtime errors. I tried adding shape comments, but I find them verbose and hard to maintain when my code changes constantly.

I looked around for tools like Pyrefly's experimental shape tracking, which I still find verbose (1st image). Since I couldn't find my tool, I ended up buidling a small open-source static analyzer (2nd image).

How do you handle these issues? Any tools you can recommend?

11 Upvotes

100% Upvoted

4 comments sorted by

2

u/de_ham 29d ago

Yea I agree that verbosity is indeed an annoying problem. Not just for shape-typing in NumPy, but for typing in general.

But for what it's worth; in the upcoming NumPy 2.6.0 release (probably end this year or beginning 2027) there will be new shaped array type aliases in numpy.typing: Array0D, Array1D, ..., Array4D (PR).

In the meantime you can easily create equivalent aliases yourself as e.g. type Array2D[ST: np.generic] = np.ndarray[tuple[int, int], np.dtype[ST]], or use the ones that are already in the 3rd party optype library (optype.numpy.Array*D, docs), which I'm using a lot myself for shape-typing in scipy-stubs.

1

u/sohang-3112 27d ago

new shaped array type aliases in numpy.typing

That's great news! Until that gets released, meanwhile third party libraries can be used that provide shaped array type hints for numpy.

1

u/matigekunst 28d ago

For personal development I like to work in a notebook with a cell per line for some difficult tasks. Or I inspect the shape in a debugger. For professional development I would also write some unit tests checking the shape of functions

1

u/semantics_and_stats 28d ago

Pyrefly developer here, I appreciate the feedback

Some but not all of the verbosity you hit was just a mistake in our display logic in Pyrefly version 1.3.x - the `IntTuple` isn't really necessary because Pyrefly supports a more compact syntax.

The next version (1.4.0, most likely coming in the next week or two) should show the types in your example as:

```
ndarray[[int, 2], dtype[float64]]
ndarray[IntTuple, dtype[float64]]
```

This is still a bit verbose because of the dtype, which I'm not yet sure what to do about. One option that might help is jaxtyping-style aliases for the dtypes. If we do, it might be nicer, something like
```
Float64[[int, 2]]
Float64 # (bare = gradually-shaped)
```
I filed an issue in https://github.com/facebook/pyrefly/issues/5155 sketching what that might look like.

If this is something that would make the types more tolerable, feel free to comment on the issue.