Activations

Activation functions are provided as layers (callable components) in the public API.

They operate on Tensor inputs and are compatible with automatic differentiation.

---

keydnn.ReLU

Bases: StatelessConfigMixin, Module

ReLU activation module.

This layer applies the rectified linear unit elementwise:

relu(x) = max(0, x)

Notes

This module is a thin wrapper around ReLUFn that provides a Module interface and attaches an autograd Context to the output when gradients are required.

parameters

parameters() -> Iterable[IParameter]

Return an iterable over this module's parameters (recursive).

Returns:

Type	Description
Iterable[IParameter]	Iterable of parameters registered on this module and all submodules.

train

train() -> Self

Set this module to training mode and recursively set all child modules to training mode.

Notes

• This toggles self.training = True.
• Modules that behave differently in training (e.g., Dropout, BatchNorm) should read self.training during forward/predict to decide behavior.
• This method is intended to mirror PyTorch's Module.train().

eval

eval() -> Self

Set this module to evaluation (inference) mode and recursively set all child modules to evaluation mode.

Notes

• This toggles self.training = False.
• In eval mode, modules such as Dropout should be disabled, and BatchNorm should use running statistics (if implemented).
• This method is intended to mirror PyTorch's Module.eval().

register_parameter

register_parameter(
    name: str, param: Optional[IParameter]
) -> None

Register a parameter with this module.

Parameters:

Name	Type	Description	Default
name	str	Name under which the parameter will be stored (e.g., "weight", "bias").	required
param	Optional[IParameter]	Parameter instance to register. If None, registration is skipped.	required

Notes

• If param is None, nothing is registered.
• If the name already exists, it is overwritten intentionally.
• This also sets the attribute on the module so self.<name> works.

register_module

register_module(
    name: str, module: Optional["Module"]
) -> None

Register a child module with this module.

Parameters:

Name	Type	Description	Default
name	str	Name under which the module will be stored.	required
module	Optional[Module]	Child module to register. If None, registration is skipped.	required

Notes

• If module is None, nothing is registered.
• This also sets the attribute on the module so self.<name> works.

named_parameters

named_parameters(
    prefix: str = "",
) -> Iterator[tuple[str, IParameter]]

Return an iterator over (name, parameter) pairs (recursive).

Parameters:

Name	Type	Description	Default
prefix	str	Prefix to prepend to parameter names (used for recursion).	''

Returns:

Type	Description
Iterator[tuple[str, IParameter]]	Iterator yielding (fully_qualified_name, parameter).

get_config

get_config() -> Dict[str, Any]

Return a JSON-serializable configuration dictionary.

For stateless modules, this method returns an empty dictionary, indicating that no parameters are required to reconstruct the object.

Returns:

Type	Description
Dict[str, Any]	An empty configuration dictionary.

from_config classmethod

from_config(cfg: Dict[str, Any]) -> Self

Reconstruct the module from a configuration dictionary.

Since stateless modules do not require any configuration parameters, the provided configuration is ignored and a default instance of the class is returned.

Parameters:

Name	Type	Description	Default
cfg	Dict[str, Any]	Configuration dictionary (unused).	required

Returns:

Type	Description
StatelessConfigMixin	A newly constructed instance of the module.

to

to(device: Device) -> 'Module'

Move this module (recursively) to device by moving all registered Parameters.

Notes

• Uses _parameters / _modules registries as the source of truth. This avoids touching properties / methods / non-parameter attributes.
• Assumes each Parameter/Tensor implements .to(Device) -> same-type-like.
• Rebinds attributes so self.weight, etc. now point to the moved objects.

to_

to_(device: Device) -> 'Module'

Move this module and all of its parameters to device in-place.

This method performs a recursive, in-place device migration of all parameters registered on this module and its submodules. Unlike Module.to(), which may rebind parameters to newly created objects, to_() attempts to preserve the identity of each parameter whenever possible.

