jaxsnn.discrete.types.PopulationFunctions

class jaxsnn.discrete.types.PopulationFunctions(init: Callable[[jax.Array], Optional[jax.Array, None]], state: Callable[], Tuple[Optional[StateT, None], Optional[jax.Array, None]]], step: Callable[[Dict[str, jax.Array], StateT, jax.Array], Tuple[Optional[StateT, None], jax.Array]])

Bases: object

Function container for population modules

__init__(init: Callable[[jax.Array], Optional[jax.Array, None]], state: Callable[], Tuple[Optional[StateT, None], Optional[jax.Array, None]]], step: Callable[[Dict[str, jax.Array], StateT, jax.Array], Tuple[Optional[StateT, None], jax.Array]]) → None

Initialize self. See help(type(self)) for accurate signature.

Methods

__init__(init, state, step)

Initialize self.

init: Callable[[jax.Array], Optional[jax.Array, None]]

state: Callable[], Tuple[Optional[StateT, None], Optional[jax.Array, None]]]

step: Callable[[Dict[str, jax.Array], StateT, jax.Array], Tuple[Optional[StateT, None], jax.Array]]