3 ms·
Try using jaxtyping: https://github.com/google/jaxtyping https://github.com/google/jaxtyping. It also supports numpy/pytorch/etc.
by patrickkidger 3y ago
Try using jaxtyping: https://github.com/google/jaxtyping https://github.com/google/jaxtyping.
It also supports numpy/pytorch/etc.