Behavior

• For each registered parameter:
  • If the parameter implements to_(), it is migrated in-place (object identity is preserved).
  • Otherwise, the parameter is migrated out-of-place via to(device) and rebound on the module as a fallback.
• All child modules are recursively migrated using the same rules.

Parameters:

Name	Type	Description	Default
device	Device	Target device to which all parameters should be moved.	required

Returns:

Type	Description
Module	This module (self), after in-place migration.

Notes

• This method relies exclusively on the _parameters and _modules registries and does not inspect arbitrary attributes.
• In-place migration is best-effort and depends on parameter support for to_(). Parameters that do not implement to_() will be replaced by newly created objects.
• Autograd context is not preserved across device transfers; parameters should be treated as graph breaks after migration.
• Optimizers that hold references to parameters remain valid only if all parameters support true in-place migration.

forward

forward(x: Tensor) -> Tensor

Apply the ReLU activation to the input tensor.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required

Returns:

Type	Description
Tensor	Output tensor containing relu(x) elementwise.

Notes

If x.requires_grad is True, the returned tensor will have an attached Context whose backward_fn delegates to ReLUFn.backward.

keydnn.Sigmoid

Bases: StatelessConfigMixin, Module

Sigmoid activation module.

This layer applies the sigmoid function elementwise:

sigmoid(x) = 1 / (1 + exp(-x))

Notes

This module is a thin wrapper around SigmoidFn that provides a Module interface and attaches an autograd Context to the output when gradients are required.

parameters

parameters() -> Iterable[IParameter]

Return an iterable over this module's parameters (recursive).

Returns:

Type	Description
Iterable[IParameter]	Iterable of parameters registered on this module and all submodules.

train

train() -> Self

Set this module to training mode and recursively set all child modules to training mode.

Notes

• This toggles self.training = True.
• Modules that behave differently in training (e.g., Dropout, BatchNorm) should read self.training during forward/predict to decide behavior.
• This method is intended to mirror PyTorch's Module.train().

eval

eval() -> Self

Set this module to evaluation (inference) mode and recursively set all child modules to evaluation mode.

Notes

• This toggles self.training = False.
• In eval mode, modules such as Dropout should be disabled, and BatchNorm should use running statistics (if implemented).
• This method is intended to mirror PyTorch's Module.eval().

register_parameter

register_parameter(
    name: str, param: Optional[IParameter]
) -> None

Register a parameter with this module.

Parameters:

Name	Type	Description	Default
name	str	Name under which the parameter will be stored (e.g., "weight", "bias").	required
param	Optional[IParameter]	Parameter instance to register. If None, registration is skipped.	required

Notes

• If param is None, nothing is registered.
• If the name already exists, it is overwritten intentionally.
• This also sets the attribute on the module so self.<name> works.

register_module

register_module(
    name: str, module: Optional["Module"]
) -> None

Register a child module with this module.

Parameters:

Name	Type	Description	Default
name	str	Name under which the module will be stored.	required
module	Optional[Module]	Child module to register. If None, registration is skipped.	required

Notes

• If module is None, nothing is registered.
• This also sets the attribute on the module so self.<name> works.

named_parameters

named_parameters(
    prefix: str = "",
) -> Iterator[tuple[str, IParameter]]

Return an iterator over (name, parameter) pairs (recursive).

Parameters:

Name	Type	Description	Default
prefix	str	Prefix to prepend to parameter names (used for recursion).	''

Returns:

Type	Description
Iterator[tuple[str, IParameter]]	Iterator yielding (fully_qualified_name, parameter).

get_config

get_config() -> Dict[str, Any]

Return a JSON-serializable configuration dictionary.

For stateless modules, this method returns an empty dictionary, indicating that no parameters are required to reconstruct the object.

Returns:

Type	Description
Dict[str, Any]	An empty configuration dictionary.

from_config classmethod

from_config(cfg: Dict[str, Any]) -> Self

