pytype.overlays.flax_overlay

module documentation
(source)

Support for flax.struct dataclasses.

Class Dataclass Implements the @dataclass decorator.
Class DataclassOverlay A custom overlay for the 'flax.struct' module.
Class LinenModuleOverlay A custom overlay for the 'flax.linen.module' module.
Class LinenOverlay A custom overlay for the 'flax.linen' module.
Class Module Construct a dataclass for any class inheriting from Module.
Class ModuleDataclass Dataclass with automatic 'name' and 'parent' members.