Fix `_ParentType` annotation

Commit

2 years ago

Fix `_ParentType` annotation Fix definition of `_ParentType` so that runtime type checkers work with nested Flax modules.