Reconstruct the module from a configuration dictionary.

Since stateless modules do not require any configuration parameters, the provided configuration is ignored and a default instance of the class is returned.

Parameters:

Name	Type	Description	Default
cfg	Dict[str, Any]	Configuration dictionary (unused).	required

Returns:

Type	Description
StatelessConfigMixin	A newly constructed instance of the module.

to

to(device: Device) -> 'Module'

Move this module (recursively) to device by moving all registered Parameters.

Notes

• Uses _parameters / _modules registries as the source of truth. This avoids touching properties / methods / non-parameter attributes.
• Assumes each Parameter/Tensor implements .to(Device) -> same-type-like.
• Rebinds attributes so self.weight, etc. now point to the moved objects.

to_

to_(device: Device) -> 'Module'

Move this module and all of its parameters to device in-place.

This method performs a recursive, in-place device migration of all parameters registered on this module and its submodules. Unlike Module.to(), which may rebind parameters to newly created objects, to_() attempts to preserve the identity of each parameter whenever possible.

Behavior

• For each registered parameter:
  • If the parameter implements to_(), it is migrated in-place (object identity is preserved).
  • Otherwise, the parameter is migrated out-of-place via to(device) and rebound on the module as a fallback.
• All child modules are recursively migrated using the same rules.

Parameters:

Name	Type	Description	Default
device	Device	Target device to which all parameters should be moved.	required

Returns:

Type	Description
Module	This module (self), after in-place migration.

Notes

• This method relies exclusively on the _parameters and _modules registries and does not inspect arbitrary attributes.
• In-place migration is best-effort and depends on parameter support for to_(). Parameters that do not implement to_() will be replaced by newly created objects.
• Autograd context is not preserved across device transfers; parameters should be treated as graph breaks after migration.
• Optimizers that hold references to parameters remain valid only if all parameters support true in-place migration.

forward

forward(x: Tensor) -> Tensor

Apply the sigmoid activation to the input tensor.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required

Returns:

Type	Description
Tensor	Output tensor containing sigmoid(x) elementwise.

Notes

If x.requires_grad is True, the returned tensor will have an attached Context whose backward_fn delegates to SigmoidFn.backward.

keydnn.Tanh

Bases: StatelessConfigMixin, Module

Hyperbolic tangent activation module.

This layer applies the tanh function elementwise:

tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))

Notes

This module is a thin wrapper around TanhFn that provides a Module interface and attaches an autograd Context to the output when gradients are required.

parameters

parameters() -> Iterable[IParameter]

Return an iterable over this module's parameters (recursive).

Returns:

Type	Description
Iterable[IParameter]	Iterable of parameters registered on this module and all submodules.

train

train() -> Self

Set this module to training mode and recursively set all child modules to training mode.

Notes

• This toggles self.training = True.
• Modules that behave differently in training (e.g., Dropout, BatchNorm) should read self.training during forward/predict to decide behavior.
• This method is intended to mirror PyTorch's Module.train().

eval

eval() -> Self

Set this module to evaluation (inference) mode and recursively set all child modules to evaluation mode.

Notes

• This toggles self.training = False.
• In eval mode, modules such as Dropout should be disabled, and BatchNorm should use running statistics (if implemented).
• This method is intended to mirror PyTorch's Module.eval().

register_parameter

register_parameter(
    name: str, param: Optional[IParameter]
) -> None

Register a parameter with this module.

Parameters:

Name	Type	Description	Default
name	str	Name under which the parameter will be stored (e.g., "weight", "bias").	required
param	Optional[IParameter]	Parameter instance to register. If None, registration is skipped.	required

Notes

• If param is None, nothing is registered.
• If the name already exists, it is overwritten intentionally.
• This also sets the attribute on the module so self.<name> works.

register_module

register_module(
    name: str, module: Optional["Module"]
) -> None

Register a child module with this module.

Parameters:

