flax

This change is needed to allow using @flax.struct.dataclass on a class that inherits itself from another dataclass created with @flax.struct.dataclass.

#1097

Merged