Try using jaxtyping: https://github.com/google/jaxtyping. It also supports numpy...

		patrickkidger on May 21, 2023 \| parent \| context \| favorite \| on: Writing Python like it's Rust Try using jaxtyping: https://github.com/google/jaxtyping. It also supports numpy/pytorch/etc.