The Pyrefly type checker is starting to work on this kind of shape hinting - so far it only works on Torch but I believe the plan is for it to work with other array packages (eg. JAX, NumPy)

https://pyrefly.org/en/docs/tensor-shapes/#how-it-works