openarchx 0.1.0__py3-none-any.whl

openarchx/layers/cnn.py

@@ -0,0 +1,145 @@
+import numpy as np
+from ..core.tensor import Tensor
+from ..nn.module import Module
+
+class Conv2d(Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
+        self.stride = stride if isinstance(stride, tuple) else (stride, stride)
+        self.padding = padding if isinstance(padding, tuple) else (padding, padding)
+        
+        # Initialize weights using He initialization
+        scale = np.sqrt(2.0 / (in_channels * self.kernel_size[0] * self.kernel_size[1]))
+        self.weight = Tensor(
+            np.random.normal(0, scale, 
+                           (out_channels, in_channels, *self.kernel_size)),
+            requires_grad=True
+        )
+        self.bias = Tensor(np.zeros(out_channels), requires_grad=True)
+    
+    def _extract_patches(self, x, k_h, k_w, stride_h, stride_w):
+        """Extract patches from input tensor efficiently"""
+        batch_size, channels, height, width = x.shape
+        
+        # Calculate output dimensions
+        out_h = (height - k_h) // stride_h + 1
+        out_w = (width - k_w) // stride_w + 1
+        
+        # Initialize patches array
+        patches = np.zeros((batch_size, out_h * out_w, channels * k_h * k_w))
+        
+        # Extract patches
+        patch_idx = 0
+        for i in range(0, height - k_h + 1, stride_h):
+            for j in range(0, width - k_w + 1, stride_w):
+                # Extract patch for all batches and channels
+                patch = x[:, :, i:i+k_h, j:j+k_w]
+                # Reshape patch to (batch_size, channels * k_h * k_w)
+                patches[:, patch_idx, :] = patch.reshape(batch_size, -1)
+                patch_idx += 1
+        
+        return patches, out_h, out_w
+    
+    def forward(self, x):
+        batch_size, C, H, W = x.data.shape
+        pad_h, pad_w = self.padding
+        stride_h, stride_w = self.stride
+        k_h, k_w = self.kernel_size
+        
+        # Add padding if needed
+        if pad_h > 0 or pad_w > 0:
+            x_padded = np.pad(x.data, ((0,0), (0,0), (pad_h,pad_h), (pad_w,pad_w)), mode='constant')
+        else:
+            x_padded = x.data
+        
+        # Extract patches
+        patches, H_out, W_out = self._extract_patches(x_padded, k_h, k_w, stride_h, stride_w)
+        
+        # Reshape weights to [out_channels, in_channels * k_h * k_w]
+        w_reshaped = self.weight.data.reshape(self.out_channels, -1)
+        
+        # Compute convolution using matrix multiplication
+        out = patches @ w_reshaped.T  # [batch_size, H_out * W_out, out_channels]
+        out = out.transpose(0, 2, 1).reshape(batch_size, self.out_channels, H_out, W_out)
+        
+        # Add bias
+        out += self.bias.data.reshape(1, -1, 1, 1)
+        
+        return Tensor(out, requires_grad=True)
+
+class MaxPool2d(Module):
+    def __init__(self, kernel_size, stride=None):
+        super().__init__()
+        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
+        self.stride = self.kernel_size if stride is None else (stride if isinstance(stride, tuple) else (stride, stride))
+    
+    def forward(self, x):
+        batch_size, C, H, W = x.data.shape
+        k_h, k_w = self.kernel_size
+        stride_h, stride_w = self.stride
+        
+        # Calculate output dimensions
+        H_out = (H - k_h) // stride_h + 1
+        W_out = (W - k_w) // stride_w + 1
+        
+        # Initialize output array
+        out = np.zeros((batch_size, C, H_out, W_out))
+        
+        # Perform max pooling
+        for b in range(batch_size):
+            for c in range(C):
+                for h in range(H_out):
+                    for w in range(W_out):
+                        h_start = h * stride_h
+                        w_start = w * stride_w
+                        h_end = h_start + k_h
+                        w_end = w_start + k_w
+                        
+                        pool_region = x.data[b, c, h_start:h_end, w_start:w_end]
+                        out[b, c, h, w] = np.max(pool_region)
+        
+        return Tensor(out, requires_grad=True)
+
+class BatchNorm2d(Module):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1):
+        super().__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        
+        # Parameters
+        self.gamma = Tensor(np.ones(num_features), requires_grad=True)
+        self.beta = Tensor(np.zeros(num_features), requires_grad=True)
+        
+        # Running estimates
+        self.running_mean = np.zeros(num_features)
+        self.running_var = np.ones(num_features)
+        
+        # Training mode flag
+        self.training = True
+    
+    def forward(self, x):
+        if self.training:
+            # Calculate batch statistics
+            batch_mean = x.data.mean(axis=(0,2,3), keepdims=True)
+            batch_var = x.data.var(axis=(0,2,3), keepdims=True)
+            
+            # Update running statistics
+            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * batch_mean.squeeze()
+            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var.squeeze()
+            
+            # Normalize
+            x_normalized = (x.data - batch_mean) / np.sqrt(batch_var + self.eps)
+        else:
+            # Use running statistics
+            x_normalized = (x.data - self.running_mean.reshape(1,-1,1,1)) / \
+                         np.sqrt(self.running_var.reshape(1,-1,1,1) + self.eps)
+        
+        # Apply scale and shift
+        out = self.gamma.data.reshape(1,-1,1,1) * x_normalized + \
+              self.beta.data.reshape(1,-1,1,1)
+        
+        return Tensor(out, requires_grad=True)

