search

patrick-kidger / jaxtyping

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/
Other
1.14k stars 59 forks source link

`install_import_hook` skip `no_type_check` #215

Closed nstarman closed 3 months ago

nstarman commented 3 months ago

If a function is decorated with typing.no_type_check then it would be great if jaxtyping.install_import_hook would skip decorating that object. Thanks!

patrick-kidger commented 3 months ago

Agreed, that sounds reasonable! See #216.

To be precise, it will still decorate the object because the import hook puts in underneath the typing.no_type_check decorator. But if the inserted decorator itself has a __no_type_check__ attribute, we disable ourselves.