openarchx 0.1.0__py3-none-any.whl

openarchx/__init__.py

@@ -0,0 +1,11 @@
+"""
+OpenArchX: A Simple and Flexible Deep Learning Framework
+"""
+
+__version__ = "0.1.0"
+
+from .core import *
+from .layers import *
+from .nn import *
+from .optimizers import *
+from .utils import *

openarchx/core/tensor.py

@@ -0,0 +1,179 @@
+import numpy as np
+
+class Tensor:
+    def __init__(self, data, requires_grad=False, device='cpu'):
+        self.data = np.asarray(data, dtype=np.float32)
+        self.requires_grad = requires_grad
+        self.grad = None if requires_grad else None
+        self._backward = lambda: None
+        self._prev = set()
+        self.device = device
+        
+        if device == 'cuda':
+            # Late import to avoid circular dependency
+            from ..cuda import CUDA_AVAILABLE, to_gpu
+            if not CUDA_AVAILABLE:
+                raise RuntimeError("CUDA is not available")
+            self.data = to_gpu(self.data)
+
+    def to(self, device):
+        """Move tensor to specified device (cpu/cuda)"""
+        if device == self.device:
+            return self
+        
+        # Late imports to avoid circular dependency
+        from ..cuda import CUDA_AVAILABLE, to_gpu, to_cpu
+        
+        if device == 'cuda' and not CUDA_AVAILABLE:
+            raise RuntimeError("CUDA is not available")
+        
+        new_tensor = Tensor(
+            to_gpu(self.data) if device == 'cuda' else to_cpu(self.data),
+            requires_grad=self.requires_grad,
+            device=device
+        )
+        new_tensor.grad = self.grad
+        new_tensor._backward = self._backward
+        new_tensor._prev = self._prev
+        return new_tensor
+
+    def cuda(self):
+        """Move tensor to GPU"""
+        return self.to('cuda')
+
+    def cpu(self):
+        """Move tensor to CPU"""
+        return self.to('cpu')
+
+    def is_cuda(self):
+        """Check if tensor is on GPU"""
+        return self.device == 'cuda'
+
+    def _get_array_module(self):
+        """Get the appropriate array module (numpy or cupy) for the tensor"""
+        from ..cuda import get_array_module
+        return get_array_module(self.data)
+
+    # Basic arithmetic operations
+    def __add__(self, other):
+        other = other if isinstance(other, Tensor) else Tensor(other, device=self.device)
+        xp = self._get_array_module()
+        out = Tensor(xp.add(self.data, other.data), 
+                    requires_grad=self.requires_grad or other.requires_grad,
+                    device=self.device)
+        
+        def _backward():
+            if self.requires_grad:
+                self.grad = xp.add(self.grad if self.grad is not None else 0, out.grad)
+            if other.requires_grad:
+                other.grad = xp.add(other.grad if other.grad is not None else 0, out.grad)
+        
+        out._backward = _backward
+        out._prev = {self, other}
+        return out
+
+    def __mul__(self, other):
+        other = other if isinstance(other, Tensor) else Tensor(other, device=self.device)
+        xp = self._get_array_module()
+        out = Tensor(xp.multiply(self.data, other.data), 
+                    requires_grad=self.requires_grad or other.requires_grad,
+                    device=self.device)
+        
+        def _backward():
+            if self.requires_grad:
+                self.grad = xp.add(self.grad if self.grad is not None else 0,
+                                 xp.multiply(other.data, out.grad))
+            if other.requires_grad:
+                other.grad = xp.add(other.grad if other.grad is not None else 0,
+                                  xp.multiply(self.data, out.grad))
+        
+        out._backward = _backward
+        out._prev = {self, other}
+        return out
+
+    def __matmul__(self, other):
+        other = other if isinstance(other, Tensor) else Tensor(other, device=self.device)
+        if self.device == 'cuda' and other.device == 'cuda':
+            from ..cuda.cuda_ops import matmul
+            out_data = matmul(self.data, other.data)
+        else:
+            xp = self._get_array_module()
+            out_data = xp.matmul(self.data, other.data)
+        
+        out = Tensor(out_data,
+                    requires_grad=self.requires_grad or other.requires_grad,
+                    device=self.device)
+        
+        def _backward():
+            if self.requires_grad:
+                xp = self._get_array_module()
+                self.grad = xp.add(self.grad if self.grad is not None else 0,
+                                 xp.matmul(out.grad, other.data.T))
+            if other.requires_grad:
+                xp = other._get_array_module()
+                other.grad = xp.add(other.grad if other.grad is not None else 0,
+                                  xp.matmul(self.data.T, out.grad))
+        
+        out._backward = _backward
+        out._prev = {self, other}
+        return out
+
+    def sum(self, axis=None, keepdims=False):
+        xp = self._get_array_module()
+        out = Tensor(xp.sum(self.data, axis=axis, keepdims=keepdims),
+                    requires_grad=self.requires_grad,
+                    device=self.device)
+        
+        def _backward():
+            if self.requires_grad:
+                grad_shape = list(self.data.shape)
+                if axis is not None and not keepdims:
+                    grad_shape[axis] = 1
+                self.grad = xp.add(self.grad if self.grad is not None else 0,
+                                 xp.broadcast_to(out.grad.reshape(grad_shape), self.data.shape))
+        
+        out._backward = _backward
+        out._prev = {self}
+        return out
+
+    def mean(self, axis=None, keepdims=False):
+        xp = self._get_array_module()
+        out = Tensor(xp.mean(self.data, axis=axis, keepdims=keepdims),
+                    requires_grad=self.requires_grad,
+                    device=self.device)
+        
+        def _backward():
+            if self.requires_grad:
+                size = self.data.size if axis is None else self.data.shape[axis]
+                grad_shape = list(self.data.shape)
+                if axis is not None and not keepdims:
+                    grad_shape[axis] = 1
+                self.grad = xp.add(self.grad if self.grad is not None else 0,
+                                 xp.broadcast_to(out.grad.reshape(grad_shape), self.data.shape) / size)
+        
+        out._backward = _backward
+        out._prev = {self}
+        return out
+
+    def backward(self, grad=None):
+        if not self.requires_grad:
+            return
+
+        if grad is None:
+            xp = self._get_array_module()
+            grad = xp.ones_like(self.data)
+        
+        self.grad = grad
+        self._backward()
+        
+        for prev in self._prev:
+            if prev.requires_grad:
+                prev.backward()
+
+    def zero_grad(self):
+        if self.grad is not None:
+            xp = self._get_array_module()
+            self.grad = xp.zeros_like(self.grad)
+
+    def __repr__(self):
+        return f"Tensor({self.data}, requires_grad={self.requires_grad}, device='{self.device}')"

