evox.problems.neuroevolution.utils

`evox.problems.neuroevolution.utils`#

Module Contents#

Classes#

ModelStateForwardResult

Functions#

get_vmap_model_state_forward

Get model state forward function for vmap and non-vmap models. When get_non_vmap is False, the function returns only vmap model state forward function. When get_non_vmap is True, the function returns both vmap and non-vmap model state forward functions.

API#

class evox.problems.neuroevolution.utils.ModelStateForwardResult[source]#

Bases: typing.NamedTuple

init_state: Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]#: None

state_forward: Callable#: None

evox.problems.neuroevolution.utils.get_vmap_model_state_forward(model: torch.nn.Module, pop_size: int, device: torch.device, in_dims=(0, None), randomness='different') → evox.problems.neuroevolution.utils.ModelStateForwardResult[source]#: Get model state forward function for vmap and non-vmap models. When get_non_vmap is False, the function returns only vmap model state forward function. When get_non_vmap is True, the function returns both vmap and non-vmap model state forward functions.