flax: 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.
About this issue
- Original URL
- State: closed
- Created 4 years ago
- Reactions: 6
- Comments: 16 (7 by maintainers)
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.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.
Let’s give this a try. I hope we can also get this supported in pytype
Hi @thomkeh – I think that’s a good idea. If you get it to work with and without an explicit
@dataclass, with good tests, I think we could probably merge it. @jheek is on vacation this week but I’d like to get his review on the idea as well.If better IDE support is very important, perhaps we could consider changing our example patterns as well.