shpe

A lightweight static analyzer for tracking and validating NumPy tensor shapes without running the code.

VS Code Extension

Preview

The VS Code extension provides real-time shape inference and error messages.

Installation

It is not yet published on VS Code Marketplace. To install the extension, clone the repository, install the package and build the VS Code extension:

git clone https://github.com/jiaquan-cheng/shpe.git
cd shpe
pip install -e .
make vscode-package

This compiles a .vsix file in the shpe-vscode directory, which you can install as an extension in VS Code.

Terminal

Installation

Prerequisites: Python 3.12+

To install shpe directly:

pip install git+https://github.com/jiaquan-cheng/shpe.git

Usage

shpe path/to/your/file_or_directory

You can use the --show-shapes flag to display the inferred shapes of all expressions in the code:

shpe path/to/your/file_or_directory --show-shapes

Features

• Infers shapes from NumPy array (np.array([1, 2, 3])) and NumPy functions (np.zeros((3, 2)), a.T).
• Validates shape annotations for NumPy arrays (c: Annotated[np.ndarray, (3, 2)]).
• Validates NumPy operations for shape compatibility (a + b, a @ b).
• Infers shapes for simple functions calls and function bodies (c = custom_func(a, b)).
• Tracks scalar variables used in shape definitions (np.zeros((dim, 2))).

Checkout examples/demo.py for a more comprehensive demonstration of shpe's capabilities.

Limitations

We prioritize soundness over completeness, so shpe might miss errors. When shpe is uncertain it, does not infer the shape.

False positive:

• We do not track inplace function modification like b = a.resize((3, 2)), so it might infer the wrong shape.

False negatives:

• Only supports a subset of NumPy arrays and functions.
• No control flow support (if, for, while), variables touched are not inferred.
• No support for recursive functions.
• To keep development simple,shpe identifies functions by suffix, so it may not trigger an error in cases like (var.expand_dims without np. prefix).

If you noticed any bugs or have any feature requests, please report them on GitHub Issues.

Development

Prerequisites: Python 3.12+, uv

To get started locally:

git clone https://github.com/jiaquan-cheng/shpe.git
cd shpe
make setup

We would recommend to using the VS Code extension or the --show-shapes flag to check the inferred shapes of your code while developing.

• make : Runs the test suite and quality checks.
• make lint : Runs Ruff and Mypy for code quality and type safety.
• make format : Auto-format code.
• make test : Runs Pytest.
• make clean : Cleans up the project by removing build artifacts and caches.