Skip to main content

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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

mypy-einsum-1.0.0.tar.gz (6.5 kB view details)

Uploaded Source

Built Distribution

mypy_einsum-1.0.0-py3-none-any.whl (5.0 kB view details)

Uploaded Python 3

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

Hashes for mypy-einsum-1.0.0.tar.gz
Algorithm Hash digest
SHA256 e79b0a2b1d77300bd1ea99a977ca2b51a35d20845efac165349f0f0e5156a7bb
MD5 47d5fda1f3d005130647c57f349ec06c
BLAKE2b-256 04a54c32eeaaf6199778f6018d1bf6814e4fbbfc3c57f2301b00cc10b2c3eb90

See more details on using hashes here.

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

Hashes for mypy_einsum-1.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 210651116788ee55bcc1b7ae441d4caa2ece52c52d8152fd192529b99b44523b
MD5 3985e66aa764099614d90994ec3d953c
BLAKE2b-256 886026e1a24d12f57be8ddef68b383a927a445b92cfe7692bf158eea0250326f

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page