Mypy plugin for NumPy/Jax/PyTorch einsum operations
Project description
MyPy Type Checking for NumPy/Jax/PyTorch Einsum Operations
mypy_einsum
is a Mypy plugin for type checking np.einsum
, jax.numpy.einsum
, and torch.einsum
operations.
The Einstein summation convention can be used to compute many multi-dimensional, linear algebraic array operations. einsum
provides a succinct way of representing these.
However, since einsum
equations are passed as a string, it is very easy to overlook typos or other bugs as linters are unable to help.
mypy_einsum
is a Mypy plugin that that is able to statically verify the correctness of einsum
equations with needing to execute the code.
Installation
mypy_einsum
can be installed with pip
:
pip install mypy-einsum
Setup
To enable the plugin, add it to you projects Mypy configuration file.
Usually mypy.ini
:
[mypy]
plugins = mypy_einsum
or pyproject.toml
:
[tool.mypy]
plugins = ["mypy_einsum"]
Example
Can you spot the 🐛 without running the code?
import numpy as np
a = np.arange(9).reshape(3, 3)
np.einsum("ik,kj->ij", a)
mypy_einsum
will catch it for you:
❯ mypy example.py --pretty
example.py:5: error: Number of einsum subscripts must be equal to the
number of operands. [einsum]
np.einsum("ik,kj->ij", a)
^~~~~~~~~~~
Found 1 error in 1 file (checked 1 source file)
After fixing it mypy
will succeed 🎉:
np.einsum("ik,kj->ij", a, a)
❯ mypy example.py
Success: no issues found in 1 source file
Supported Operations
Reporting Issues and Contributing
mypy_einsum
aims to never raise warnings for valid einsum
operations. If you encounter a warning that you believe is incorrect, or think mypy_einsum
is not reporting an error please let us know. Contributions are very welcome!
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
File details
Details for the file mypy-einsum-1.0.0.tar.gz
.
File metadata
- Download URL: mypy-einsum-1.0.0.tar.gz
- Upload date:
- Size: 6.5 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.11.6
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | e79b0a2b1d77300bd1ea99a977ca2b51a35d20845efac165349f0f0e5156a7bb |
|
MD5 | 47d5fda1f3d005130647c57f349ec06c |
|
BLAKE2b-256 | 04a54c32eeaaf6199778f6018d1bf6814e4fbbfc3c57f2301b00cc10b2c3eb90 |
File details
Details for the file mypy_einsum-1.0.0-py3-none-any.whl
.
File metadata
- Download URL: mypy_einsum-1.0.0-py3-none-any.whl
- Upload date:
- Size: 5.0 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.11.6
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 210651116788ee55bcc1b7ae441d4caa2ece52c52d8152fd192529b99b44523b |
|
MD5 | 3985e66aa764099614d90994ec3d953c |
|
BLAKE2b-256 | 886026e1a24d12f57be8ddef68b383a927a445b92cfe7692bf158eea0250326f |