gradix 0.1.0__py3-none-any.whl

gradix/__init__.py

@@ -0,0 +1,35 @@
+import numpy as np
+
+try:
+	import cupy as cp
+	import cupyx
+except ImportError:
+	cp = None
+	cupyx = None
+
+GPU_AVAILABLE = cp is not None and cp.cuda.is_available() and cp.cuda.runtime.getDeviceCount() > 0
+DEVICE = 'cuda' if GPU_AVAILABLE else 'cpu'
+
+#dtypes
+float32 = np.float32
+float64 = np.float64
+float16 = np.float16
+int32   = np.int32
+int64   = np.int64
+bool_   = np.bool_
+
+from .tensor import Tensor
+from . import nn
+
+__all__ = [
+	"Tensor",
+	"nn",
+	"GPU_AVAILABLE",
+	"DEVICE",
+	"float32",
+	"float64",
+	"float16",
+	"int32",
+	"int64",
+	"bool_",
+]

gradix/init.py

@@ -0,0 +1,33 @@
+import numpy as np
+
+def uniform(shape, left_bound = 0.0, right_bound = 1.0, dtype = np.float32): 
+    return np.random.uniform(left_bound, right_bound, size=shape).astype(dtype)
+
+def kaiming_uniform(shape, a = np.sqrt(5.0)):
+    fan_in, fan_out = _get_fan_in_and_fan_out(shape)
+    gain = np.sqrt(2.0 / (1.0 + a ** 2))
+
+    std = gain / np.sqrt(fan_in)
+
+    weight_bound  = np.sqrt(3.0) * std
+
+    return uniform(shape, -weight_bound, weight_bound)
+
+
+def _get_fan_in_and_fan_out(shape): 
+    assert isinstance(shape, (tuple, list))
+
+    if len(shape) < 2: 
+        raise ValueError('Can not implement fan in and fan out for shape less than 2')
+    fan_in = shape[1]
+    fan_out = shape[0]
+
+    receptive_field_size = 1
+
+    if len(shape) > 2: 
+        for s in shape[2:]: 
+            receptive_field_size *= s
+    fan_in *= receptive_field_size
+    fan_out *= receptive_field_size
+
+    return fan_in, fan_out

gradix/nn/__init__.py

@@ -0,0 +1,17 @@
+from . import functional, loss, modules, optim
+from .loss import cross_entropy
+from .modules import Conv2D, Linear, Module, Upsample
+from .optim import SGD
+
+__all__ = [
+	"functional",
+	"loss",
+	"modules",
+	"optim",
+	"Module",
+	"Linear",
+	"Conv2D",
+	"Upsample",
+	"cross_entropy",
+	"SGD",
+]

gradix/nn/functional.py

@@ -0,0 +1,108 @@
+import numpy as np
+from gradix.tensor import Tensor
+
+def _validate_int_tuple(*dims, func_name = ''): 
+    dims = tuple(dims)
+    if len(dims) == 1 and isinstance(dims, (tuple, list)): 
+        dims = tuple(dims[0])
+
+    if any((not isinstance(d, int) or d < 0) for d in dims): 
+        raise ValueError('Expected Positive int in : {func_name}')
+    return dims
+
+
+def sigmoid(x: Tensor): 
+    return 1 / (1 + (-x).exp())
+
+def softmax(x: Tensor): 
+    max_x = x.max()
+    x = x - max_x
+    ex = x.exp()
+    sum_ex = ex.sum()
+    return ex / sum_ex
+
+def log_softmax(x: Tensor, axis = 1): 
+    max_x = x.max(axis = axis, keepdims = True)
+    shifted = x - max_x
+    sum_ex = shifted.exp().sum(axis = axis, keepdims = True)
+    return shifted - sum_ex.log()
+
+
+def im2td(x, kernel, stride = 1, padding = 0): 
+    if not isinstance(x, Tensor) or x.ndim != 4:
+        raise ValueError("unfold2d expects a Tensor shaped (batch, channels, height, width)")
+    if not isinstance(kernel, int) or kernel <= 0:
+        raise ValueError("kernel_size must be a positive integer")
+    if not isinstance(stride, int) or stride <= 0:
+        raise ValueError("stride must be a positive integer")
+    if not isinstance(padding, int) or padding < 0:
+        raise ValueError("padding must be a non-negative integer")
+
+    batch_size, channels, height, width = x.shape
+    output_height = (height + 2 * padding - kernel) // stride + 1
+    output_width = (width + 2 * padding - kernel) // stride + 1
+
+    if output_height <= 0 or output_width <= 0:
+        raise ValueError("kernel_size cannot exceed the padded input dimensions")
+    
+
+    batch, out_y, out_x, channel, kernel_y, kernel_x = np.meshgrid(
+        np.arange(batch_size), 
+        np.arange(output_height), 
+        np.arange(output_width), 
+        np.arange(channels), 
+        np.arange(kernel), 
+        np.arange(kernel), 
+        indexing='ij'
+    )    
+
+    out_y = out_y * stride + kernel_y - padding
+    out_x = out_x * stride + kernel_x - padding
+
+    valid = (out_y >= 0) & (out_x >= 0) & (out_y < height) & (out_x < width)
+
+    clipped_x = np.clip(out_x, 0, width - 1)
+    clipped_y = np.clip(out_y, 0, height - 1)
+
+    indices = ((batch * channels + channel) * height + clipped_y) * width + clipped_x
+
+    gathered_out = x.gather(indices)
+
+    out = gathered_out * Tensor(valid.astype(np.float32))
+
+    return out.reshape(
+        batch_size * output_height * output_width, 
+        channels * kernel * kernel
+    )
+
+def repeat(x: Tensor, *dims): 
+    dims = _validate_int_tuple(dims)
+    if len(dims) < x.ndim: 
+        raise ValueError(f"repeat expects at least {x.ndim} but got {len(dims)}")
+    padded_shape = (1, ) * (len(dims) - x.ndim) + x.shape
+    reshape_shape = tuple(
+        value 
+        for size in padded_shape
+        for value in (1, size)
+    )
+    x = x.reshape(reshape_shape)
+    broadcast_shape = tuple(
+        value 
+        for repeat, shape in zip(dims, padded_shape)
+        for value in (repeat, shape)
+    )
+    x = x.broadcast_to(broadcast_shape)
+    output_shape = tuple(repeat * size for repeat, size in zip(dims, padded_shape))
+    return x.reshape(output_shape)
+
+
+def Upsample(x, scale_factor = 2): 
+    assert isinstance(scale_factor, int)
+    assert isinstance(x, Tensor)
+    assert x.ndim == 4
+
+    b, c, h, w = x.shape
+    x = x.reshape(b, c, h, 1, w, 1)
+    x = repeat(x, 1, 1, 1, scale_factor, 1, scale_factor)
+    print(x.shape)
+    return x.reshape(b, c, h * scale_factor, w * scale_factor)

