module

module

A theseus nn module, which is nn module, execpt it asks you to dump all the types of constitutent parts so we can configure it.

Module

Bases: Module

plot(intermediates: Any) -> Dict[str, Any] staticmethod

intermediates -> [figure]

sharding() -> List[Tuple[str, Optional[Any]]]

Return the sharding configuration for this module.

Returns:

Type	Description
List[Tuple[str, Optional[Any]]]	Tuple[Tuple[str, Optional[Axis]]]: A dictionary mapping Axes to sharding dimensions.

components() -> List[Type[Any]] classmethod

Return the types of constituent parts of this module.

Returns:

Name	Type	Description
Type	List[Type[Any]]	A type or tuple of types representing the constituent parts.

gather() -> List[Type[Any]] classmethod

Depth-first search of all constituent parts of this module.

Returns:

Type	Description
List[Type[Any]]	List[Type]: A list of all constituent part types.