patrick-kidger / equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Apache License 2.0
2.12k stars 142 forks source link

AbstractVar property: interaction bug with static type-checker #882

Open EtaoinWu opened 1 month ago

EtaoinWu commented 1 month ago

The following code:

import equinox as eqx

class A(eqx.Module):
    x: eqx.AbstractVar[int]

class B(A):
    y: int

    @property
    def x(self):
        return self.y

b = B(y=3)
b.x

runs without a problem, but gives error in pyright and mypy. Pyright complains

⮞ mypy ./test.py
test.py:9: error: Dataclass attribute may only be overridden by another attribute  [misc]
test.py:13: error: Missing positional argument "x" in call to "B"  [call-arg]
Found 2 errors in 1 file (checked 1 source file)

and mypy complains

⮞ pyright ./qwq.py
/path/to/test.py
  /path/to/test.py:10:9 - error: "x" overrides symbol of same name in class "A"
    "property" is not assignable to "int" (reportIncompatibleVariableOverride)
  /path/to/test.py:13:5 - error: Argument missing for parameter "x" (reportCallIssue)
2 errors, 0 warnings, 0 informations

My pylance in vscode somehow only reports the second error ("x" missing) but not the first (property cannot override int).

I believe that AbstractVar property is an intentional feature, as the doc says

An AbstractVar[T] must be overridden by an attribute annotated with AbstractVar[T], AbstractClassVar[T], ClassVar[T], T, or a property returning T.

I suggest this to be fixed by _ModuleMeta adding something (to its __new__?) at typechecking time to modify the __annotations__ of its initializer.

patrick-kidger commented 1 month ago

Yup, implementing an AbstractVar with a property is an intentional feature.

That said, I don't know of a way to make this static-type-checking compatible. If you think you have a way to do that then I'd be happy to see a PR doing so!