gradix/nn/loss/__init__.py

@@ -0,0 +1,3 @@
+from .crossentropyloss import cross_entropy
+
+__all__ = ["cross_entropy"]

gradix/nn/loss/crossentropyloss.py

@@ -0,0 +1,19 @@
+import numpy as np
+from gradix.nn.functional import log_softmax
+from gradix.tensor import Tensor
+
+def cross_entropy(logits, targets): 
+    assert isinstance(logits, Tensor)
+    assert isinstance(targets, Tensor)
+    if logits.ndim != 2: 
+        raise ValueError(f'cross_entropy expects logits of shape (batch, classes) but found: {logits.shape}')
+    batch_size, num_classes = logits.shape
+    class_indices = targets.data.astype(int).reshape(-1)
+    if class_indices.size != batch_size: 
+        raise ValueError(f'Expected {batch_size} targets but found: {class_indices.size}')
+    if (class_indices < 0).any() or (class_indices >= num_classes).any(): 
+        raise IndexError(f'target class out of range [0, {num_classes})')
+    log_probs = log_softmax(logits)
+    row_offsets = num_classes * logits.xp.arange(batch_size)
+    corrected_log_probs = log_probs.gather(row_offsets + class_indices)
+    return -corrected_log_probs.mean()

gradix/nn/modules/__init__.py

@@ -0,0 +1,6 @@
+from .conv2d import Conv2D
+from .linear import Linear
+from .module import Module
+from .upsample import Upsample
+
+__all__ = ["Module", "Linear", "Conv2D", "Upsample"]

gradix/nn/modules/conv2d.py

@@ -0,0 +1,66 @@
+from gradix.tensor import Tensor
+import numpy as np
+
+from gradix.nn.modules.module import Module
+from gradix.init import _get_fan_in_and_fan_out, kaiming_uniform, uniform
+from gradix.nn import functional as F
+
+
+class Conv2D(Module): 
+    def __init__(self, in_channels, out_channels, kernel, stride = 1, padding = 0, use_bias = False): 
+        super().__init__()
+
+        if not isinstance(in_channels, int) or in_channels <= 0:
+            raise ValueError("in_channels must be a positive integer")
+        if not isinstance(out_channels, int) or out_channels <= 0:
+            raise ValueError("out_channels must be a positive integer")
+        if not isinstance(kernel, int) or kernel <= 0:
+            raise ValueError("kernel_size must be a positive integer")
+        if not isinstance(stride, int) or stride <= 0:
+            raise ValueError("stride must be a positive integer")
+        if not isinstance(padding, int) or padding < 0:
+            raise ValueError("padding must be a non-negative integer")
+        
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.kernel = kernel
+        self.stride = stride
+        self.padding = padding
+        self.bias = None
+        self.weight = Tensor(
+            kaiming_uniform((out_channels, in_channels, kernel, kernel), np.sqrt(5.0)), 
+            requires_grad = True, 
+            dtype = np.float32
+        )
+        if use_bias: 
+            fan_in, _ = _get_fan_in_and_fan_out(self.weight.shape)
+            bound = (1.0 / np.sqrt(fan_in)) if fan_in > 0 else 0
+            self.bias = Tensor(uniform((self.out_channels, ), -bound, bound), requires_grad=True, dtype=np.float32)
+
+    def parameters(self): 
+        return [self.weight] if self.bias is None else [self.weight, self.bias]
+
+    def forward(self, x): 
+        if not isinstance(x, Tensor) or x.ndim != 4: 
+            raise ValueError(f'Conv2D expects a tensor of shape (batch, channel, height, width) but found one of shape: {x.shape}')
+        
+        batch_size, in_channels, height, width = x.shape
+
+        if in_channels != self.in_channels: 
+            raise ValueError(f'Expected {self.in_channels} input channels but received: {in_channels}')
+
+
+        output_height = (height + 2 * self.padding - self.kernel) // self.stride + 1
+        output_width = (width + 2 * self.padding - self.kernel) // self.stride + 1
+
+        flattened_weight = self.weight.reshape(self.out_channels, -1)
+        patches = F.im2td(x, self.kernel, self.stride, self.padding)
+
+        product = patches @ flattened_weight.T
+
+        output = product.reshape(batch_size, output_height, output_width, self.out_channels).transpose(0, 3, 1, 2)
+
+        if self.bias is not None: 
+            output = output + self.bias.reshape(1, self.out_channels, 1, 1)
+        
+        return output

