PyPie, A DSL for Tensor Programming

PyPie is a subset of Python built on a dependent type system to statically validate tensor shapes.

We have made the "type surface" just rigorous enough but not too much: it catches shape mismatches without breaking the illusion that we are still writing Python. It's also seasoned with rank polymorphism and some basic arithmetic rewriting.

This is our early demo focused on language semantics. Under the hood, PyPie currently compiles to JAX, which does the heavy lifting for optimizations.