Here's a Python library for static type checking of tensor shapes: https://pypie.dev/

Thanks. Did not know about it.

Looks a little FiShy. In a good way [0].

[0] https://link.springer.com/article/10.1007/s100090050037