gradix/nn/modules/linear.py

@@ -0,0 +1,26 @@
+from gradix.nn.modules.module import Module
+from gradix.init import kaiming_uniform, uniform, _get_fan_in_and_fan_out
+import numpy as np
+from gradix.tensor import Tensor
+
+class Linear(Module): 
+    def __init__(self, in_features, out_features, bias = False):
+        super().__init__()
+        self.shape = (out_features, in_features)
+        self.weights = Tensor(kaiming_uniform(self.shape, np.sqrt(5.0)), requires_grad=True, dtype=np.float32)
+        self.bias = None
+        if bias: 
+            fan_in, _ = _get_fan_in_and_fan_out(self.shape)
+            bound = 1 / np.sqrt(fan_in) if fan_in > 0 else 0
+            self.bias = Tensor(uniform((out_features, ), -bound, bound), requires_grad=True, dtype=np.float32)
+
+    def parameters(self): 
+        if self.bias is not None: 
+            return [self.weights, self.bias]
+        return [self.weights]
+
+    def forward(self, x): 
+        out = x @ self.weights.T
+        if self.bias is not None: 
+            out = out + self.bias
+        return out

gradix/nn/modules/module.py

@@ -0,0 +1,14 @@
+from gradix.tensor import Tensor
+class Module: 
+    def zero_grad(self): 
+        for p in self.parameters(): 
+            p.grad = p.xp.zeros_like(p.data)
+
+    def __call__(self, *args, **kwargs): 
+        return self.forward(*args, **kwargs)
+
+    def forward(self, *args, **kwargs): 
+        raise NotImplementedError
+
+    def parameters(self): 
+        return []

gradix/nn/modules/upsample.py

@@ -0,0 +1,20 @@
+
+from gradix.nn.modules.module import Module
+
+
+class Upsample(Module): 
+    def __init__(self, scale_factor: int = 1, mode = 'nearest'): 
+        self.scale_factor = scale_factor
+        self.mode = mode
+
+    def _nearest_interpolate(self, x): 
+        pass
+
+    def _uniform_interpolate(self, x): 
+        pass
+
+    def parameters(self): 
+        return []
+
+    def forward(self, x): 
+        return x

gradix/nn/optim/SGD.py

@@ -0,0 +1,14 @@
+class SGD: 
+    def __init__(self, parameters, lr = 0.001): 
+        self.parameters = list(parameters)
+        self.lr = lr
+
+    def step(self): 
+        for p in self.parameters: 
+            if p.grad is not None: 
+                p.data -= self.lr * p.grad
+
+    def zero_grad(self): 
+        for p in self.parameters: 
+            if p.grad is not None: 
+                p.grad = p.xp.zeros_like(p.data)

gradix/nn/optim/__init__.py

@@ -0,0 +1,3 @@
+from .SGD import SGD
+
+__all__ = ["SGD"]

gradix/tensor.py