openarchx/layers/transformer.py

@@ -0,0 +1,131 @@
+import numpy as np
+from ..core.tensor import Tensor
+from ..nn.module import Module
+from .base import Linear
+
+def get_positional_encoding(seq_length, d_model):
+    """Generate positional encodings for transformer input"""
+    position = np.arange(seq_length)[:, np.newaxis]
+    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))
+    
+    pos_encoding = np.zeros((seq_length, d_model))
+    pos_encoding[:, 0::2] = np.sin(position * div_term)
+    pos_encoding[:, 1::2] = np.cos(position * div_term)
+    
+    return Tensor(pos_encoding[np.newaxis, :, :])  # Add batch dimension
+
+class PositionalEncoding(Module):
+    def __init__(self, d_model, max_seq_length=5000):
+        super().__init__()
+        self.pos_encoding = get_positional_encoding(max_seq_length, d_model)
+    
+    def forward(self, x):
+        return x + self.pos_encoding[:, :x.data.shape[1], :]
+
+class MultiHeadAttention(Module):
+    def __init__(self, embed_dim, num_heads):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        
+        self.q_proj = Linear(embed_dim, embed_dim)
+        self.k_proj = Linear(embed_dim, embed_dim)
+        self.v_proj = Linear(embed_dim, embed_dim)
+        self.out_proj = Linear(embed_dim, embed_dim)
+    
+    def split_heads(self, x, batch_size):
+        # [batch_size, seq_len, embed_dim] -> [batch_size, seq_len, num_heads, head_dim]
+        new_shape = (batch_size, -1, self.num_heads, self.head_dim)
+        x = x.reshape(*new_shape)
+        # [batch_size, seq_len, num_heads, head_dim] -> [batch_size, num_heads, seq_len, head_dim]
+        return x.transpose(0, 2, 1, 3)
+    
+    def merge_heads(self, x, batch_size, seq_len):
+        # [batch_size, num_heads, seq_len, head_dim] -> [batch_size, seq_len, num_heads, head_dim]
+        x = x.transpose(0, 2, 1, 3)
+        # [batch_size, seq_len, num_heads, head_dim] -> [batch_size, seq_len, embed_dim]
+        return x.reshape(batch_size, seq_len, self.embed_dim)
+    
+    def forward(self, query, key, value, mask=None):
+        batch_size = query.data.shape[0]
+        q_len, k_len = query.data.shape[1], key.data.shape[1]
+        
+        # Linear projections and split heads
+        q = self.split_heads(self.q_proj.forward(query), batch_size)  # [batch, heads, q_len, head_dim]
+        k = self.split_heads(self.k_proj.forward(key), batch_size)    # [batch, heads, k_len, head_dim]
+        v = self.split_heads(self.v_proj.forward(value), batch_size)  # [batch, heads, v_len, head_dim]
+        
+        # Scaled dot-product attention
+        # [batch, heads, q_len, head_dim] @ [batch, heads, head_dim, k_len]
+        scores = (q @ k.transpose(0, 1, 3, 2)) / np.sqrt(self.head_dim)
+        
+        if mask is not None:
+            scores.data = scores.data + mask.data * -1e9
+        
+        # Apply softmax and attention
+        attn = self._softmax(scores)  # [batch, heads, q_len, k_len]
+        out = attn @ v  # [batch, heads, q_len, head_dim]
+        
+        # Merge heads and project
+        out = self.merge_heads(out, batch_size, q_len)  # [batch, q_len, embed_dim]
+        return self.out_proj.forward(out)
+    
+    def _softmax(self, x):
+        exp_x = Tensor(np.exp(x.data - np.max(x.data, axis=-1, keepdims=True)))
+        return exp_x / exp_x.sum(axis=-1, keepdims=True)
+
+class LayerNorm(Module):
+    def __init__(self, normalized_shape, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.gamma = Tensor(np.ones(normalized_shape), requires_grad=True)
+        self.beta = Tensor(np.zeros(normalized_shape), requires_grad=True)
+    
+    def forward(self, x):
+        mean = x.mean(axis=-1, keepdims=True)
+        var = ((x - mean) ** 2).mean(axis=-1, keepdims=True)
+        return self.gamma * (x - mean) / (var + self.eps).sqrt() + self.beta
+    
+    def parameters(self):
+        return [self.gamma, self.beta]
+
+class TransformerEncoderLayer(Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1):
+        super().__init__()
+        self.self_attn = MultiHeadAttention(d_model, nhead)
+        self.pos_encoding = PositionalEncoding(d_model)
+        self.linear1 = Linear(d_model, dim_feedforward)
+        self.linear2 = Linear(dim_feedforward, d_model)
+        self.norm1 = LayerNorm(d_model)
+        self.norm2 = LayerNorm(d_model)
+        self.dropout = dropout
+    
+    def forward(self, src, src_mask=None):
+        # Add positional encoding
+        src = self.pos_encoding.forward(src)
+        
+        # Multi-head self-attention
+        attn_output = self.self_attn.forward(src, src, src, mask=src_mask)
+        attn_output = self._dropout(attn_output)
+        out1 = self.norm1.forward(src + attn_output)
+        
+        # Position-wise feed-forward network
+        ff_output = self.linear1.forward(out1)
+        ff_output = self._relu(ff_output)
+        ff_output = self._dropout(ff_output)
+        ff_output = self.linear2.forward(ff_output)
+        ff_output = self._dropout(ff_output)
+        
+        return self.norm2.forward(out1 + ff_output)
+    
+    def _dropout(self, x):
+        if self.dropout > 0:
+            mask = np.random.binomial(1, 1-self.dropout, x.data.shape)
+            return Tensor(x.data * mask / (1-self.dropout))
+        return x
+    
+    def _relu(self, x):
+        return Tensor(np.maximum(0, x.data))

openarchx/nn/__init__.py

@@ -0,0 +1,26 @@
+# Core module
+from .module import Module
+
+# Activation functions
+from .activations import (
+    ReLU, LeakyReLU, Sigmoid, Tanh, GELU, SiLU, 
+    ELU, SELU, Softmax, LogSoftmax
+)
+
+# Core layers
+from .layers import (
+    Linear, Conv1d, Conv2d, LayerNorm,
+    Embedding
+)
+
+# Pooling layers
+from .pooling import (
+    MaxPool1d, MaxPool2d, AvgPool1d, AvgPool2d,
+    AdaptiveAvgPool2d, AdaptiveMaxPool2d
+)
+
+# Transformer components
+from ..layers.transformer import PositionalEncoding
+
+# Container modules
+from .containers import Sequential, ModuleList

openarchx/nn/activations.py

@@ -0,0 +1,127 @@
+import numpy as np
+from ..core.tensor import Tensor
+from .module import Module
+
+class ReLU(Module):
+    def forward(self, x):
+        return Tensor(np.maximum(0, x.data), requires_grad=True)
+
+class LeakyReLU(Module):
+    def __init__(self, negative_slope=0.01):
+        super().__init__()
+        self.negative_slope = negative_slope
+
+    def forward(self, x):
+        return Tensor(np.where(x.data > 0, x.data, self.negative_slope * x.data), requires_grad=True)
+
+class PReLU(Module):
+    def __init__(self, num_parameters=1, init=0.25):
+        super().__init__()
+        self.weight = Tensor(np.full(num_parameters, init), requires_grad=True)
+
+    def forward(self, x):
+        return Tensor(np.where(x.data > 0, x.data, self.weight.data * x.data), requires_grad=True)
+
+class ELU(Module):
+    def __init__(self, alpha=1.0):
+        super().__init__()
+        self.alpha = alpha
+
+    def forward(self, x):
+        return Tensor(np.where(x.data > 0, x.data, self.alpha * (np.exp(x.data) - 1)), requires_grad=True)
+
+class GELU(Module):
+    def forward(self, x):
+        return Tensor(0.5 * x.data * (1 + np.tanh(np.sqrt(2 / np.pi) * (x.data + 0.044715 * x.data**3))), requires_grad=True)
+
+class Sigmoid(Module):
+    def forward(self, x):
+        return Tensor(1 / (1 + np.exp(-x.data)), requires_grad=True)
+
+class Tanh(Module):
+    def forward(self, x):
+        return Tensor(np.tanh(x.data), requires_grad=True)
+
+class Softmax(Module):
+    def __init__(self, dim=-1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        exp_x = np.exp(x.data - np.max(x.data, axis=self.dim, keepdims=True))
+        return Tensor(exp_x / np.sum(exp_x, axis=self.dim, keepdims=True), requires_grad=True)
+
+class LogSoftmax(Module):
+    def __init__(self, dim=-1):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x):
+        exp_x = np.exp(x.data - np.max(x.data, axis=self.dim, keepdims=True))
+        softmax = exp_x / np.sum(exp_x, axis=self.dim, keepdims=True)
+        return Tensor(np.log(softmax), requires_grad=True)
+
+class SELU(Module):
+    def __init__(self):
+        super().__init__()
+        self.alpha = 1.6732632423543772848170429916717
+        self.scale = 1.0507009873554804934193349852946
+
+    def forward(self, x):
+        return Tensor(self.scale * np.where(x.data > 0, x.data, 
+                                          self.alpha * (np.exp(x.data) - 1)), requires_grad=True)
+
+class Hardtanh(Module):
+    def __init__(self, min_val=-1.0, max_val=1.0):
+        super().__init__()
+        self.min_val = min_val
+        self.max_val = max_val
+
+    def forward(self, x):
+        return Tensor(np.clip(x.data, self.min_val, self.max_val), requires_grad=True)
+
+class SiLU(Module):  # Also known as Swish
+    def forward(self, x):
+        return Tensor(x.data * (1 / (1 + np.exp(-x.data))), requires_grad=True)
+
+class Mish(Module):
+    def forward(self, x):
+        return Tensor(x.data * np.tanh(np.log(1 + np.exp(x.data))), requires_grad=True)
+
+class ActX(Module):
+    """
+    Advanced activation function that combines multiple activation types with learnable parameters.
+    ActX(x) = α * GELU(x) + β * SiLU(x) + γ * tanh(λx)
+    where α, β, γ, and λ are learnable parameters
+    """
+    def __init__(self, num_parameters=1, init_alpha=0.5, init_beta=0.5, init_gamma=0.25, init_lambda=1.0):
+        super().__init__()
+        self.num_parameters = num_parameters
+        
+        # Initialize learnable parameters
+        self.alpha = Tensor(np.full(num_parameters, init_alpha), requires_grad=True)
+        self.beta = Tensor(np.full(num_parameters, init_beta), requires_grad=True)
+        self.gamma = Tensor(np.full(num_parameters, init_gamma), requires_grad=True)
+        self.lambda_param = Tensor(np.full(num_parameters, init_lambda), requires_grad=True)
+
+    def forward(self, x):
+        # GELU component
+        gelu = 0.5 * x.data * (1 + np.tanh(np.sqrt(2 / np.pi) * (x.data + 0.044715 * x.data**3)))
+        
+        # SiLU (Swish) component
+        silu = x.data * (1 / (1 + np.exp(-x.data)))
+        
+        # Tanh component with learnable frequency
+        tanh = np.tanh(self.lambda_param.data.reshape(-1, 1, 1) * x.data)
+        
+        # Combine components with learnable weights
+        alpha = self.alpha.data.reshape(-1, 1, 1)
+        beta = self.beta.data.reshape(-1, 1, 1)
+        gamma = self.gamma.data.reshape(-1, 1, 1)
+        
+        result = alpha * gelu + beta * silu + gamma * tanh
+        
+        return Tensor(result, requires_grad=True)
+
+    def parameters(self):
+        return [self.alpha, self.beta, self.gamma, self.lambda_param]

openarchx/nn/containers.py

@@ -0,0 +1,174 @@
+import numpy as np
+from ..core.tensor import Tensor
+from .module import Module
+from collections import OrderedDict
+
+class Sequential(Module):
+    def __init__(self, *args):
+        super().__init__()
+        self.modules = []
+        for arg in args:
+            if isinstance(arg, (list, tuple)):
+                self.modules.extend(arg)
+            elif isinstance(arg, dict):
+                self.modules.extend(arg.values())
+            else:
+                self.modules.append(arg)
+
+    def forward(self, x):
+        for module in self.modules:
+            x = module(x)
+        return x
+
+    def parameters(self):
+        params = []
+        for module in self.modules:
+            params.extend(module.parameters())
+        return params
+
+class ModuleList(Module):
+    def __init__(self, modules=None):
+        super().__init__()
+        self.modules = []
+        if modules is not None:
+            self.extend(modules)
+
+    def __getitem__(self, idx):
+        return self.modules[idx]
+
+    def __setitem__(self, idx, module):
+        self.modules[idx] = module
+
+    def __len__(self):
+        return len(self.modules)
+
+    def append(self, module):
+        self.modules.append(module)
+
+    def extend(self, modules):
+        if isinstance(modules, (list, tuple)):
+            self.modules.extend(modules)
+        else:
+            self.modules.extend(list(modules))
+
+    def parameters(self):
+        params = []
+        for module in self.modules:
+            params.extend(module.parameters())
+        return params
+
+class ModuleDict(Module):
+    def __init__(self, modules=None):
+        super().__init__()
+        self.modules = OrderedDict()
+        if modules is not None:
+            self.update(modules)
+
+    def __getitem__(self, key):
+        return self.modules[key]
+
+    def __setitem__(self, key, module):
+        self.modules[key] = module
+
+    def __delitem__(self, key):
+        del self.modules[key]
+
+    def __len__(self):
+        return len(self.modules)
+
+    def __iter__(self):
+        return iter(self.modules)
+
+    def keys(self):
+        return self.modules.keys()
+
+    def items(self):
+        return self.modules.items()
+
+    def values(self):
+        return self.modules.values()
+
+    def update(self, modules):
+        if isinstance(modules, dict):
+            self.modules.update(modules)
+        else:
+            for key, module in modules:
+                self.modules[key] = module
+
+    def parameters(self):
+        params = []
+        for module in self.modules.values():
+            params.extend(module.parameters())
+        return params
+
+class ParameterList(Module):
+    def __init__(self, parameters=None):
+        super().__init__()
+        self.parameters_list = []
+        if parameters is not None:
+            self.extend(parameters)
+
+    def __getitem__(self, idx):
+        return self.parameters_list[idx]
+
+    def __setitem__(self, idx, parameter):
+        self.parameters_list[idx] = parameter
+
+    def __len__(self):
+        return len(self.parameters_list)
+
+    def append(self, parameter):
+        if not isinstance(parameter, Tensor):
+            parameter = Tensor(parameter, requires_grad=True)
+        self.parameters_list.append(parameter)
+
+    def extend(self, parameters):
+        for param in parameters:
+            self.append(param)
+
+    def parameters(self):
+        return self.parameters_list
+
+class ParameterDict(Module):
+    def __init__(self, parameters=None):
+        super().__init__()
+        self.parameters_dict = OrderedDict()
+        if parameters is not None:
+            self.update(parameters)
+
+    def __getitem__(self, key):
+        return self.parameters_dict[key]
+
+    def __setitem__(self, key, parameter):
+        if not isinstance(parameter, Tensor):
+            parameter = Tensor(parameter, requires_grad=True)
+        self.parameters_dict[key] = parameter
+
+    def __delitem__(self, key):
+        del self.parameters_dict[key]
+
+    def __len__(self):
+        return len(self.parameters_dict)
+
+    def __iter__(self):
+        return iter(self.parameters_dict)
+
+    def keys(self):
+        return self.parameters_dict.keys()
+
+    def items(self):
+        return self.parameters_dict.items()
+
+    def values(self):
+        return self.parameters_dict.values()
+
+    def update(self, parameters):
+        if isinstance(parameters, dict):
+            for key, param in parameters.items():
+                self[key] = param
+        else:
+            for key, param in parameters:
+                self[key] = param
+
+    def parameters(self):
+        return list(self.parameters_dict.values())