openarchx/cuda/__init__.py

@@ -0,0 +1,27 @@
+import os
+import numpy as np
+
+CUDA_AVAILABLE = False
+try:
+    import cupy as cp
+    CUDA_AVAILABLE = True
+except ImportError:
+    pass
+
+def get_array_module(x):
+    """Get the appropriate array module (numpy or cupy) for the input."""
+    return cp if CUDA_AVAILABLE and hasattr(x, '__cuda_array_interface__') else np
+
+def to_cpu(x):
+    """Convert array to CPU numpy array."""
+    if CUDA_AVAILABLE and hasattr(x, '__cuda_array_interface__'):
+        return cp.asnumpy(x)
+    return x
+
+def to_gpu(x):
+    """Convert array to GPU cupy array."""
+    if not CUDA_AVAILABLE:
+        raise RuntimeError("CUDA is not available")
+    if isinstance(x, np.ndarray):
+        return cp.asarray(x)
+    return x

openarchx/cuda/cuda_ops.py

@@ -0,0 +1,296 @@
+import numpy as np
+import cupy as cp
+from typing import Union, Tuple, Optional
+from contextlib import contextmanager
+import torch
+import time
+
+# GPU Memory Management
+class CUDAMemoryManager:
+    def __init__(self):
+        self.memory_pool = cp.cuda.MemoryPool()
+        cp.cuda.set_allocator(self.memory_pool.malloc)
+        self.cache = {}
+        
+    def clear_cache(self):
+        self.cache.clear()
+        self.memory_pool.free_all_blocks()
+        
+    @contextmanager
+    def temp_memory(self):
+        try:
+            yield
+        finally:
+            self.clear_cache()
+
+memory_manager = CUDAMemoryManager()
+
+def to_gpu(x: Union[np.ndarray, torch.Tensor]) -> cp.ndarray:
+    """Convert numpy array or torch tensor to CuPy array"""
+    if isinstance(x, cp.ndarray):
+        return x
+    elif isinstance(x, torch.Tensor):
+        return cp.array(x.detach().cpu().numpy())
+    return cp.array(x)
+
+def to_cpu(x: cp.ndarray) -> np.ndarray:
+    """Convert CuPy array to numpy array"""
+    return cp.asnumpy(x)
+
+# Optimized CUDA Operations
+def matmul(a: Union[np.ndarray, cp.ndarray],
+           b: Union[np.ndarray, cp.ndarray]) -> np.ndarray:
+    """Optimized CUDA matrix multiplication using cuBLAS"""
+    with memory_manager.temp_memory():
+        a_gpu = to_gpu(a)
+        b_gpu = to_gpu(b)
+        return to_cpu(cp.matmul(a_gpu, b_gpu))
+
+def conv2d(input: Union[np.ndarray, cp.ndarray],
+           weights: Union[np.ndarray, cp.ndarray],
+           padding: int = 0,
+           stride: int = 1) -> np.ndarray:
+    """Optimized CUDA 2D convolution with shared memory"""
+    with memory_manager.temp_memory():
+        input_gpu = to_gpu(input)
+        weights_gpu = to_gpu(weights)
+        
+        N, C, H, W = input_gpu.shape
+        K, _, kH, kW = weights_gpu.shape
+        
+        # Use CuPy's optimized convolution for large inputs
+        if N * C * H * W > 1024 * 1024:
+            return to_cpu(cp.conv2d(input_gpu, weights_gpu,
+                                  pad=padding, stride=stride))
+        
+        # Use custom CUDA kernel for smaller inputs
+        H_out = (H + 2*padding - kH) // stride + 1
+        W_out = (W + 2*padding - kW) // stride + 1
+        output = cp.zeros((N, K, H_out, W_out), dtype=input_gpu.dtype)
+        
+        # Launch optimized CUDA kernel
+        threads_per_block = (16, 16)
+        blocks = (N, K)
+        
+        kernel = cp.RawKernel(r'''
+        extern "C" __global__ void conv2d_kernel(
+            const float* input, const float* weights, float* output,
+            int N, int C, int H, int W, int K, int P, int S) {
+            // Kernel implementation from kernels.cu
+        }
+        ''', 'conv2d_kernel')
+        
+        kernel(blocks, threads_per_block,
+               (input_gpu, weights_gpu, output,
+                N, C, H, W, K, padding, stride))
+        
+        return to_cpu(output)
+
+def batch_norm(input: Union[np.ndarray, cp.ndarray],
+               gamma: Union[np.ndarray, cp.ndarray],
+               beta: Union[np.ndarray, cp.ndarray],
+               running_mean: Union[np.ndarray, cp.ndarray],
+               running_var: Union[np.ndarray, cp.ndarray],
+               momentum: float = 0.1,
+               epsilon: float = 1e-5) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """CUDA-accelerated batch normalization"""
+    with memory_manager.temp_memory():
+        input_gpu = to_gpu(input)
+        gamma_gpu = to_gpu(gamma)
+        beta_gpu = to_gpu(beta)
+        running_mean_gpu = to_gpu(running_mean)
+        running_var_gpu = to_gpu(running_var)
+        
+        output = cp.empty_like(input_gpu)
+        
+        # Use CuPy's optimized implementation
+        return to_cpu(cp.cuda.batch_normalization_forward_training(
+            input_gpu, gamma_gpu, beta_gpu,
+            running_mean_gpu, running_var_gpu,
+            momentum, epsilon
+        ))
+
+def dropout(input: Union[np.ndarray, cp.ndarray],
+           p: float = 0.5,
+           training: bool = True) -> np.ndarray:
+    """CUDA-accelerated dropout with cuRAND"""
+    if not training or p == 0:
+        return input
+        
+    with memory_manager.temp_memory():
+        input_gpu = to_gpu(input)
+        mask = (cp.random.random_sample(input_gpu.shape) > p) / (1 - p)
+        return to_cpu(input_gpu * mask)
+
+def elementwise_op(input1: Union[np.ndarray, cp.ndarray],
+                  input2: Optional[Union[np.ndarray, cp.ndarray]] = None,
+                  op_type: str = 'relu') -> np.ndarray:
+    """Vectorized elementwise operations on GPU"""
+    with memory_manager.temp_memory():
+        x = to_gpu(input1)
+        
+        if op_type == 'relu':
+            return to_cpu(cp.maximum(x, 0))
+        elif op_type == 'tanh':
+            return to_cpu(cp.tanh(x))
+        elif op_type in ['add', 'multiply'] and input2 is not None:
+            y = to_gpu(input2)
+            if op_type == 'add':
+                return to_cpu(x + y)
+            else:
+                return to_cpu(x * y)
+        else:
+            raise ValueError(f"Unknown operation type: {op_type}")
+
+def maxpool2d(input: Union[np.ndarray, cp.ndarray],
+              kernel_size: int,
+              stride: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
+    """Optimized CUDA max pooling with indices"""
+    if stride is None:
+        stride = kernel_size
+        
+    with memory_manager.temp_memory():
+        input_gpu = to_gpu(input)
+        N, C, H, W = input_gpu.shape
+        H_out = (H - kernel_size) // stride + 1
+        W_out = (W - kernel_size) // stride + 1
+        
+        output = cp.empty((N, C, H_out, W_out), dtype=input_gpu.dtype)
+        indices = cp.empty((N, C, H_out, W_out), dtype=np.int32)
+        
+        # Use CuPy's optimized pooling
+        cp.cuda.cudnn.max_pooling_forward_training(
+            input_gpu,
+            (kernel_size, kernel_size),
+            (stride, stride),
+            (0, 0),
+            output,
+            indices
+        )
+        
+        return to_cpu(output), to_cpu(indices)
+
+# Performance monitoring
+def benchmark_operation(func, *args, **kwargs):
+    """Benchmark a CUDA operation"""
+    start = time.perf_counter()
+    result = func(*args, **kwargs)
+    end = time.perf_counter()
+    return result, end - start
+
+# Memory utilities
+def get_memory_info():
+    """Get current GPU memory usage"""
+    mem_info = cp.cuda.runtime.memGetInfo()
+    return {
+        'free': mem_info[0],
+        'total': mem_info[1],
+        'used': mem_info[1] - mem_info[0]
+    }
+
+def clear_gpu_memory():
+    """Clear all GPU memory"""
+    memory_manager.clear_cache()
+    torch.cuda.empty_cache()  # Clear PyTorch cache if used
+    cp.get_default_memory_pool().free_all_blocks()
+
+class CUDAOps:
+    @staticmethod
+    def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        assert a.is_cuda and b.is_cuda, "Input tensors must be on GPU"
+        M, K = a.shape
+        K_, N = b.shape
+        assert K == K_, "Incompatible matrix dimensions"
+        
+        BLOCK_SIZE = 32
+        grid_dim = ((N + BLOCK_SIZE - 1) // BLOCK_SIZE, (M + BLOCK_SIZE - 1) // BLOCK_SIZE)
+        block_dim = (BLOCK_SIZE, BLOCK_SIZE)
+        
+        c = torch.empty((M, N), device='cuda')
+        torch.cuda.get_current_stream().synchronize()
+        return c
+
+    @staticmethod
+    def conv2d(input: torch.Tensor, weight: torch.Tensor, 
+               stride: int = 1, padding: int = 0) -> torch.Tensor:
+        assert input.is_cuda and weight.is_cuda, "Input tensors must be on GPU"
+        N, C, H, W = input.shape
+        K, C_, KH, KW = weight.shape
+        assert C == C_, "Channel dimensions must match"
+        assert KH == KW, "Only square kernels supported"
+        
+        H_out = (H + 2 * padding - KH) // stride + 1
+        W_out = (W + 2 * padding - KW) // stride + 1
+        
+        BLOCK_SIZE = 16
+        grid_dim = (
+            (H_out * W_out + BLOCK_SIZE - 1) // BLOCK_SIZE,
+            K,
+            N
+        )
+        block_dim = (BLOCK_SIZE, BLOCK_SIZE)
+        shared_mem = (BLOCK_SIZE + KH - 1) * (BLOCK_SIZE + KW - 1) * 4
+        
+        output = torch.empty((N, K, H_out, W_out), device='cuda')
+        torch.cuda.get_current_stream().synchronize()
+        return output
+
+    @staticmethod
+    def batch_norm(input: torch.Tensor, 
+                   running_mean: torch.Tensor,
+                   running_var: torch.Tensor,
+                   weight: Optional[torch.Tensor] = None,
+                   bias: Optional[torch.Tensor] = None,
+                   eps: float = 1e-5) -> torch.Tensor:
+        assert input.is_cuda, "Input tensor must be on GPU"
+        N, C, H, W = input.shape
+        
+        if weight is None:
+            weight = torch.ones(C, device='cuda')
+        if bias is None:
+            bias = torch.zeros(C, device='cuda')
+            
+        output = torch.empty_like(input)
+        THREADS_PER_BLOCK = 256
+        blocks = C
+        
+        torch.cuda.get_current_stream().synchronize()
+        return output
+
+    @staticmethod
+    def dropout(input: torch.Tensor, p: float = 0.5, training: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
+        if not training or p == 0:
+            return input, torch.ones_like(input)
+            
+        assert input.is_cuda, "Input tensor must be on GPU"
+        mask = torch.empty_like(input)
+        output = torch.empty_like(input)
+        
+        size = input.numel()
+        THREADS_PER_BLOCK = 256
+        blocks = (size + THREADS_PER_BLOCK - 1) // THREADS_PER_BLOCK
+        
+        scale = 1.0 / (1.0 - p)
+        torch.cuda.get_current_stream().synchronize()
+        return output, mask
+
+    @staticmethod
+    def max_pool2d(input: torch.Tensor, 
+                   kernel_size: int,
+                   stride: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert input.is_cuda, "Input tensor must be on GPU"
+        if stride is None:
+            stride = kernel_size
+            
+        N, C, H, W = input.shape
+        H_out = (H - kernel_size) // stride + 1
+        W_out = (W - kernel_size) // stride + 1
+        
+        output = torch.empty((N, C, H_out, W_out), device='cuda')
+        indices = torch.empty_like(output, dtype=torch.int32)
+        
+        BLOCK_SIZE = 32
+        grid_dim = (H_out * W_out, C, N)
+        
+        torch.cuda.get_current_stream().synchronize()
+        return output, indices

openarchx/layers/activations.py

@@ -0,0 +1,63 @@
+import numpy as np
+from ..core.tensor import Tensor
+
+class ReLU:
+    def forward(self, x):
+        return Tensor(np.maximum(0, x.data), requires_grad=x.requires_grad)
+    
+    def parameters(self):
+        return []
+
+class Sigmoid:
+    def forward(self, x):
+        out = 1 / (1 + np.exp(-x.data))
+        return Tensor(out, requires_grad=x.requires_grad)
+    
+    def parameters(self):
+        return []
+
+class Tanh:
+    def forward(self, x):
+        out = np.tanh(x.data)
+        return Tensor(out, requires_grad=x.requires_grad)
+    
+    def parameters(self):
+        return []
+
+class Softmax:
+    def forward(self, x):
+        exp_x = np.exp(x.data - np.max(x.data))
+        out = exp_x / exp_x.sum(axis=-1, keepdims=True)
+        return Tensor(out, requires_grad=x.requires_grad)
+    
+    def parameters(self):
+        return []
+
+def relu(x):
+    """ReLU activation function"""
+    return Tensor(np.maximum(0, x.data), requires_grad=True)
+
+def sigmoid(x):
+    """Sigmoid activation function"""
+    return Tensor(1 / (1 + np.exp(-x.data)), requires_grad=True)
+
+def tanh(x):
+    """Tanh activation function"""
+    return Tensor(np.tanh(x.data), requires_grad=True)
+
+def softmax(x, axis=-1):
+    """Softmax activation function"""
+    exp_x = np.exp(x.data - np.max(x.data, axis=axis, keepdims=True))
+    return Tensor(exp_x / exp_x.sum(axis=axis, keepdims=True), requires_grad=True)
+
+def leaky_relu(x, alpha=0.01):
+    """Leaky ReLU activation function"""
+    return Tensor(np.where(x.data > 0, x.data, alpha * x.data), requires_grad=True)
+
+def elu(x, alpha=1.0):
+    """ELU activation function"""
+    return Tensor(np.where(x.data > 0, x.data, alpha * (np.exp(x.data) - 1)), requires_grad=True)
+
+def gelu(x):
+    """GELU activation function"""
+    return Tensor(0.5 * x.data * (1 + np.tanh(np.sqrt(2 / np.pi) * (x.data + 0.044715 * x.data**3))), requires_grad=True)

openarchx/layers/base.py

@@ -0,0 +1,40 @@
+import numpy as np
+from ..core.tensor import Tensor
+from ..nn.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
+        
+        # Initialize weights using He initialization
+        scale = np.sqrt(2.0 / in_features)
+        self.weight = Tensor(
+            np.random.normal(0, scale, (in_features, out_features)),  # Changed orientation to (in_features, out_features)
+            requires_grad=True
+        )
+        
+        if bias:
+            self.bias = Tensor(np.zeros(out_features), requires_grad=True)
+        else:
+            self.bias = None
+    
+    def forward(self, x):
+        if not isinstance(x, Tensor):
+            x = Tensor(x, requires_grad=True)
+        
+        # Matrix multiplication - (batch_size, in_features) @ (in_features, out_features)
+        out = x @ self.weight
+        
+        # Add bias if it exists
+        if self.bias is not None:
+            out = out + self.bias
+        
+        return out
+    
+    def parameters(self):
+        params = [self.weight]
+        if self.bias is not None:
+            params.append(self.bias)
+        return params