Strange behavior with eqx.Module

[knyazer]

I encountered a strange behavior when trying to install_import_hooks for a class inherited from eqx.Module: there appears a warning that states about typechecker inability to wrap the init method, which to me seems strange, since eqx.Module is (promised to be) a dataclass.

Warning message:

/home/user/.local/lib/python3.10/site-packages/jaxtyping/_decorator.py:139: InstrumentationWarning: @typechecked only supports instrumenting functions wrapped with @classmethod, @staticmethod or @property -- not typechecking source.A.__init__
  init = jaxtyped(typechecker(kls.__init__))

Minimal example:

##### main.py #####
from jaxtyping import install_import_hook
with install_import_hook("source", typechecker="typeguard.typechecked"):
    from source import *

##### source.py #####
import equinox as eqx

class A(eqx.Module):
    x: int
    def __init__(self, y: int):
        self.x = y

I reckon there is a high chance I am simply not able to use the method correctly, and if so, I would appreciate any general feedback, or a short example of how to do it correctly. The problem seems to be directly related to issue 46.

[patrick-kidger]

What version of typeguard do you have installed?

[knyazer]

pip3 freeze | grep -E "typeguard|equinox|jax|jaxtyping"

gives

equinox==0.10.11
jax==0.4.14
jaxlib==0.4.14+cuda12.cudnn89
jaxopt==0.7
jaxtyping==0.2.20
typeguard==4.1.0

I believe 4.1.0 version of typeguard is the latest pypi version available. I use Python 3.10.4, OS Ubuntu 22.04.

[patrick-kidger]

Thanks! Okay, so I think this is a typeguard bug. If you write A("hi"), then you'll see the type error being raised correctly, despite the warning. E.g. A("hi") will raise an error.

For now I'd suggest either ignoring the warning, or switching to typeguard version 2.*, or using beartype instead.

[knyazer]

Thanks for debugging the problem for me. Yes, even in my original code this warning seems to be meaningless, since both jax-types and python-types mismatches are raised correctly.