@@ -0,0 +1,336 @@
+import math
+import numpy as np
+from . import *
+
+class Tensor:
+    float32 = float32
+    float64 = float64
+    float16 = float16
+    int32 = int32
+    int64 = int64
+    bool_ = bool_
+    def __init__(self, data: np.ndarray, _childs: tuple = (), requires_grad = False, _op = None, device = DEVICE, dtype = None, name = None): 
+        if device == 'cuda' and cp is None:
+            raise RuntimeError("CUDA support requires installing the 'gpu' extra")
+        self.device = DEVICE if device not in ['cpu', 'cuda'] else device
+        self.xp = cp if self.device == 'cuda' else np
+        self.dtype = dtype
+        self.requires_grad = requires_grad
+        self._init_data(data)
+        self._backward = lambda: None
+        self.grad = self.xp.zeros_like(self.data) if requires_grad else None
+        self._prevs = set(_childs)
+        self._op = _op
+        self._name = None
+        self.name = name
+
+    @property
+    def name(self):
+        return self._name
+
+    @name.setter
+    def name(self, value):
+        self._name = None if value is None else str(value)
+
+    def _init_data(self, data):
+        if self.dtype is None: 
+            self.dtype = getattr(data, 'dtype', self.float32)
+            if self.requires_grad and np.dtype(self.dtype).kind != 'f':
+                self.dtype = self.float32
+        if isinstance(data, self.xp.ndarray):
+            self.data = data.astype(self.dtype, copy = False)
+        else: 
+            self.data = self.xp.asarray(data, dtype = self.dtype) 
+
+    def _scatter_add(self, target, key, values):
+        if self.device == 'cuda':
+            if cupyx is None:
+                raise RuntimeError("CUDA support requires installing the 'gpu' extra")
+            cupyx.scatter_add(target, key, values)
+            return
+        np.add.at(target, key, values)
+
+    def _const(self, other):
+        if isinstance(other, Tensor):
+            return other
+        return Tensor(other, device = self.device, dtype = self.dtype)
+
+    def _norm_axes(self, axis):
+        if axis is None:
+            return tuple(range(self.ndim))
+        axes = axis if isinstance(axis, tuple) else (axis,)
+        return tuple(ax % self.ndim for ax in axes)
+
+    @property
+    def T(self): 
+        return self.transpose()
+    
+    def transpose(self, *axes):
+        if not axes:
+            axes = tuple(reversed(range(self.data.ndim)))
+        elif len(axes) == 1 and isinstance(axes[0], (tuple, list)):
+            axes = tuple(axes[0])
+        else:
+            axes = tuple(axes)
+        axes = tuple(ax % self.ndim for ax in axes)
+        out = Tensor(self.xp.transpose(self.data, axes), (self,), self.requires_grad, _op='transpose', device = self.device)
+        def _backward(): 
+            inverse_axes = tuple(int(i) for i in np.argsort(axes))
+            self._accumulate_grad(self.xp.transpose(out.grad, inverse_axes))
+        out._backward = _backward
+        return out
+
+    def reshape(self, *shape):
+        if len(shape) == 1 and isinstance(shape[0], (tuple, list)):
+            shape = tuple(shape[0])
+        out = Tensor(self.data.reshape(*shape), (self,), self.requires_grad, _op='reshape', device = self.device)
+        def _backward():
+            self._accumulate_grad(out.grad.reshape(self.shape))
+        out._backward = _backward
+        return out
+
+    def gather(self, indices):
+        indices = self.xp.asarray(indices)
+        if not np.issubdtype(indices.dtype, np.integer):
+            raise TypeError("gather indices must be integers")
+        indices = indices.astype(self.xp.intp, copy=False)
+        if indices.size and (int(indices.min()) < 0 or int(indices.max()) >= self.data.size):
+            raise IndexError("gather index is outside the flattened Tensor")
+
+        out = Tensor(self.data.reshape(-1)[indices], (self,), self.requires_grad, _op='gather', device = self.device)
+
+        def _backward():
+            self._scatter_add(self.grad.reshape(-1), indices.reshape(-1), out.grad.reshape(-1))
+
+        out._backward = _backward
+        return out
+
+    def broadcast_to(self, shape: tuple): 
+        if len(shape) < self.ndim: 
+            raise ValueError(f"Can not broadcast to fewer dimension from {self.ndim} to {len(shape)}")
+
+        for source, target in zip((1,) * (len(shape) - self.ndim) + self.shape, shape): 
+            if source != target and source != 1: 
+                raise ValueError(f"Can not Broadcast shape: from {self.shape} to {shape}")
+
+        out = Tensor(self.xp.broadcast_to(self.data, shape), (self, ), requires_grad= self.requires_grad, _op ='broadcast_to', device = self.device)
+        def _backward(): 
+            self._accumulate_grad(out.grad)
+        out._backward = _backward
+        return out
+
+    def __getitem__(self, key):
+        out = Tensor(self.data[key], (self, ), requires_grad= self.requires_grad, _op='getitem', device = self.device)
+        def _backward(): 
+            self._scatter_add(self.grad, key, out.grad)
+        out._backward = _backward
+        return out
+
+    def requires_grad_(self, requires_grad: bool):
+        self.requires_grad = requires_grad
+        self.grad = self.xp.zeros_like(self.data) if self.requires_grad else None
+        return self
+
+    def __add__(self, other):
+        other = self._const(other)
+        requires_grad = self.requires_grad or other.requires_grad
+        out = Tensor(self.data + other.data, (self, other), requires_grad=requires_grad, _op='+', device = self.device)
+        def _backward(): 
+            self._accumulate_grad(out.grad)
+            other._accumulate_grad(out.grad)
+        out._backward = _backward
+        return out
+
+    def __mul__(self, other): 
+        other = self._const(other)
+        requires_grad = self.requires_grad or other.requires_grad
+        out = Tensor(self.data * other.data, (self, other), requires_grad=requires_grad, _op='*', device = self.device)
+        def _backward(): 
+            self._accumulate_grad(other.data * out.grad)
+            other._accumulate_grad(self.data * out.grad)
+        out._backward = _backward
+        return out
+
+    def __matmul__(self, other): 
+        other = self._const(other)
+        requires_grad = self.requires_grad or other.requires_grad
+        out = Tensor(self.data @ other.data, (self, other), requires_grad=requires_grad, _op='@', device = self.device)
+        def _backward(): 
+            a, b, grad = self.data, other.data, out.grad
+            if b.ndim == 1:
+                b, grad = b[:, None], self.xp.expand_dims(grad, -1)
+            if a.ndim == 1:
+                a, grad = a[None, :], self.xp.expand_dims(grad, -2)
+            grad_a = grad @ b.swapaxes(-1, -2)
+            grad_b = a.swapaxes(-1, -2) @ grad
+            if self.ndim == 1:
+                grad_a = grad_a[..., 0, :]
+            if other.ndim == 1:
+                grad_b = grad_b[..., 0]
+            self._accumulate_grad(grad_a)
+            other._accumulate_grad(grad_b)
+        out._backward = _backward
+        return out
+
+    def __pow__(self, value): 
+        assert isinstance(value, (int, float))
+        out = Tensor(self.xp.power(self.data, value), (self, ), requires_grad = self.requires_grad, _op=f'^{value}', device = self.device)
+        def _backward(): 
+            self._accumulate_grad(value * self.xp.power(self.data, value - 1) * out.grad)
+        out._backward = _backward
+        return out
+
+    def reciprocal(self): 
+        out = Tensor(1 / self.data, (self, ), requires_grad=self.requires_grad, _op='reciprocal', device = self.device)
+        def _backward(): 
+            self._accumulate_grad(-1 * out.grad / (self.data ** 2))
+        out._backward = _backward
+        return out 
+
+    def __div__(self, other): 
+        other = self._const(other)
+        return self * other.reciprocal()
+
+    def __truediv__(self, other):
+        return self.__div__(other)
+
+    def leaky_relu(self, lvalue = 0.1):
+        cdata = self.data.copy()
+        lindexes = self.data <= 0
+        cdata[lindexes] *= lvalue
+        out = Tensor(cdata, (self, ), requires_grad=self.requires_grad, _op=f'leaky_relu: {lvalue}', device = self.device)
+        def _backward():
+            grad = self.xp.ones_like(self.data)
+            grad[self.data <= 0] = lvalue
+            self._accumulate_grad(out.grad * grad)
+        out._backward = _backward
+        return out
+
+    def relu(self): 
+        return self.leaky_relu(0)
+
+    def exp(self):
+        out = Tensor(self.xp.exp(self.data), (self,), requires_grad=self.requires_grad, _op='exp', device = self.device)
+        def _backward():
+            self._accumulate_grad(out.grad * out.data)
+        out._backward = _backward
+        return out
+
+    def log(self): 
+        out = Tensor(self.xp.log(self.data), (self, ), requires_grad = self.requires_grad, _op='log', device = self.device)
+        def _backward(): 
+            self._accumulate_grad(out.grad / self.data)
+        out._backward = _backward
+        return out
+
+    def tanh(self): 
+        out = Tensor(self.xp.tanh(self.data), (self, ), requires_grad= self.requires_grad, _op='tanh', device = self.device)
+        def _backward(): 
+            self._accumulate_grad(out.grad * (1 - out.data ** 2))
+        out._backward = _backward
+        return out
+
+    def sigmoid(self): 
+        out = Tensor(0.5 * (1 + self.xp.tanh(self.data * 0.5)), (self, ), requires_grad= self.requires_grad, _op='sigmoid', device = self.device)
+        def _backward(): 
+            self._accumulate_grad(out.grad * out.data * (1 - out.data))
+        out._backward = _backward
+        return out
+
+    def _extremum(self, fn, axis, keepdims, _op):
+        out = Tensor(fn(self.data, axis = axis, keepdims = keepdims), (self, ), requires_grad=self.requires_grad, _op=_op, device = self.device)
+        def _backward(): 
+            kept = fn(self.data, axis = axis, keepdims = True)
+            grad = out.grad.reshape(kept.shape)
+            mask = (self.data == kept).astype(self.data.dtype)
+            mask = mask / mask.sum(axis = axis, keepdims = True)
+            self._accumulate_grad(grad * mask)
+        out._backward = _backward
+        return out
+
+    def max(self, axis = None, keepdims = False): 
+        return self._extremum(self.xp.max, axis, keepdims, 'max')
+
+    def min(self, axis = None, keepdims = False): 
+        return self._extremum(self.xp.min, axis, keepdims, 'min')
+
+    def sum(self, axis = None, keepdims = False):
+        out = Tensor(self.xp.sum( self.data, axis = axis, keepdims = keepdims),(self,),self.requires_grad, _op='sum', device = self.device)
+
+        def _backward():
+            axes = self._norm_axes(axis)
+            kept_shape = tuple(1 if i in axes else s for i, s in enumerate(self.shape))
+            grad = out.grad.reshape(kept_shape)
+            self._accumulate_grad(self.xp.broadcast_to(grad, self.shape))
+
+        out._backward = _backward
+        return out
+
+    def mean(self, axis = None, keepdims = False):
+        count = math.prod(self.shape[ax] for ax in self._norm_axes(axis))
+        return self.sum(axis = axis, keepdims = keepdims) / count
+
+    def __neg__(self): 
+        return self * -1
+    def __sub__(self, other): 
+        return self + (-other)
+    def __radd__(self, other): 
+        return self + other
+    def __rsub__(self, other): 
+        return (-self) + other
+    def __rmul__(self, other): 
+        return self * other
+    def __rtruediv__(self, other):
+        return other * self.reciprocal()
+
+    def zero_grad(self):
+        for node in self._build_topo():
+            if node.requires_grad:
+                node.grad = self.xp.zeros_like(node.data)
+    
+    def backward(self): 
+        if not self.requires_grad:
+            raise RuntimeError("backward() called on a Tensor that does not require grad")
+        nodes = list(self._build_topo())
+        for node in nodes:
+            if node._prevs:
+                node.grad = self.xp.zeros_like(node.data)
+        self.grad = self.xp.ones_like(self.data)
+        for node in nodes:
+            node._backward() 
+
+    def _build_topo(self): 
+        nodes, visited = [], set()
+        def build_topo(node): 
+            if (node not in visited) and (node.requires_grad): 
+                visited.add(node)
+                for child in node._prevs: 
+                    build_topo(child)
+                nodes.append(node)
+        build_topo(self)
+        return reversed(nodes)
+    
+    def _accumulate_grad(self, grad):
+        if not self.requires_grad: 
+            return
+        while grad.ndim > self.grad.ndim:
+            grad = grad.sum(axis=0)
+        axes = tuple(
+            i
+            for i, (gdim, tdim) in enumerate(zip(grad.shape, self.grad.shape))
+            if tdim == 1 and gdim != 1
+        )
+        if axes:
+            grad = grad.sum(axis=axes, keepdims=True)
+        self.grad += grad
+
+    @property
+    def shape(self): 
+        return self.data.shape
+
+    @property
+    def ndim(self): 
+        return self.data.ndim
+
+    def __repr__(self): 
+        return f'Tensor({self.data}, dtype: {self.data.dtype}, Shape: {list(self.data.shape)})'

