openarchx 0.1.0__py3-none-any.whl

openarchx/nn/dropout.py

@@ -0,0 +1,121 @@
+import numpy as np
+from ..core.tensor import Tensor
+from .module import Module
+
+class Dropout(Module):
+    def __init__(self, p=0.5, inplace=False):
+        super().__init__()
+        self.p = p
+        self.inplace = inplace
+        self.training = True
+
+    def forward(self, x):
+        if not self.training or self.p == 0:
+            return x
+        
+        mask = np.random.random(x.data.shape) > self.p
+        if self.inplace:
+            x.data = x.data * mask / (1 - self.p)
+            return x
+        return Tensor(x.data * mask / (1 - self.p), requires_grad=True)
+
+class Dropout2d(Module):
+    def __init__(self, p=0.5, inplace=False):
+        super().__init__()
+        self.p = p
+        self.inplace = inplace
+        self.training = True
+
+    def forward(self, x):
+        if not self.training or self.p == 0:
+            return x
+        
+        # Create mask for entire channels
+        mask = np.random.random((x.data.shape[0], x.data.shape[1], 1, 1)) > self.p
+        mask = np.broadcast_to(mask, x.data.shape)
+        
+        if self.inplace:
+            x.data = x.data * mask / (1 - self.p)
+            return x
+        return Tensor(x.data * mask / (1 - self.p), requires_grad=True)
+
+class Dropout3d(Module):
+    def __init__(self, p=0.5, inplace=False):
+        super().__init__()
+        self.p = p
+        self.inplace = inplace
+        self.training = True
+
+    def forward(self, x):
+        if not self.training or self.p == 0:
+            return x
+        
+        # Create mask for entire 3D feature maps
+        mask = np.random.random((x.data.shape[0], x.data.shape[1], 1, 1, 1)) > self.p
+        mask = np.broadcast_to(mask, x.data.shape)
+        
+        if self.inplace:
+            x.data = x.data * mask / (1 - self.p)
+            return x
+        return Tensor(x.data * mask / (1 - self.p), requires_grad=True)
+
+class AlphaDropout(Module):
+    def __init__(self, p=0.5, inplace=False):
+        super().__init__()
+        self.p = p
+        self.inplace = inplace
+        self.training = True
+        # SELU parameters
+        self.alpha = 1.6732632423543772848170429916717
+        self.scale = 1.0507009873554804934193349852946
+        
+        self.alpha_p = -self.alpha * self.scale
+
+    def forward(self, x):
+        if not self.training or self.p == 0:
+            return x
+        
+        # Keep mean and variance the same during training and evaluation
+        mask = np.random.random(x.data.shape) > self.p
+        
+        # Calculate the affine transformation parameters
+        a = ((1 - self.p) + self.p * self.alpha_p ** 2) ** (-0.5)
+        b = -a * self.p * self.alpha_p
+        
+        if self.inplace:
+            x.data = mask * x.data + (1 - mask) * self.alpha_p
+            x.data = a * x.data + b
+            return x
+        return Tensor(a * (mask * x.data + (1 - mask) * self.alpha_p) + b, requires_grad=True)
+
+class FeatureAlphaDropout(Module):
+    def __init__(self, p=0.5, inplace=False):
+        super().__init__()
+        self.p = p
+        self.inplace = inplace
+        self.training = True
+        # SELU parameters
+        self.alpha = 1.6732632423543772848170429916717
+        self.scale = 1.0507009873554804934193349852946
+        
+        self.alpha_p = -self.alpha * self.scale
+
+    def forward(self, x):
+        if not self.training or self.p == 0:
+            return x
+        
+        # Create mask for entire features
+        shape = list(x.data.shape)
+        shape[1:] = [1] * (len(shape) - 1)
+        mask = np.random.random(shape) > self.p
+        mask = np.broadcast_to(mask, x.data.shape)
+        
+        # Calculate the affine transformation parameters
+        a = ((1 - self.p) + self.p * self.alpha_p ** 2) ** (-0.5)
+        b = -a * self.p * self.alpha_p
+        
+        if self.inplace:
+            x.data = mask * x.data + (1 - mask) * self.alpha_p
+            x.data = a * x.data + b
+            return x
+        return Tensor(a * (mask * x.data + (1 - mask) * self.alpha_p) + b, requires_grad=True)

