Neural Networks

from neurenix.nn import Module, Linear, ReLU, Sequential
from neurenix.tensor import Tensor

class SimpleNet(Module):
    def __init__(self):
        super().__init__()
        self.fc1 = Linear(784, 256)
        self.relu = ReLU()
        self.fc2 = Linear(256, 10)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# Create and use the model
model = SimpleNet()
output = model(Tensor.randn((32, 784)))

from neurenix.nn import Module
from neurenix.tensor import Tensor
import numpy as np

class CustomLayer(Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        
        # Register parameters
        weight_data = np.random.randn(out_features, in_features) * 0.01
        self.weight = Tensor(weight_data, requires_grad=True)
        self.register_parameter('weight', self.weight)
        
        bias_data = np.zeros(out_features)
        self.bias = Tensor(bias_data, requires_grad=True)
        self.register_parameter('bias', self.bias)
    
    def forward(self, x):
        return x.matmul(self.weight.transpose(0, 1)) + self.bias

from neurenix.nn import Linear
from neurenix.tensor import Tensor

# Create linear layer
layer = Linear(
    in_features=784,   # Input size
    out_features=128,  # Output size
    bias=True          # Include bias term (default)
)

# Forward pass
x = Tensor.randn((32, 784))  # Batch of 32 samples
output = layer(x)             # Shape: (32, 128)

print(f"Weight shape: {layer.weight.shape}")  # (128, 784)
print(f"Bias shape: {layer.bias.shape}")      # (128,)

from neurenix.nn import Linear
from neurenix.tensor import DType
from neurenix.device import Device, DeviceType

layer = Linear(
    in_features=512,
    out_features=256,
    bias=True,
    dtype=DType.FLOAT64,
    device=Device(DeviceType.CUDA, 0)
)

from neurenix.nn import ReLU, LeakyReLU
from neurenix.tensor import Tensor

# In-place operations save memory
relu_inplace = ReLU(inplace=True)
leaky_inplace = LeakyReLU(negative_slope=0.2, inplace=True)

x = Tensor.randn((100, 100))
relu_inplace(x)  # Modifies x in-place

from neurenix.nn import Conv2d
from neurenix.tensor import Tensor

# Create 2D convolution layer
conv = Conv2d(
    in_channels=3,      # RGB input
    out_channels=64,    # 64 feature maps
    kernel_size=3,      # 3x3 kernel
    stride=1,           # Stride of 1
    padding=1           # Same padding
)

# Process batch of images
images = Tensor.randn((32, 3, 224, 224))  # Batch, Channels, Height, Width
features = conv(images)                    # Shape: (32, 64, 224, 224)

print(f"Output shape: {features.shape}")
print(f"Weight shape: {conv.weight.shape}")  # (64, 3, 3, 3)

from neurenix.nn import Conv2d

# Strided convolution (downsampling)
strided_conv = Conv2d(
    in_channels=64,
    out_channels=128,
    kernel_size=3,
    stride=2,        # Halves spatial dimensions
    padding=1
)

# Dilated convolution (expanded receptive field)
dilated_conv = Conv2d(
    in_channels=64,
    out_channels=64,
    kernel_size=3,
    stride=1,
    padding=2,
    dilation=2       # Increases receptive field
)

# Grouped convolution (efficient)
grouped_conv = Conv2d(
    in_channels=64,
    out_channels=64,
    kernel_size=3,
    stride=1,
    padding=1,
    groups=8         # 8 separate convolutions
)

from neurenix.nn import MSELoss
from neurenix.tensor import Tensor

preds = Tensor([[1, 2], [3, 4]])
targets = Tensor([[2, 2], [3, 3]])

# No reduction - returns loss per element
loss_none = MSELoss(reduction='none')
print(loss_none(preds, targets))  # Shape: (2, 2)

# Mean reduction
loss_mean = MSELoss(reduction='mean')
print(loss_mean(preds, targets))  # Scalar

# Sum reduction
loss_sum = MSELoss(reduction='sum')
print(loss_sum(preds, targets))   # Scalar

from neurenix.nn import Sequential, Linear, ReLU, CrossEntropyLoss
from neurenix.tensor import Tensor
from neurenix.device import Device, DeviceType
from neurenix.optim import Adam

# Define model
model = Sequential(
    Linear(784, 256),
    ReLU(),
    Linear(256, 128),
    ReLU(),
    Linear(128, 10)
)

# Move to GPU if available
from neurenix.device import get_device_count
if get_device_count(DeviceType.CUDA) > 0:
    device = Device(DeviceType.CUDA, 0)
    model.to(device)
else:
    device = Device(DeviceType.CPU)

# Setup training
loss_fn = CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=0.001)

# Training loop
for epoch in range(10):
    model.train()
    
    # Get batch (placeholder)
    inputs = Tensor.randn((32, 784), device=device)
    targets = Tensor([i % 10 for i in range(32)], dtype="int64", device=device)
    
    # Forward pass
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
    
    # Backward pass
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    
    if epoch % 2 == 0:
        print(f"Epoch {epoch}, Loss: {loss.numpy():.4f}")

# Evaluation
model.eval()
with Tensor.no_grad():
    test_inputs = Tensor.randn((10, 784), device=device)
    predictions = model(test_inputs)
    print(f"Predictions shape: {predictions.shape}")

from neurenix.nn import Sequential, Linear, ReLU

model = Sequential(
    Linear(100, 50),
    ReLU(),
    Linear(50, 10)
)

# Count parameters
total_params = sum(p.size for p in model.parameters())
print(f"Total parameters: {total_params}")

# Iterate through modules
for i, module in enumerate(model):
    print(f"Layer {i}: {module}")

# Get parameter shapes
for i, param in enumerate(model.parameters()):
    print(f"Parameter {i} shape: {param.shape}")