I think we're talking at cross purposes. The reason I'm excited about the Pyrefly work is that it leverages the type system to infer array shapes statically, which makes it simpler to reason about them when you're writing the code and catch bugs. The fact that people have developed these janky approaches to shape tracking suggests that there's a gap to be filled.
Jax and Torch don't do that statically. They obviously have to do it at runtime, but that doesn't address this particular issue. I mention Eigen because array shape hinting is generally useful for any linalg library, not purely for ML applications.