Inheritance of linen Modules does not work

[PgLoLo]

Consider this simple example of three consequent inheritances:

from flax import linen


class First(linen.Module):
    param0: int


class Second(First):
    param1: int
    
    
class Third(Second):
    param2: int

On the definition of the third classes the following exception is raised:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-2-060fb3489d03> in <module>
     11 
     12 
---> 13 class Third(Second):
     14     param2: int

~/data/miniconda3/envs/ml/lib/python3.8/site-packages/flax/linen/module.py in __init_subclass__(cls)
    421     # automatically transform Modules into dataclasses at subclass creation
    422     # time, and we set the last dataclass arguments to `parent` and `name`.
--> 423     cls._customized_dataclass_transform()
    424     # We wrap user-defined methods including setup and __call__ to enforce
    425     # a number of different checks and to provide clear error messages.

~/data/miniconda3/envs/ml/lib/python3.8/site-packages/flax/linen/module.py in _customized_dataclass_transform(cls)
    455     cls.__annotations__ = annotations
    456     # Now apply dataclass transform (which operates in-place).
--> 457     dataclasses.dataclass(cls, unsafe_hash=True, repr=False)  # pytype: disable=wrong-keyword-args
    458     cls.__hash__ = _wrap_hash(cls.__hash__)
    459     # Restore original base class __dataclass_fields__.

~/data/miniconda3/envs/ml/lib/python3.8/dataclasses.py in dataclass(cls, init, repr, eq, order, unsafe_hash, frozen)
   1017 
   1018     # We're called as @dataclass without parens.
-> 1019     return wrap(cls)
   1020 
   1021 

~/data/miniconda3/envs/ml/lib/python3.8/dataclasses.py in wrap(cls)
   1009 
   1010     def wrap(cls):
-> 1011         return _process_class(cls, init, repr, eq, order, unsafe_hash, frozen)
   1012 
   1013     # See if we're being called as @dataclass or @dataclass().

~/data/miniconda3/envs/ml/lib/python3.8/dataclasses.py in _process_class(cls, init, repr, eq, order, unsafe_hash, frozen)
    923                 if f._field_type in (_FIELD, _FIELD_INITVAR)]
    924         _set_new_attribute(cls, '__init__',
--> 925                            _init_fn(flds,
    926                                     frozen,
    927                                     has_post_init,

~/data/miniconda3/envs/ml/lib/python3.8/dataclasses.py in _init_fn(fields, frozen, has_post_init, self_name, globals)
    500                 seen_default = True
    501             elif seen_default:
--> 502                 raise TypeError(f'non-default argument {f.name!r} '
    503                                 'follows default argument')
    504 

TypeError: non-default argument 'param1' follows default argument

What could be a problem? Please note, that two consequent inheritance does work.

[marcvanzee]

I am actually not sure what the bug was, but it appears that this will be fixed in the next release of Flax. When I run your code on a Colab and pip install from our repo the error doesn't show up anymore: https://colab.research.google.com/drive/1sAZSkirM_MbK6FiMGqkbTxObuuTbQMs9?usp=sharing

I've added it to the CHANGELOG: https://github.com/google/flax/blob/master/CHANGELOG.md