Try using jaxtyping: https://github.com/google/jaxtyping.

It also supports numpy/pytorch/etc.