Module

class Module:
    def __init__(self)

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

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

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

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

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

def parameters(self) -> List[Tensor]

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

def eval(self) -> Module

def is_training(self) -> bool

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

def clone(self) -> Module

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()