gradix/utils.py

@@ -0,0 +1,34 @@
+from graphviz import Digraph
+
+def visualize_computation_graph(output_tensor, filename='computation_graph', format='png', view = True, c_view_options = []):
+    dot = Digraph(format=format, graph_attr={'rankdir': 'LR'}, )
+    visited = set()
+    
+    def add_nodes(tensor):
+        if tensor in visited:
+            return
+        visited.add(tensor)
+        show_data = ('data' in c_view_options) and tensor.ndim == 0
+        show_grad = ('grad' in c_view_options) and tensor.ndim == 0
+        tensor_label = f'Tensor: {tensor.name}' if tensor.name is not None else 'Tensor'
+        label = f'{tensor_label}\n shape: {tensor.shape}'
+        if show_data:
+            label += f'\n data: {tensor.data}'
+        if show_grad:
+            label += f'\n grad: {tensor.grad}'
+        dot.node(str(id(tensor)), label=label, shape='box')
+        if hasattr(tensor, '_op') and tensor._op:
+            dot.node(str(id(tensor)) + '_op', label=f'Op: {tensor._op}', shape='ellipse')
+            dot.edge(str(id(tensor)) + '_op', str(id(tensor)))
+
+        for parent in tensor._prevs:
+            add_nodes(parent)
+            if hasattr(tensor, '_op') and tensor._op:
+                dot.edge(str(id(parent)), str(id(tensor)) + '_op')
+            else: 
+                dot.edge(str(id(parent)), str(id(tensor)))
+
+    add_nodes(output_tensor)
+    if view is False:
+        dot.render(filename, view=view)
+    return dot