openarchx/nn/layers.py

@@ -0,0 +1,338 @@
+import numpy as np
+from ..core.tensor import Tensor
+from .module import Module
+
+class Linear(Module):
+    def __init__(self, in_features, out_features, bias=True):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        scale = np.sqrt(2.0 / in_features)
+        self.weight = Tensor(np.random.normal(0, scale, (in_features, out_features)), requires_grad=True)
+        self.bias = Tensor(np.zeros(out_features), requires_grad=True) if bias else None
+
+    def forward(self, x):
+        if not isinstance(x, Tensor):
+            x = Tensor(x, requires_grad=True)
+        out = x @ self.weight
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+class Conv1d(Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=True):
+        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,)
+        self.stride = stride if isinstance(stride, tuple) else (stride,)
+        self.padding = padding if isinstance(padding, tuple) else (padding,)
+        
+        scale = np.sqrt(2.0 / (in_channels * kernel_size))
+        self.weight = Tensor(
+            np.random.normal(0, scale, (out_channels, in_channels, kernel_size)),
+            requires_grad=True
+        )
+        self.bias = Tensor(np.zeros(out_channels), requires_grad=True) if bias else None
+
+    def forward(self, x):
+        # Implementation for 1D convolution
+        pass  # TODO: Implement actual convolution
+
+class Conv2d(Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=True):
+        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)
+        
+        scale = np.sqrt(2.0 / (in_channels * kernel_size * kernel_size))
+        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) if bias else None
+
+    def forward(self, x):
+        # Implementation for 2D convolution
+        pass  # TODO: Implement actual convolution
+
+class BatchNorm1d(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
+        
+        self.gamma = Tensor(np.ones(num_features), requires_grad=True)
+        self.beta = Tensor(np.zeros(num_features), requires_grad=True)
+        self.running_mean = Tensor(np.zeros(num_features), requires_grad=False)
+        self.running_var = Tensor(np.ones(num_features), requires_grad=False)
+
+    def forward(self, x):
+        if self.training:
+            mean = x.mean(axis=0)
+            var = x.var(axis=0)
+            
+            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean
+            self.running_var = (1 - self.momentum) * self.running_var + self.momentum * var
+        else:
+            mean = self.running_mean
+            var = self.running_var
+        
+        x_norm = (x - mean) / np.sqrt(var + self.eps)
+        return self.gamma * x_norm + self.beta
+
+class LayerNorm(Module):
+    def __init__(self, normalized_shape, eps=1e-5):
+        super().__init__()
+        self.normalized_shape = normalized_shape
+        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.var(axis=-1, keepdims=True)
+        return self.gamma * (x - mean) / np.sqrt(var + self.eps) + self.beta
+
+class Embedding(Module):
+    def __init__(self, num_embeddings, embedding_dim):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.weight = Tensor(
+            np.random.normal(0, 0.02, (num_embeddings, embedding_dim)),
+            requires_grad=True
+        )
+
+    def forward(self, x):
+        return self.weight[x]
+
+class ConvTranspose1d(Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, bias=True):
+        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,)
+        self.stride = stride if isinstance(stride, tuple) else (stride,)
+        self.padding = padding if isinstance(padding, tuple) else (padding,)
+        self.output_padding = output_padding if isinstance(output_padding, tuple) else (output_padding,)
+        
+        scale = np.sqrt(2.0 / (in_channels * kernel_size))
+        self.weight = Tensor(
+            np.random.normal(0, scale, (in_channels, out_channels, *self.kernel_size)),
+            requires_grad=True
+        )
+        if bias:
+            self.bias = Tensor(np.zeros(out_channels), requires_grad=True)
+        else:
+            self.bias = None
+
+    def forward(self, x):
+        # TODO: Implement actual transposed convolution
+        pass
+
+class ConvTranspose2d(Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, bias=True):
+        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)
+        self.output_padding = output_padding if isinstance(output_padding, tuple) else (output_padding, output_padding)
+        
+        scale = np.sqrt(2.0 / (in_channels * kernel_size * kernel_size))
+        self.weight = Tensor(
+            np.random.normal(0, scale, (in_channels, out_channels, *self.kernel_size)),
+            requires_grad=True
+        )
+        if bias:
+            self.bias = Tensor(np.zeros(out_channels), requires_grad=True)
+        else:
+            self.bias = None
+
+    def forward(self, x):
+        # TODO: Implement actual transposed convolution
+        pass
+
+class ConvTranspose3d(Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, output_padding=0, bias=True):
+        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, kernel_size)
+        self.stride = stride if isinstance(stride, tuple) else (stride, stride, stride)
+        self.padding = padding if isinstance(padding, tuple) else (padding, padding, padding)
+        self.output_padding = output_padding if isinstance(output_padding, tuple) else (output_padding, output_padding, output_padding)
+        
+        scale = np.sqrt(2.0 / (in_channels * kernel_size * kernel_size * kernel_size))
+        self.weight = Tensor(
+            np.random.normal(0, scale, (in_channels, out_channels, *self.kernel_size)),
+            requires_grad=True
+        )
+        if bias:
+            self.bias = Tensor(np.zeros(out_channels), requires_grad=True)
+        else:
+            self.bias = None
+
+    def forward(self, x):
+        # TODO: Implement actual transposed convolution
+        pass
+
+class InstanceNorm1d(Module):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
+        super().__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        
+        if affine:
+            self.weight = Tensor(np.ones(num_features), requires_grad=True)
+            self.bias = Tensor(np.zeros(num_features), requires_grad=True)
+
+    def forward(self, x):
+        mean = np.mean(x.data, axis=(2,), keepdims=True)
+        var = np.var(x.data, axis=(2,), keepdims=True)
+        
+        x_norm = (x.data - mean) / np.sqrt(var + self.eps)
+        
+        if self.affine:
+            x_norm = self.weight.data.reshape(-1, 1) * x_norm + self.bias.data.reshape(-1, 1)
+        
+        return Tensor(x_norm, requires_grad=True)
+
+class InstanceNorm2d(Module):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
+        super().__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        
+        if affine:
+            self.weight = Tensor(np.ones(num_features), requires_grad=True)
+            self.bias = Tensor(np.zeros(num_features), requires_grad=True)
+
+    def forward(self, x):
+        mean = np.mean(x.data, axis=(2, 3), keepdims=True)
+        var = np.var(x.data, axis=(2, 3), keepdims=True)
+        
+        x_norm = (x.data - mean) / np.sqrt(var + self.eps)
+        
+        if self.affine:
+            x_norm = self.weight.data.reshape(-1, 1, 1) * x_norm + self.bias.data.reshape(-1, 1, 1)
+        
+        return Tensor(x_norm, requires_grad=True)
+
+class InstanceNorm3d(Module):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
+        super().__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.momentum = momentum
+        self.affine = affine
+        
+        if affine:
+            self.weight = Tensor(np.ones(num_features), requires_grad=True)
+            self.bias = Tensor(np.zeros(num_features), requires_grad=True)
+
+    def forward(self, x):
+        mean = np.mean(x.data, axis=(2, 3, 4), keepdims=True)
+        var = np.var(x.data, axis=(2, 3, 4), keepdims=True)
+        
+        x_norm = (x.data - mean) / np.sqrt(var + self.eps)
+        
+        if self.affine:
+            x_norm = self.weight.data.reshape(-1, 1, 1, 1) * x_norm + self.bias.data.reshape(-1, 1, 1, 1)
+        
+        return Tensor(x_norm, requires_grad=True)
+
+class GroupNorm(Module):
+    def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
+        super().__init__()
+        self.num_groups = num_groups
+        self.num_channels = num_channels
+        self.eps = eps
+        self.affine = affine
+        
+        if affine:
+            self.weight = Tensor(np.ones(num_channels), requires_grad=True)
+            self.bias = Tensor(np.zeros(num_channels), requires_grad=True)
+
+    def forward(self, x):
+        N, C, *spatial = x.data.shape
+        x_reshaped = x.data.reshape(N, self.num_groups, -1)
+        
+        mean = np.mean(x_reshaped, axis=2, keepdims=True)
+        var = np.var(x_reshaped, axis=2, keepdims=True)
+        
+        x_norm = (x_reshaped - mean) / np.sqrt(var + self.eps)
+        x_norm = x_norm.reshape(N, C, *spatial)
+        
+        if self.affine:
+            shape = (1, -1) + (1,) * len(spatial)
+            x_norm = self.weight.data.reshape(shape) * x_norm + self.bias.data.reshape(shape)
+        
+        return Tensor(x_norm, requires_grad=True)
+
+class LocalResponseNorm(Module):
+    def __init__(self, size, alpha=1e-4, beta=0.75, k=1.0):
+        super().__init__()
+        self.size = size
+        self.alpha = alpha
+        self.beta = beta
+        self.k = k
+
+    def forward(self, x):
+        N, C, *spatial = x.data.shape
+        half_size = self.size // 2
+        
+        square = x.data ** 2
+        for i in range(C):
+            start = max(0, i - half_size)
+            end = min(C, i + half_size + 1)
+            scale = self.k + self.alpha * np.sum(square[:, start:end], axis=1, keepdims=True)
+            x.data[:, i] /= scale ** self.beta
+        
+        return Tensor(x.data, requires_grad=True)
+
+class EmbeddingBag(Module):
+    def __init__(self, num_embeddings, embedding_dim, mode='mean'):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.mode = mode
+        self.weight = Tensor(
+            np.random.normal(0, 0.02, (num_embeddings, embedding_dim)),
+            requires_grad=True
+        )
+
+    def forward(self, x, offsets=None):
+        embeddings = self.weight.data[x.data]
+        if offsets is None:
+            # Treat entire input as a single bag
+            if self.mode == 'mean':
+                return Tensor(np.mean(embeddings, axis=1), requires_grad=True)
+            elif self.mode == 'sum':
+                return Tensor(np.sum(embeddings, axis=1), requires_grad=True)
+            else:  # max
+                return Tensor(np.max(embeddings, axis=1), requires_grad=True)
+        
+        # Handle multiple bags using offsets
+        result = []
+        for i in range(len(offsets) - 1):
+            start, end = offsets[i:i+2]
+            if self.mode == 'mean':
+                bag = np.mean(embeddings[start:end], axis=0)
+            elif self.mode == 'sum':
+                bag = np.sum(embeddings[start:end], axis=0)
+            else:  # max
+                bag = np.max(embeddings[start:end], axis=0)
+            result.append(bag)
+        
+        return Tensor(np.stack(result), requires_grad=True)

