Annotate Module fields

[cgarciae]

What does this PR do?

Fixes #2416. Adds a base class _ModuleBase which contains the annotations for Module's fields, this pattern avoids a couple of pitfalls with dataclasses and makes static analysis tools happy. Thanks @Conchylicultor for proposing the solution!

The approach works but adding the TYPE_CHECKING check over __getattr__ opened the lid to a barricade of pytype errors, question is do we want to solve all typing errors right now?

[codecov-commenter]

Codecov Report

Merging #2447 (149672e) into main (45c2955) will increase coverage by 0.00%. The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main    #2447   +/-   ##
=======================================
  Coverage   78.81%   78.82%           
=======================================
  Files          49       48    -1     
  Lines        5065     5067    +2     
=======================================
+ Hits         3992     3994    +2     
  Misses       1073     1073           

Impacted Files	Coverage Δ
flax/linen/module.py	92.86% <100.00%> (+0.11%)	:arrow_up:
flax/__init__.py

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[Conchylicultor]

Why was it closed ?