I tried to do this but got bitten by the fact that Python didn't support varardic type arguments... somehow never thought to just pass the arguments as a tuple!

I'll try and put together a Numpy/JAX wrapper for this, because I've been looking for something that does compile-time shape checking properly for a long time!

Still wish Python's type system was as powerful as Typescript's... it's got potential that it just doesn't live up to.