This is sort of the polar opposite of how modern high performance ML frameworks are built. Skimming the code, there's a ton of boilerplate for the various operations that could be library-ized and generified, if that makes sense.

I actually really like minimal implementations of state-of-the-art systems because the code is much easier to understand (modern frameworks are super-complex) but I wonder what it means long-term if you don't need frameworks.