Module

Base class for neural network modules

Overview

The Module class is the base class for all neural network modules in Neurenix. It provides parameter management, training/evaluation modes, and hierarchical organization of submodules.

Class Definition

class Module:
    def __init__(self)

Initialize a new module.

Core Methods

forward

def forward(self, *args, **kwargs) -> Any

Forward pass of the module. This method should be overridden by all subclasses.

*args · Any

Input arguments to the module.

**kwargs · Any

Keyword arguments to the module.

return · Any

Output of the forward pass.

call

def __call__(self, *args, **kwargs) -> Any

Call the module as a function. Internally calls forward().

Parameter Management

register_parameter

def register_parameter(self, name: str, param: Optional[Tensor]) -> None

Register a parameter with the module.

name · str · required

The name of the parameter.

param · Optional[Tensor] · required

The parameter tensor, or None to remove the parameter.

register_module

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

Register a submodule with the module.

name · str · required

The name of the submodule.

module · Optional[Module] · required

The submodule, or None to remove the submodule.

register_buffer

def register_buffer(self, name: str, tensor: Optional[Tensor]) -> None

Register a buffer with the module. Buffers are module states that should be saved along with parameters but are not parameters (e.g., running mean in batch normalization).

name · str · required

The name of the buffer.

tensor · Optional[Tensor] · required

The tensor to register as buffer, or None to remove the buffer.

parameters

def parameters(self) -> List[Tensor]

Get all parameters of the module and its submodules.

return · List[Tensor]

A list of all parameter tensors.

Training Mode

train

def train(self, mode: bool = True) -> Module

Set the module in training mode.

mode · bool · default: "True"

Whether to set training mode (True) or evaluation mode (False).

return · Module

The module itself.

eval

def eval(self) -> Module

Set the module in evaluation mode.

return · Module

The module itself.

is_training

def is_training(self) -> bool

Check if the module is in training mode.

return · bool

True if the module is in training mode, False otherwise.

Device Management

to

def to(self, device: Device) -> Module

Move the module and its parameters to the specified device.

device · Device · required

The device to move to.

return · Module

The module itself.

clone

def clone(self) -> Module

Create a clone of this module with the same parameters.

return · Module

A new module with the same parameters.

Example Usage

import neurenix as nx
from neurenix.nn import Module, Linear

# Define a custom module
class MyNetwork(Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        
        # Register submodules
        self.fc1 = Linear(input_size, hidden_size)
        self.fc2 = Linear(hidden_size, output_size)
        
        # Register a buffer
        self.register_buffer('running_mean', nx.Tensor.zeros((hidden_size,)))
    
    def forward(self, x):
        x = self.fc1(x)
        x = x.relu()
        x = self.fc2(x)
        return x

# Create the network
model = MyNetwork(784, 256, 10)

# Get parameters
params = model.parameters()
print(f"Number of parameters: {len(params)}")

# Set training mode
model.train()
assert model.is_training()

# Set evaluation mode
model.eval()
assert not model.is_training()

# Move to device
device = nx.Device(nx.DeviceType.CUDA, 0)
model.to(device)

# Forward pass
input_tensor = nx.Tensor.randn((32, 784), device=device)
output = model(input_tensor)
print(output.shape)  # (32, 10)

# Clone the model
cloned_model = model.clone()

Built-in Modules

Neurenix provides many built-in modules:

• Linear layers: Linear
• Convolutional layers: Conv1d, Conv2d, Conv3d
• Recurrent layers: RNN, LSTM, GRU
• Activation functions: ReLU, Sigmoid, Tanh, Softmax
• Pooling layers: MaxPool2d
• Regularization: Dropout
• Containers: Sequential

See individual documentation pages for details on each module type.