fix(lib): import_hook now correctly handles @no_type_check decorator

[jeertmans]

Hi!

I recently decided to use NVIDIA's warp-lang package for performance reason, and noticed that beartype would complain that warp's type hints are not PEP-compliant. I then tried to put a @no_type_check decorator, but it didn't fix the error.

After some investigation, it appears that this bug is caused by the import hook mechanism implemented by jaxtyping. Fixing it was relatively easy :-)

You can run a MWE with uv run mwe.py.

mwe.py:

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "beartype>=0.22.9",
#     "jaxtyping",
#     "warp-lang>=1.14.0",
# ]
#
# [tool.uv.sources]
# jaxtyping = { path = ".", editable = true }
# ///
from jaxtyping import install_import_hook

with install_import_hook("warp_example", "beartype.beartype"):
    from warp_example import foo

    print(foo)

warp_example.py

import warp as wp
from typing import no_type_check


@no_type_check
@wp.kernel
def foo(x: wp.array[wp.vec3], y: wp.array[float]) -> None:
    tid = wp.tid()
    x[tid] = y[tid] + 1

[patrick-kidger]

LGTM! Thank you for the fix :)

We can do a version bump if you'd like, I don't have anything else outstanding for jaxtyping at the moment?

[jeertmans]

Thanks for merging @patrick-kidger!

Yes that would be great to release a new version so I can merge my PR without having to ship a temporary hack :)