google-deepmind / dm-haiku

JAX-based neural network library
https://dm-haiku.readthedocs.io
Apache License 2.0
2.91k stars 231 forks source link

Prevent module creation pyright type errors #702

Closed NeilGirdhar closed 1 year ago

NeilGirdhar commented 1 year ago

Creating any Haiku module gives type errors on the latest Pyright due to the module metaclass specification (https://github.com/microsoft/pyright/discussions/5561). The issue is that the ideal type annotation is an intersection of type[T] and ModuleMetaclass where T is the return type, but Python doesn't have type intersections yet. The authors of the code chose type[T], but Pyright requires ModuleMetaclass.

NeilGirdhar commented 1 year ago

Incidentally, I think most if not all of what is done in the metaclass should be moved into Module.__init_subclass__. Metaclasses have significant compatibility issues (you can't inherit from two classes with different metaclasses). It's best to avoid them.

superbobry commented 1 year ago

Thanks, Neil!