We're working on statically checking Jaxtyping annotations in Pyrefly, but it's incomplete and not ready to use yet :)