Comment by dwohnitmok

This isn't static, but jaxtyping gives you at least runtime checks and also a standardized form of documenting those types. https://github.com/patrick-kidger/jaxtyping