openarchx/nn/losses.py

@@ -0,0 +1,156 @@
+import numpy as np
+from ..core.tensor import Tensor
+from .module import Module
+
+class MSELoss(Module):
+    def __init__(self, reduction='mean'):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, pred, target):
+        if not isinstance(target, Tensor):
+            target = Tensor(target, requires_grad=False)
+        
+        loss = (pred - target) ** 2
+        if self.reduction == 'mean':
+            return Tensor(np.mean(loss.data), requires_grad=True)
+        elif self.reduction == 'sum':
+            return Tensor(np.sum(loss.data), requires_grad=True)
+        return loss
+
+class CrossEntropyLoss(Module):
+    def __init__(self, reduction='mean'):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, pred, target):
+        if not isinstance(target, Tensor):
+            target = Tensor(target, requires_grad=False)
+        
+        # Apply log softmax
+        max_val = np.max(pred.data, axis=1, keepdims=True)
+        exp_x = np.exp(pred.data - max_val)
+        softmax = exp_x / np.sum(exp_x, axis=1, keepdims=True)
+        log_softmax = np.log(softmax + 1e-10)
+        
+        # Compute cross entropy
+        batch_size = pred.data.shape[0]
+        loss = -np.sum(target.data * log_softmax) / batch_size if self.reduction == 'mean' \
+               else -np.sum(target.data * log_softmax)
+        
+        return Tensor(loss, requires_grad=True)
+
+class BCELoss(Module):
+    def __init__(self, reduction='mean'):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, pred, target):
+        if not isinstance(target, Tensor):
+            target = Tensor(target, requires_grad=False)
+        
+        eps = 1e-12
+        loss = -(target.data * np.log(pred.data + eps) + (1 - target.data) * np.log(1 - pred.data + eps))
+        
+        if self.reduction == 'mean':
+            return Tensor(np.mean(loss), requires_grad=True)
+        elif self.reduction == 'sum':
+            return Tensor(np.sum(loss), requires_grad=True)
+        return Tensor(loss, requires_grad=True)
+
+class BCEWithLogitsLoss(Module):
+    def __init__(self, reduction='mean'):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, pred, target):
+        if not isinstance(target, Tensor):
+            target = Tensor(target, requires_grad=False)
+        
+        # Compute sigmoid and BCE in a numerically stable way
+        max_val = np.maximum(0, pred.data)
+        loss = pred.data - pred.data * target.data + max_val + \
+               np.log(np.exp(-max_val) + np.exp(pred.data - max_val))
+        
+        if self.reduction == 'mean':
+            return Tensor(np.mean(loss), requires_grad=True)
+        elif self.reduction == 'sum':
+            return Tensor(np.sum(loss), requires_grad=True)
+        return Tensor(loss, requires_grad=True)
+
+class L1Loss(Module):
+    def __init__(self, reduction='mean'):
+        super().__init__()
+        self.reduction = reduction
+
+    def forward(self, pred, target):
+        if not isinstance(target, Tensor):
+            target = Tensor(target, requires_grad=False)
+        
+        loss = np.abs(pred.data - target.data)
+        if self.reduction == 'mean':
+            return Tensor(np.mean(loss), requires_grad=True)
+        elif self.reduction == 'sum':
+            return Tensor(np.sum(loss), requires_grad=True)
+        return Tensor(loss, requires_grad=True)
+
+class SmoothL1Loss(Module):
+    def __init__(self, reduction='mean', beta=1.0):
+        super().__init__()
+        self.reduction = reduction
+        self.beta = beta
+
+    def forward(self, pred, target):
+        if not isinstance(target, Tensor):
+            target = Tensor(target, requires_grad=False)
+        
+        diff = np.abs(pred.data - target.data)
+        loss = np.where(diff < self.beta, 0.5 * diff ** 2 / self.beta, diff - 0.5 * self.beta)
+        
+        if self.reduction == 'mean':
+            return Tensor(np.mean(loss), requires_grad=True)
+        elif self.reduction == 'sum':
+            return Tensor(np.sum(loss), requires_grad=True)
+        return Tensor(loss, requires_grad=True)
+
+class KLDivLoss(Module):
+    def __init__(self, reduction='mean', log_target=False):
+        super().__init__()
+        self.reduction = reduction
+        self.log_target = log_target
+
+    def forward(self, pred, target):
+        if not isinstance(target, Tensor):
+            target = Tensor(target, requires_grad=False)
+        
+        if self.log_target:
+            loss = np.exp(target.data) * (target.data - pred.data)
+        else:
+            loss = target.data * (np.log(target.data + 1e-10) - pred.data)
+        
+        if self.reduction == 'mean':
+            return Tensor(np.mean(loss), requires_grad=True)
+        elif self.reduction == 'sum':
+            return Tensor(np.sum(loss), requires_grad=True)
+        return Tensor(loss, requires_grad=True)
+
+class CosineEmbeddingLoss(Module):
+    def __init__(self, margin=0.0, reduction='mean'):
+        super().__init__()
+        self.margin = margin
+        self.reduction = reduction
+
+    def forward(self, x1, x2, y):
+        if not isinstance(y, Tensor):
+            y = Tensor(y, requires_grad=False)
+        
+        cos_sim = np.sum(x1.data * x2.data, axis=1) / \
+                 (np.sqrt(np.sum(x1.data ** 2, axis=1)) * np.sqrt(np.sum(x2.data ** 2, axis=1)) + 1e-10)
+        
+        loss = np.where(y.data == 1, 1 - cos_sim, np.maximum(0, cos_sim - self.margin))
+        
+        if self.reduction == 'mean':
+            return Tensor(np.mean(loss), requires_grad=True)
+        elif self.reduction == 'sum':
+            return Tensor(np.sum(loss), requires_grad=True)
+        return Tensor(loss, requires_grad=True)

openarchx/nn/module.py

@@ -0,0 +1,18 @@
+class Module:
+    def __init__(self):
+        self._parameters = []
+        
+    def parameters(self):
+        params = []
+        for attr in self.__dict__.values():
+            if isinstance(attr, Module):
+                params.extend(attr.parameters())
+            elif hasattr(attr, 'parameters'):
+                params.extend(attr.parameters())
+        return params + self._parameters
+    
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
+        
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)