Module Class

Module Class#

class evox.Stateful[source]#

Base class for all evox modules.

This module allow easy managing of states.

All the constants (e.g. hyperparameters) are initialized in the __init__, and mutated states are initialized in the setup method.

The init method will automatically call the setup of the current module and recursively call setup methods of all submodules.

Currently, there are two special metadata that can be used to control the behavior of the module initialization: - stack: If set to True, the module will be initialized multiple times, and the states will be stacked together. - nested: If set to True, the a list of modules, that is [module1, module2, …], will be iterated and initialized.

__init__()[source]#
init(key=None, no_state=False, re_init=False)[source]#

Initialize this module and all submodules

This method should not be overwritten.

Parameters:

  • key (Array | None) – A PRNGKey.

  • no_state (bool)

  • re_init (bool)

Returns:

The state of this module and all submodules combined.

Return type:

State

parallel_init(key, num_copies, no_state=False)[source]#

Initialize multiple copies of this module in parallel

This method should not be overwritten.

Parameters:

  • key (Array) – A PRNGKey.

  • num_copies (int) – The number of copies to be initialized

  • no_state (bool) – Whether to skip the state initialization

Returns:

The state of this module and all submodules combined, and the last node_id

Return type:

Tuple[State, int]

setup(key)[source]#

Setup mutable state here

The state it self is immutable, but it act as a mutable state by returning new state each time.

Parameters:

key (Array) – A PRNGKey.

Returns:

The state of this module.

Return type:

State

Contents