gradix-0.1.0.dist-info/METADATA

@@ -0,0 +1,312 @@
+Metadata-Version: 2.4
+Name: gradix
+Version: 0.1.0
+Summary: A lightweight NumPy-based automatic differentiation and deep learning library
+License-Expression: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.24
+Requires-Dist: graphviz>=0.20
+Requires-Dist: matplotlib>=3.8
+Provides-Extra: gpu
+Requires-Dist: cupy-cuda12x>=12; extra == "gpu"
+Provides-Extra: all
+Requires-Dist: cupy-cuda12x>=12; extra == "all"
+Requires-Dist: graphviz>=0.20; extra == "all"
+Requires-Dist: matplotlib>=3.8; extra == "all"
+Dynamic: license-file
+
+# Gradix
+
+Gradix is a small, educational deep-learning engine built around a `Tensor` class and automatic differentiation. It lets you compose tensor operations, build neural-network layers, calculate gradients with `backward()`, and update parameters with SGD.
+
+The project is intentionally compact and is useful for learning how computation graphs and training loops fit together. It is not a drop-in replacement for PyTorch or another production deep-learning framework.
+
+## What You Can Do
+
+- Create tensors and combine them with arithmetic, matrix multiplication, reductions, and activations.
+- Track operations in a computation graph and calculate gradients with backpropagation.
+- Build models with `Linear` and `Conv2D` modules.
+- Calculate classification loss with `cross_entropy` and update parameters with `SGD`.
+- Visualize a tensor's computation graph with Graphviz.
+
+## Install
+
+Gradix requires Python 3.10 or newer. Install the published package and its required Python dependencies with:
+
+```bash
+pip install gradix
+```
+
+The base installation includes NumPy, the Python Graphviz binding, and Matplotlib. Graph rendering also requires the Graphviz system executable (`dot`), which must be installed separately through your operating system's package manager.
+
+For example, on Ubuntu or Debian:
+
+```bash
+sudo apt install graphviz
+```
+
+CuPy is optional. The `gpu` extra currently requests the CUDA 12.x CuPy build:
+
+```bash
+pip install "gradix[gpu]"
+```
+
+You also need a compatible NVIDIA driver and CUDA setup. CPU usage is the recommended path for the current version; CUDA support is experimental, and some operations, particularly Conv2D's patch-index creation, still use NumPy internally.
+
+### Install From a Local Checkout
+
+When working in a clone of this repository, install the current source in editable mode:
+
+```bash
+python -m pip install -e .
+```
+
+This makes `import gradix` work from that Python environment without adding `src` to `sys.path` manually. In Jupyter, choose the same environment as the notebook kernel.
+
+## Your First Tensor
+
+```python
+from gradix import Tensor
+
+left = Tensor([1.0, 2.0, 3.0], requires_grad=True, name="left input")
+right = Tensor([4.0, 5.0, 6.0], requires_grad=True, name="right input")
+
+result = (left * right).sum()
+result.backward()
+
+print(result.data)  # 32.0
+print(left.grad)    # [4. 5. 6.]
+print(right.grad)   # [1. 2. 3.]
+```
+
+`requires_grad=True` asks Gradix to track operations that contribute to that tensor. Calling `backward()` on a result propagates its gradient backward through those operations. Gradients are stored in each participating tensor's `.grad` attribute.
+
+Tensor names can be assigned at construction or afterward. Any non-`None` value is converted to a string:
+
+```python
+image = Tensor([0.0, 1.0], name=1)
+image.name = 2.5
+assert image.name == "2.5"
+```
+
+## Public Imports
+
+Gradix re-exports common classes at convenient package levels:
+
+```python
+from gradix import Tensor
+from gradix.nn import Linear, Conv2D, cross_entropy, SGD
+from gradix.nn.modules import Module, Linear, Conv2D, Upsample
+from gradix.nn.loss import cross_entropy
+from gradix.nn.optim import SGD
+```
+
+## How Backpropagation Flows
+
+Each operation creates a result tensor that remembers its input tensors and a local backward function. For example, matrix multiplication records its operands. When `backward()` is called, Gradix traverses the graph and applies those local derivatives from the output toward the inputs.
+
+```mermaid
+flowchart LR
+	A[Input a] --> M[Multiply]
+	B[Input b] --> M
+	M --> S[Sum]
+	S --> L[Scalar loss]
+	L -->|backward| S
+	S --> M
+	M --> GA[Gradient for a]
+	M --> GB[Gradient for b]
+```
+
+The core idea in code is:
+
+```python
+prediction = model(batch)
+loss = cross_entropy(prediction, targets)
+optimizer.zero_grad()
+loss.backward()
+optimizer.step()
+```
+
+`zero_grad()` clears gradients from the previous batch. `backward()` calculates gradients for this batch. `step()` uses those gradients to update parameters.
+
+## Visualize a Computation Graph
+
+The graph examples in `demos/digraph_test.ipynb` build longer expressions and ask Graphviz to display their operation dependencies. Here is a smaller version:
+
+```python
+from gradix import Tensor
+from gradix.utils import visualize_computation_graph
+
+a = Tensor(1.0, requires_grad=True, name="a")
+b = Tensor(2.0, requires_grad=True, name="b")
+c = Tensor(3.0, requires_grad=True, name="c")
+
+result = (a + b) * c
+result.backward()
+
+graph = visualize_computation_graph(
+	result,
+	filename="example_graph",
+	c_view_options=["data", "grad"],
+)
+graph
+```
+
+Each box represents a tensor and each ellipse represents an operation. Edges connect an operation to its result and its inputs. Tensor names, shapes, scalar data, and scalar gradients can help identify nodes. Graphviz's `dot` executable must be installed for rendering.
+
+The attention-style example in `demos/digraph_test.ipynb` follows this flow:
+
+1. Create input and weight tensors for query, key, value, and projection operations.
+2. Compute query-key scores with matrix multiplication and a transpose.
+3. Apply a numerically stabilized softmax using max, exponentiation, and sum.
+4. Multiply attention weights by values, then apply a projection.
+5. Add a residual connection and normalize using mean and variance operations.
+6. Pass through a small MLP, combine the residual output, reduce to a scalar loss, and visualize the full graph.
+
+The notebook also builds a recurrent, gated sequence example. Each time-step appends more operations to the same graph through the hidden state; reducing the final hidden state to a scalar and visualizing it shows how dependencies span time steps.
+
+`demos/demo.ipynb` contains a smaller arithmetic graph. The corresponding examples are notebooks, so run cells from top to bottom after installing Gradix and Graphviz.
+
+## Define a Model
+
+Models subclass `Module`, implement `forward`, and return all trainable tensors from `parameters`. The `Linear` layer accepts a batch whose last dimension matches `in_features`.
+
+```python
+from gradix.nn.modules import Linear, Module
+
+
+class SmallClassifier(Module):
+	def __init__(self):
+		self.hidden = Linear(784, 128, bias=True)
+		self.output = Linear(128, 10, bias=True)
+
+	def parameters(self):
+		return self.hidden.parameters() + self.output.parameters()
+
+	def forward(self, x):
+		x = x.reshape(x.shape[0], 784)
+		x = self.hidden(x).relu()
+		return self.output(x)
+```
+
+The MNIST examples in `demos/MnistModel.py` show two model shapes:
+
+- `MNISTModel`: flatten each 28-by-28 image, apply a 784-to-128 linear layer with ReLU, then produce 10 class scores.
+- `MNISTConvModel`: reshape to `(batch, 1, 28, 28)`, apply three Conv2D layers with ReLU, flatten the result, and produce 10 class scores.
+
+The convolutional model's `parameters()` method combines the parameters from each child layer. Gradix modules do not currently automatically register nested modules, so a custom model should explicitly return its parameters as above.
+
+## Train a Model
+
+Gradix's `cross_entropy` expects logits shaped `(batch, classes)` and integer class labels in a target tensor. The targets do not need gradients. This self-contained example demonstrates the same batch flow used in `demos/train.ipynb` without downloading a dataset:
+
+```python
+import numpy as np
+
+from gradix import Tensor
+from gradix.nn.loss import cross_entropy
+from gradix.nn.modules import Linear, Module
+from gradix.nn.optim import SGD
+
+
+class TinyClassifier(Module):
+	def __init__(self):
+		self.layer = Linear(2, 2, bias=True)
+
+	def parameters(self):
+		return self.layer.parameters()
+
+	def forward(self, x):
+		return self.layer(x)
+
+
+x = Tensor(np.array([[1, 0], [0, 1], [1, 1], [0, 0]], dtype=np.float32))
+targets = Tensor(np.array([0, 1, 0, 1]))
+
+model = TinyClassifier()
+optimizer = SGD(model.parameters(), lr=0.1)
+
+for step in range(20):
+	logits = model(x)
+	loss = cross_entropy(logits, targets)
+
+	optimizer.zero_grad()
+	loss.backward()
+	optimizer.step()
+
+	print(f"step {step + 1}: loss={float(loss.data):.4f}")
+```
+
+In each step, the model produces one score per class, cross-entropy compares those scores with the labels, and backpropagation computes parameter gradients. SGD then updates the parameters. This is a minimal educational loop; it does not include features such as batching helpers, schedulers, regularization, checkpointing, or automatic device placement.
+
+## MNIST Notebook Workflow
+
+`demos/train.ipynb` demonstrates the larger workflow:
+
+1. Load MNIST with scikit-learn's `fetch_openml` and split pixels and labels into training, validation, and test sets.
+2. Convert pixel arrays to `float32`, scale pixels into the 0-to-1 range, and convert labels to integers.
+3. Create `MNISTModel` or `MNISTConvModel`, then pass `model.parameters()` to `SGD`.
+4. Iterate over mini-batches. For each batch, create input and target tensors, run the model, calculate cross-entropy, clear gradients, call `backward()`, and update parameters.
+5. Calculate validation loss without calling `backward()`, then evaluate on the test data after training.
+6. Plot training and validation loss with Matplotlib and optionally visualize the model's graph.
+
+The notebook uses additional packages for the example workflow. Install them in the notebook's Python environment if they are not already present:
+
+```bash
+python -m pip install scikit-learn tqdm
+```
+
+Matplotlib is already a required Gradix dependency. Fetching MNIST requires an internet connection the first time. Training the full dataset, particularly with the pure-Python Conv2D implementation, may take a while.
+
+## API at a Glance
+
+| API | Purpose |
+| --- | --- |
+| `Tensor(data, requires_grad=False, name=None)` | Store array data and optionally track gradients. |
+| `tensor.backward()` | Compute gradients from a scalar result back through its graph. |
+| `tensor.zero_grad()` | Clear gradients in the tensor's graph. |
+| `Module` | Base class for models with `forward()` and `parameters()`. |
+| `Linear(in_features, out_features, bias=False)` | Fully connected transformation. |
+| `Conv2D(in_channels, out_channels, kernel, stride=1, padding=0, use_bias=False)` | Two-dimensional convolution. |
+| `cross_entropy(logits, targets)` | Mean multiclass cross-entropy for integer class labels. |
+| `SGD(parameters, lr=0.001)` | Basic stochastic gradient descent optimizer. |
+| `visualize_computation_graph(tensor, ...)` | Return or render the tensor's operation graph. |
+
+## Project Layout
+
+```text
+src/gradix/             Library implementation
+  tensor.py             Tensor and automatic differentiation
+  nn/                   Functions, modules, losses, and optimizer
+demos/MnistModel.py     Example fully connected and convolutional models
+demos/demo.ipynb        Small tensor and graph examples
+demos/digraph_test.ipynb Attention-style and recurrent graph examples
+demos/train.ipynb       MNIST data preparation, training, and evaluation
+tests/test_conv2d.py    Conv2D comparison against PyTorch
+```
+
+## Development Checks
+
+Install the project and run the Conv2D comparison from the repository root:
+
+```bash
+python -m pip install -e .
+python -m pip install torch
+PYTHONPATH=src python tests/test_conv2d.py
+```
+
+The test uses PyTorch only as a reference implementation. The library itself does not depend on PyTorch.
+
+## Current Limitations
+
+- Gradix is educational software and implements a small subset of tensor and neural-network functionality.
+- CPU execution is the most suitable current path. CUDA support is experimental; Conv2D currently constructs some intermediate indices with NumPy and has not been validated on a real CUDA device.
+- Models explicitly implement `parameters()`; nested layers are not automatically discovered.
+- Dataset loading and notebook conveniences such as scikit-learn and tqdm are separate from the core package dependencies.
+- The Graphviz Python package does not include the Graphviz system executable.
+
+## License
+
+See [LICENSE](LICENSE).

