jaxsnn.discrete.types.GeneratorFunctions

class jaxsnn.discrete.types.GeneratorFunctions(*args, **kwargs)

Bases: Protocol

Base protocol for all module function containers

__init__(*args, **kwargs)

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

Methods

__init__(*args, **kwargs)

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]]