Name	Type	Description	Default
name	str	Name under which the module will be stored.	required
module	Optional[Module]	Child module to register. If None, registration is skipped.	required

Notes

• If module is None, nothing is registered.
• This also sets the attribute on the module so self.<name> works.

named_parameters

named_parameters(
    prefix: str = "",
) -> Iterator[tuple[str, IParameter]]

Return an iterator over (name, parameter) pairs (recursive).

Parameters:

Name	Type	Description	Default
prefix	str	Prefix to prepend to parameter names (used for recursion).	''

Returns:

Type	Description
Iterator[tuple[str, IParameter]]	Iterator yielding (fully_qualified_name, parameter).

get_config

get_config() -> Dict[str, Any]

Return a JSON-serializable configuration dictionary.

For stateless modules, this method returns an empty dictionary, indicating that no parameters are required to reconstruct the object.

Returns:

Type	Description
Dict[str, Any]	An empty configuration dictionary.

from_config classmethod

from_config(cfg: Dict[str, Any]) -> Self

Reconstruct the module from a configuration dictionary.

Since stateless modules do not require any configuration parameters, the provided configuration is ignored and a default instance of the class is returned.

Parameters:

Name	Type	Description	Default
cfg	Dict[str, Any]	Configuration dictionary (unused).	required

Returns:

Type	Description
StatelessConfigMixin	A newly constructed instance of the module.

to

to(device: Device) -> 'Module'

Move this module (recursively) to device by moving all registered Parameters.

Notes

• Uses _parameters / _modules registries as the source of truth. This avoids touching properties / methods / non-parameter attributes.
• Assumes each Parameter/Tensor implements .to(Device) -> same-type-like.
• Rebinds attributes so self.weight, etc. now point to the moved objects.

to_

to_(device: Device) -> 'Module'

Move this module and all of its parameters to device in-place.

This method performs a recursive, in-place device migration of all parameters registered on this module and its submodules. Unlike Module.to(), which may rebind parameters to newly created objects, to_() attempts to preserve the identity of each parameter whenever possible.

Behavior

• For each registered parameter:
  • If the parameter implements to_(), it is migrated in-place (object identity is preserved).
  • Otherwise, the parameter is migrated out-of-place via to(device) and rebound on the module as a fallback.
• All child modules are recursively migrated using the same rules.

Parameters:

Name	Type	Description	Default
device	Device	Target device to which all parameters should be moved.	required

Returns:

Type	Description
Module	This module (self), after in-place migration.

Notes

• This method relies exclusively on the _parameters and _modules registries and does not inspect arbitrary attributes.
• In-place migration is best-effort and depends on parameter support for to_(). Parameters that do not implement to_() will be replaced by newly created objects.
• Autograd context is not preserved across device transfers; parameters should be treated as graph breaks after migration.
• Optimizers that hold references to parameters remain valid only if all parameters support true in-place migration.

forward

forward(x: Tensor) -> Tensor

Apply the tanh activation to the input tensor.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required

Returns:

Type	Description
Tensor	Output tensor containing tanh(x) elementwise.

Notes

If x.requires_grad is True, the returned tensor will have an attached Context whose backward_fn delegates to TanhFn.backward.

keydnn.Softmax

Bases: Module

Softmax activation module.

This module applies the softmax function to its input tensor along a specified axis, producing a normalized probability distribution.

By default, softmax is applied over the last dimension, which is the standard convention for classification outputs.

parameters

parameters() -> Iterable[IParameter]

Return an iterable over this module's parameters (recursive).

Returns:

Type	Description
Iterable[IParameter]	Iterable of parameters registered on this module and all submodules.

train

train() -> Self

Set this module to training mode and recursively set all child modules to training mode.

Notes

• This toggles self.training = True.
• Modules that behave differently in training (e.g., Dropout, BatchNorm) should read self.training during forward/predict to decide behavior.
• This method is intended to mirror PyTorch's Module.train().

