Module-level import hook

Hi Patrick,

I would love to have a module-level import hook, similarly to the one implemented in beartype (discussed here, beartype.claw.beartype_this_package) that could be used just by calling a single function on top of __init__.py:

# __init__.py
from jaxtyping import module_wise_import_hook
module_wise_import_hook("beartype.beartype")

from .some_file import some_function # some_function gets typechecked
...

I would say that this is should be a really low priority feature, since I don't think it's that critical, even though it is convenient: I can simulate the same behaviour by doing something like:

# __init.py__
from jaxtyping import install_import_hook
with install_import_hook("module name", "beartype.beartype"):
    from .some_file import some_function

Besides, the implementation of this feature in a safe manner seems to be quite complex, judging by the fact that I was not able to force myself through reading all the source code of beartype implementation.

But, I think that having this written down is important: maybe at some point someone will have enough time to actually implement this. And, of course, there is always hope that you will propose a solution that would take only a single line of code.

Start by reading the existing install_import_hook API and the beartype.claw.beartype_this_package implementation referenced in the issue. Determine how a module-level function could provide the requested init.py usage without the context manager. Done means the proposed hook works for imports within a package and has coverage for its intended behavior.