Module

Module is the base class for all neural network layers in simplegrad. Subclass it and implement forward() to define a custom layer. Parameter tensors and sub-modules are discovered automatically via attribute introspection, so you can call .parameters() or .summary() on any composite model without wiring anything by hand.

import simplegrad as sg
from simplegrad.core import Module

class MyLayer(Module):
    def __init__(self):
        super().__init__()
        self.weight = sg.normal((4, 4), requires_grad=True)

    def forward(self, x):
        return x @ self.weight

layer = MyLayer()
out = layer(sg.ones((2, 4)))

Module

Base class for all neural network layers.

Subclass this and implement forward() to create custom layers. Parameters (leaf Tensor attributes) and sub-modules are discovered automatically via attribute introspection.

Attributes

Attribute	Type	Description
.label	str	Display name of the module. Defaults to the class name.
.eval_mode	bool	True when in evaluation mode (set via .set_eval_mode()).

Methods

Method	Description
.forward()	Define the forward pass. Must be implemented by subclasses.
.parameters()	Return all parameter tensors in this module and its sub-modules.
.submodules()	Return all direct sub-modules as a named dict.
.to_device()	Move all parameters to the target device in-place.
.summary()	Print a table of all parameters, their shapes, and total parameter count.
.set_train_mode()	Switch this module and all sub-modules to training mode.
.set_eval_mode()	Switch this module and all sub-modules to evaluation mode.