eval

eval() -> Self

Set this module to evaluation (inference) mode and recursively set all child modules to evaluation mode.

Notes

• This toggles self.training = False.
• In eval mode, modules such as Dropout should be disabled, and BatchNorm should use running statistics (if implemented).
• This method is intended to mirror PyTorch's Module.eval().

register_parameter

register_parameter(
    name: str, param: Optional[IParameter]
) -> None

Register a parameter with this module.

Parameters:

Name	Type	Description	Default
name	str	Name under which the parameter will be stored (e.g., "weight", "bias").	required
param	Optional[IParameter]	Parameter instance to register. If None, registration is skipped.	required

Notes

• If param is None, nothing is registered.
• If the name already exists, it is overwritten intentionally.
• This also sets the attribute on the module so self.<name> works.

register_module

register_module(
    name: str, module: Optional["Module"]
) -> None

Register a child module with this module.

Parameters:

Name	Type	Description	Default
name	str	Name under which the module will be stored.	required
module	Optional[Module]	Child module to register. If None, registration is skipped.	required

Notes

• If module is None, nothing is registered.
• This also sets the attribute on the module so self.<name> works.

named_parameters

named_parameters(
    prefix: str = "",
) -> Iterator[tuple[str, IParameter]]

Return an iterator over (name, parameter) pairs (recursive).

Parameters:

Name	Type	Description	Default
prefix	str	Prefix to prepend to parameter names (used for recursion).	''

Returns:

Type	Description
Iterator[tuple[str, IParameter]]	Iterator yielding (fully_qualified_name, parameter).

to

to(device: Device) -> 'Module'

Move this module (recursively) to device by moving all registered Parameters.

Notes

• Uses _parameters / _modules registries as the source of truth. This avoids touching properties / methods / non-parameter attributes.
• Assumes each Parameter/Tensor implements .to(Device) -> same-type-like.
• Rebinds attributes so self.weight, etc. now point to the moved objects.

to_

to_(device: Device) -> 'Module'

Move this module and all of its parameters to device in-place.

This method performs a recursive, in-place device migration of all parameters registered on this module and its submodules. Unlike Module.to(), which may rebind parameters to newly created objects, to_() attempts to preserve the identity of each parameter whenever possible.

Behavior

• For each registered parameter:
  • If the parameter implements to_(), it is migrated in-place (object identity is preserved).
  • Otherwise, the parameter is migrated out-of-place via to(device) and rebound on the module as a fallback.
• All child modules are recursively migrated using the same rules.

Parameters:

Name	Type	Description	Default
device	Device	Target device to which all parameters should be moved.	required

Returns:

Type	Description
Module	This module (self), after in-place migration.

Notes

• This method relies exclusively on the _parameters and _modules registries and does not inspect arbitrary attributes.
• In-place migration is best-effort and depends on parameter support for to_(). Parameters that do not implement to_() will be replaced by newly created objects.
• Autograd context is not preserved across device transfers; parameters should be treated as graph breaks after migration.
• Optimizers that hold references to parameters remain valid only if all parameters support true in-place migration.

forward

forward(x: Tensor) -> Tensor

Apply the softmax activation to the input tensor.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor to which softmax will be applied.	required

Returns:

Type	Description
Tensor	A tensor of the same shape as x, where values along the specified axis form a probability distribution (sum to 1).

Notes

• This method delegates the numerical computation to SoftmaxFn, attaching a backward Context when gradient tracking is enabled.
• The returned tensor will participate in autograd only if x.requires_grad is True.
• Gradient propagation is implemented via a Jacobian–vector product in SoftmaxFn.backward, avoiding explicit Jacobian construction.

---

Notes

• Unless otherwise documented, activations preserve input shape.
• Softmax typically takes a dim/axis argument to specify which dimension to normalize over.
• For numerical stability, prefer using losses that accept logits directly (if your API supports that).