gradix-0.1.0.dist-info/RECORD

@@ -0,0 +1,20 @@
+gradix/__init__.py,sha256=xXZ0omtpooFzkQyyqvHMnxaC1ZmDoyMTtsaUvawuoyI,548
+gradix/init.py,sha256=GI2JmWUXLcvDArKQIddKrrp1hnKqJg4B4xW9eunwOjc,909
+gradix/tensor.py,sha256=7vB61EMfRenYUvWW5arjAwUZjJnFgqyJAFvAewYAyTc,13056
+gradix/utils.py,sha256=QfmnE8EYhIXH7WvjZur3EConSgjTmYTTGtehZ1taHyQ,1408
+gradix/nn/__init__.py,sha256=lLxMRv99l_Y97XdYoJgSUObIDdOyhp1OQHT9SEcuVms,289
+gradix/nn/functional.py,sha256=7lYDRynk-f8yse0IteUT05AchXXrCmeuJNIaZhAxVf0,3488
+gradix/nn/loss/__init__.py,sha256=COTJEVrxp_TnKZNkP8G8mOPdpwUgkMaEcYNa7g4UPzs,72
+gradix/nn/loss/crossentropyloss.py,sha256=7sBLVlYk9jGQqnoSL61XgXksCe5YC1xDk3y9ZwRraOU,937
+gradix/nn/modules/__init__.py,sha256=I3btKSOvnVogXeBfAfj2pmu9IYkKi9bQ2iZd9nst7n0,166
+gradix/nn/modules/conv2d.py,sha256=xdUITDRnEACCQnlhAnm6YrJYVKF8xOCLYWPu8aLny_k,2836
+gradix/nn/modules/linear.py,sha256=96tk0ELPHkjYYhHQB7w_sP5HgFGXDrRl7oh8oSZaEME,1007
+gradix/nn/modules/module.py,sha256=mk1pLooenoWMr3LeyqPIhWKp-IVqn1FRLH40EeuVgFU,365
+gradix/nn/modules/upsample.py,sha256=GHkL_8FowMnlWfN41Zz9E7GVi9rJoPTiIQCluQnpF-w,402
+gradix/nn/optim/SGD.py,sha256=ZI4c6ZKvyvR0dnZm5LDvGIOKd0vHcVEjLSeEsfBIEAU,407
+gradix/nn/optim/__init__.py,sha256=s4ZojzOft3rmfAlZzq7M-zJS-Jyi88rAQnIA46QUMy4,39
+gradix-0.1.0.dist-info/licenses/LICENSE,sha256=KTS90TmJspuDVTKJ2bWxw6khU19Dh2LYNb2kfqX8TQk,1076
+gradix-0.1.0.dist-info/METADATA,sha256=BTmytZ4S0PkcpAooUSCTpAE26AfPmDB3CCMMi9jTGi4,12541
+gradix-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+gradix-0.1.0.dist-info/top_level.txt,sha256=yUMOJsHaTXqqfKsHI-x15ERiCLfs3khC7Cl42NYx43k,7
+gradix-0.1.0.dist-info/RECORD,,

gradix-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.2.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

gradix-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Gradix contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

gradix-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+gradix