Documentation Index
Fetch the complete documentation index at: https://mintlify.com/MilesONerd/neurenix/llms.txt
Use this file to discover all available pages before exploring further.
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)
Core Methods
forward
def forward(self, *args, **kwargs) -> Any
Input arguments to the module.
Keyword arguments to the module.
Output of the forward pass.
call
def __call__(self, *args, **kwargs) -> Any
forward().
Parameter Management
register_parameter
def register_parameter(self, name: str, param: Optional[Tensor]) -> None
The name of the parameter.
The parameter tensor, or None to remove the parameter.
register_module
def register_module(self, name: str, module: Optional[Module]) -> None
The name of the submodule.
The submodule, or None to remove the submodule.
register_buffer
def register_buffer(self, name: str, tensor: Optional[Tensor]) -> None
The tensor to register as buffer, or None to remove the buffer.
parameters
def parameters(self) -> List[Tensor]
A list of all parameter tensors.
Training Mode
train
def train(self, mode: bool = True) -> Module
Whether to set training mode (True) or evaluation mode (False).
eval
Set the module in evaluation mode.
is_training
def is_training(self) -> bool
True if the module is in training mode, False otherwise.
Device Management
def to(self, device: Device) -> Module
clone
def clone(self) -> 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.