Type checking + autocomplete don't work with the Linen API

Modules in Linen are defined similar to dataclasses yet they aren’t dataclasses as such so the type checkers and the IDEs have no way on knowing what arguments each Module expects. This yields a lot of type checking errors and makes it harder to get automatic documentation.

Since the API is new maybe this could be taken into consideration early to avoid changes in the future.

Also, how does one perform logic similar to the one you would do in __init__ in this API? I know about setup but sometimes you e.g. accept union types in the constructor and transform them to a standardize type for a field.

[thomkeh]

The pyright team has a proposal for how to specify “dataclass-like semantics” for custom classes, which would solve the problem in this issue. They intend to submit this as a PEP, but they have also already implemented this spec in pyright and pylance, so VS Code users can directly make use of it. pydantic has a PR for supporting this spec.

Would the flax team be open to supporting this? As far as I can tell, it just involves creating a decorator called __dataclass_transform__ (that doesn’t do anything at runtime) and applying it to the base class.

[NeilGirdhar]

FYI: The dataclass transform PEP is now in draft.

I mentioned Flax’s dataclass during the discussion, and they agreed to support applying the decorator to classes directly (instead of just to metaclasses), and pyright has already implemented that suggestion.