Weird KeyError: '0' when using None typechecker for install_import_hook

[EtaoinWu]

Problem

The following code gives a weird error:

from jaxtyping import install_import_hook
with install_import_hook('temp', None):
    import temp

With the file temp.py:

from jaxtyping import Float, Scalar

def fn(x: Float[Scalar, ""]):
    return x + x

Expected behavior: import success without error.

Actual behavior: a KeyError: '0' occurred. The complete error message:

Traceback (most recent call last):
  File "[directory redacted]\import_test.py", line 3, in <module>
    import temp
  File "<frozen importlib._bootstrap>", line 1354, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1325, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 929, in _load_unlocked
  File "[conda directory redacted]\envs\jaxtest\Lib\site-packages\jaxtyping\_import_hook.py", line 232, in exec_module
    return super().exec_module(module)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "[directory redacted]\temp.py", line 1, in <module>
    from jaxtyping import Float, Scalar
KeyError: '0'

Environment

Windows x64, but platform should be irrevalent.

micromamba create -n jaxtest python=3.12 # or use conda/mamba, should be the same
micromamba activate jaxtest
pip install jax[cpu] jaxtyping # jax==0.4.20, jaxtyping==0.2.23 (both latest as of the issue is created)

[EtaoinWu]

Explanation for the issue's title: the KeyError does not occur when a typechecker, e.g. 'beartype.beartype' or 'typeguard.typechecked' is provided to install_import_hook.

[EtaoinWu]

Workaround

Create a trivial decorator...

# trivial.py
def decorator(x):
    print('trivial decorator called on', x) # optional
    return x

... and feed it into install_import_hook.

with install_import_hook('temp', 'trivial.decorator'):
    import temp

[patrick-kidger]

Thanks for the report! I've just fixed this in #146.

[patrick-kidger]

Closing as fixed in the v0.2.24 release!