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.