I use jaxtyping as documentation, but the fact it can only be used for runtime checking (in a slightly clunky manner) and can't infer shapes based on ops really limits its utility imo.