deep-atomic 0.1.1__tar.gz → 0.2.0__tar.gz

{deep_atomic-0.1.1 → deep_atomic-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: deep_atomic
-Version: 0.1.1
+Version: 0.2.0
 Summary: A simple deep learning framework built upon numpy only
 Author-email: Elliot Zhang <elliot_zh@proton.me>
 License: MIT
@@ -18,12 +18,21 @@ Requires-Dist: pytest>=7.0; extra == 'dev'
 Requires-Dist: ruff>=0.1.0; extra == 'dev'
 Description-Content-Type: text/markdown
 
-# deep-atomic
+# Deep Atomic
 
 A simple deep learning framework built upon numpy only. Mainly for practice and learning.
 
 ## Usage
 
+### Installation
+
+```bash
+pip install deep-atomic
+
+# or using uv
+uv add deep-atomic
+```
+
 ### Import
 
 ```python
@@ -48,44 +57,84 @@ Most essential deep‑learning operations are implemented. For those that also e
 ```python
 a, b = da.Tensor(np.random.rand(3, 4)), da.Tensor(np.random.rand(3, 4))
 
-# Arithmetic & math
 c = a + b                     # addition
 c = a - b                     # subtraction
 c = a * b                     # element‑wise multiplication
 c = a / b                     # element‑wise division
 c = a ** b                    # element‑wise power
 c = a @ b                     # matrix multiplication
+
 c = da.exp(a)
 c = da.log(a)
+c = da.sin(a)
+# the same for cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh, arcsinh, arccosh, arctanh
+
+d = a < b                     # element-wise comparison, create a boolean tensor
+e = a <= b
+c = a > b
+c = a >= b
+c = a == b
+c = a != b
+c = da.fmax(a, b)             # IMPORTANT: here da.fmax is identical to da.maximum, for simplicity. same for da.fmin / da.minimum
+c = da.maximum(a, b)
+c = da.fmin(a, b)
+c = da.minimum(a, b)
+
+c = d & e                     # element-wise and
+c = d | e                     # element-wise or
+c = d ^ e                     # element-wise xor
+c = da.logical_not(d)         # element-wise not
+c = d.all(axis=-1)            # reduction of and operation
+c = d.any(axis=-1)            # reduction of or operation
+c = da.where(d, a, b)         # return elements chosen from a or b depending on condition
+
+c = da.topk(a, 2, axis=-1, largest=True)    # same as pytorch. axis=-1, largest=True by default
 
-# Reductions
 c = da.sum(a)                               # shape: (1,)
 c = da.sum(a, axis=1)                       # shape: (3,)
 c = da.sum(a, axis=1, keepdims=True)        # shape: (3, 1)
 # min, max, argmin, argmax follow the same signature
 
-# Softmax
-c = da.softmax(a, axis=-1)
-c = da.log_softmax(a, axis=-1)
+c = da.softmax(a, axis=-1, temperature=0.6) # support temperature. temperature=1 by default
+c = da.log_softmax(a, axis=-1, temperature=0.6)
+
+c = da.sigmoid(a)
+c = da.silu(a)
+c = da.relu(a)
+c = da.gelu(a) # Deep Atomic uses the tanh approximation for speed and convenience
 
-# Shape manipulations
 c = a.reshape(2, 6)
 c = a.reshape(1, 12).squeeze(0)               # shape: (12,)
-c = da.expand_dims(a, -1)                   # shape: (3, 4, 1)
+c = da.expand_dims(a, -1)                     # shape: (3, 4, 1)
 c = a.expand_dims(-1)                         # method‑style alternative
 c = a.repeat(2, axis=1)                       # shape: (3, 8)
-c = da.tile(a, (2, 2))                      # shape: (6, 8)
+c = da.tile(a, (2, 2))                        # shape: (6, 8)
 c = a.tile(2, 2)                              # method‑style alternative
 ```
 
+### Autograd
+
+Autograd is supported via computational graph.
+_Currently only support scalar source points._
+
+```python
+x = Tensor(np.random.rand(3, 4)) # requires_grad == True by default
+res = ... # some calculation related to x. res is a **scalar** result
+res.backward()
+print(res.grad) # gradient get!
+```
+
 ## Todo
 
 - [ ] more basic operations and their autograd
+    - [x] topk
+    - [x] boolean operations
+    - [x] gather or take_along_axis
+    - [ ] scatter
     - [ ] convolution and 2d convolution
-    - [ ] topk
-    - [ ] boolean and masked operation
+    - [ ] direct masking and indexing via `[]` syntax
     - [ ] einsum
-    - [ ] scatter
+- [ ] support backward with Vector-Jacobian Product like pytorch
 - [ ] basic neural network classes
 - [ ] optimizers and loss functions
 - [ ] full training test
@@ -93,25 +142,23 @@ c = a.tile(2, 2)                              # method‑style alternative
 - [ ] attention layers and full LLM training
 - [ ] finer type annotations, comments and documentation
 
-## Installation
-
-```bash
-pip install deep-atomic
-
-# or with uv
-uv add deep-atomic
-```
-
 ## Development
 
-Install in editable mode with development dependencies:
+Recommend manage dependencies using [uv](https://github.com/astral-sh/uv).
 
 ```bash
-pip install -e ".[dev]"
+uv sync
 ```
 
 Run tests:
 
 ```bash
+cd tests
 pytest
 ```
+
+Build wheels:
+
+```bash
+uv build
+```

{deep_atomic-0.1.1 → deep_atomic-0.2.0}/README.md

@@ -1,9 +1,18 @@
-# deep-atomic
+# Deep Atomic
 
 A simple deep learning framework built upon numpy only. Mainly for practice and learning.
 
 ## Usage
 
+### Installation
+
+```bash
+pip install deep-atomic
+
+# or using uv
+uv add deep-atomic
+```
+
 ### Import
 
 ```python
@@ -28,44 +37,84 @@ Most essential deep‑learning operations are implemented. For those that also e
 ```python
 a, b = da.Tensor(np.random.rand(3, 4)), da.Tensor(np.random.rand(3, 4))
 
-# Arithmetic & math
 c = a + b                     # addition
 c = a - b                     # subtraction
 c = a * b                     # element‑wise multiplication
 c = a / b                     # element‑wise division
 c = a ** b                    # element‑wise power
 c = a @ b                     # matrix multiplication
+
 c = da.exp(a)
 c = da.log(a)
+c = da.sin(a)
+# the same for cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh, arcsinh, arccosh, arctanh
+
+d = a < b                     # element-wise comparison, create a boolean tensor
+e = a <= b
+c = a > b
+c = a >= b
+c = a == b
+c = a != b
+c = da.fmax(a, b)             # IMPORTANT: here da.fmax is identical to da.maximum, for simplicity. same for da.fmin / da.minimum
+c = da.maximum(a, b)
+c = da.fmin(a, b)
+c = da.minimum(a, b)
+
+c = d & e                     # element-wise and
+c = d | e                     # element-wise or
+c = d ^ e                     # element-wise xor
+c = da.logical_not(d)         # element-wise not
+c = d.all(axis=-1)            # reduction of and operation
+c = d.any(axis=-1)            # reduction of or operation
+c = da.where(d, a, b)         # return elements chosen from a or b depending on condition
+
+c = da.topk(a, 2, axis=-1, largest=True)    # same as pytorch. axis=-1, largest=True by default
 
-# Reductions
 c = da.sum(a)                               # shape: (1,)
 c = da.sum(a, axis=1)                       # shape: (3,)
 c = da.sum(a, axis=1, keepdims=True)        # shape: (3, 1)
 # min, max, argmin, argmax follow the same signature
 
-# Softmax
-c = da.softmax(a, axis=-1)
-c = da.log_softmax(a, axis=-1)
+c = da.softmax(a, axis=-1, temperature=0.6) # support temperature. temperature=1 by default
+c = da.log_softmax(a, axis=-1, temperature=0.6)
+
+c = da.sigmoid(a)
+c = da.silu(a)
+c = da.relu(a)
+c = da.gelu(a) # Deep Atomic uses the tanh approximation for speed and convenience
 
-# Shape manipulations
 c = a.reshape(2, 6)
 c = a.reshape(1, 12).squeeze(0)               # shape: (12,)
-c = da.expand_dims(a, -1)                   # shape: (3, 4, 1)
+c = da.expand_dims(a, -1)                     # shape: (3, 4, 1)
 c = a.expand_dims(-1)                         # method‑style alternative
 c = a.repeat(2, axis=1)                       # shape: (3, 8)
-c = da.tile(a, (2, 2))                      # shape: (6, 8)
+c = da.tile(a, (2, 2))                        # shape: (6, 8)
 c = a.tile(2, 2)                              # method‑style alternative
 ```
 
+### Autograd
+
+Autograd is supported via computational graph.
+_Currently only support scalar source points._
+
+```python
+x = Tensor(np.random.rand(3, 4)) # requires_grad == True by default
+res = ... # some calculation related to x. res is a **scalar** result
+res.backward()
+print(res.grad) # gradient get!
+```
+
 ## Todo
 
 - [ ] more basic operations and their autograd
+    - [x] topk
+    - [x] boolean operations
+    - [x] gather or take_along_axis
+    - [ ] scatter
     - [ ] convolution and 2d convolution
-    - [ ] topk
-    - [ ] boolean and masked operation
+    - [ ] direct masking and indexing via `[]` syntax
     - [ ] einsum
-    - [ ] scatter
+- [ ] support backward with Vector-Jacobian Product like pytorch
 - [ ] basic neural network classes
 - [ ] optimizers and loss functions
 - [ ] full training test
@@ -73,25 +122,23 @@ c = a.tile(2, 2)                              # method‑style alternative
 - [ ] attention layers and full LLM training
 - [ ] finer type annotations, comments and documentation
 
-## Installation
-
-```bash
-pip install deep-atomic
-
-# or with uv
-uv add deep-atomic
-```
-
 ## Development
 
-Install in editable mode with development dependencies:
+Recommend manage dependencies using [uv](https://github.com/astral-sh/uv).
 
 ```bash
-pip install -e ".[dev]"
+uv sync
 ```
 
 Run tests:
 
 ```bash
+cd tests
 pytest
 ```
+
+Build wheels:
+
+```bash
+uv build
+```

{deep_atomic-0.1.1 → deep_atomic-0.2.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "deep_atomic"
-version = "0.1.1"
+version = "0.2.0"
 description = "A simple deep learning framework built upon numpy only"
 readme = "README.md"
 requires-python = ">=3.8"

{deep_atomic-0.1.1 → deep_atomic-0.2.0}/src/deep_atomic/__init__.py

@@ -1,6 +1,6 @@
 """deep-atomic - A simple deep learning framework built upon numpy only"""
 
-__version__ = "0.1.1"
+__version__ = "0.2.0"
 __author__ = "elliot_zh@proton.me"
 
 from .tensor import *

{deep_atomic-0.1.1 → deep_atomic-0.2.0}/src/deep_atomic/graph.py

@@ -97,6 +97,115 @@ class Log(SingleOp):
         self.input.backward(1 / self.input.to_np() * grad)
 
 
+class Sigmoid(SingleOp):
+    def __init__(self, input: Tensor, output_np: np.ndarray):
+        super().__init__(input)
+        self.output_np = output_np
+
+    def backward(self, grad):
+        # TODO: input's actual value is no longer useful during backward, and can be released to increase memory efficiency
+        self.input.backward(self.output_np * (1 - self.output_np) * grad)
+        
+        
+class Sin(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(np.cos(self.input.to_np()) * grad)
+        
+
+class Cos(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(-np.sin(self.input.to_np()) * grad)
+        
+        
+class Tan(SingleOp):
+    def __init__(self, input: Tensor, output_np: np.ndarray):
+        super().__init__(input)
+        self.output_np = output_np
+        
+    def backward(self, grad):
+        self.input.backward((1 + self.output_np ** 2) * grad) # faster since only polynomials are concerned
+        
+        
+class Arcsin(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(1 / (1 - self.input.to_np() ** 2) ** 0.5 * grad)
+        
+        
+class Arccos(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(-1 / (1 - self.input.to_np() ** 2) ** 0.5 * grad)
+        
+        
+class Arctan(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(1 / (self.input.to_np() ** 2 + 1) * grad)
+        
+
+class Sinh(SingleOp):
+    def __init__(self, input: Tensor, output_np: np.ndarray):
+        super().__init__(input)
+        self.output_np = output_np
+        
+    def backward(self, grad):
+        self.input.backward((1 + self.output_np ** 2) ** 0.5 * grad) # faster since only polynomials are concerned
+        
+        
+class Cosh(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(np.sinh(self.input.to_np()) * grad)
+
+
+class Tanh(SingleOp):
+    def __init__(self, input: Tensor, output_np: np.ndarray):
+        super().__init__(input)
+        self.output_np = output_np
+        
+    def backward(self, grad):
+        self.input.backward((1 - self.output_np ** 2) * grad) # faster since only polynomials are concerned
+        
+        
+class Arcsinh(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(1 / (self.input.to_np() ** 2 + 1) ** 0.5 * grad)
+
+
+class Arccosh(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(1 / (self.input.to_np() ** 2 - 1) ** 0.5 * grad)
+
+
+class Arctanh(SingleOp):
+    def __init__(self, input: Tensor):
+        super().__init__(input)
+        
+    def backward(self, grad):
+        self.input.backward(1 / (1 - self.input.to_np() ** 2) * grad)
+
+
 class Pow(TwoOp):
     def __init__(self, a1, a2):
         super().__init__(a1, a2)
@@ -111,6 +220,7 @@ class Pow(TwoOp):
 class Sum(SingleOp):
     def __init__(self, input: Tensor, axis, keepdims):
         super().__init__(input)
+        # TODO: need to handle when axis is a tuple of int that means reducing mutiple dims at a time
         self.axis = axis
         self.keepdims = keepdims
 
@@ -244,3 +354,28 @@ class Tile(SingleOp):
         if grad.ndim < len(self.reps):
             grad = grad.reshape(self.input.shape)
         self.input.backward(grad)
+
+# the condition is not differentiatable, so still two-operanded
+class Where(TwoOp):
+    def __init__(self, condition, a1, a2):
+        super().__init__(a1, a2)
+        if isinstance(condition, Tensor):
+            condition = condition.to_np()
+        self.condition = condition
+        
+    def backward(self, grad):
+        if self.a1 is not None:
+            self.a1.backward(np.where(self.condition, grad, 0))
+        if self.a2 is not None:
+            self.a2.backward(np.where(self.condition, 0, grad))
+            
+
+class TakeAlongAxis(SingleOp):
+    def __init__(self, input: Tensor, indices, axis):
+        super().__init__(input)
+        self.indices, self.axis = indices, axis
+        
+    def backward(self, grad):
+        grad_ = np.zeros(self.input.shape)
+        np.put_along_axis(grad_, self.indices, grad, self.axis)
+        self.input.backward(grad_)

deep_atomic-0.2.0/src/deep_atomic/op.py

@@ -0,0 +1,230 @@
+import numpy as np
+from .tensor import *
+from .graph import *
+from .utils import *
+
+
+# implement add, sub, mul, div etc. here
+def add(a1, a2):
+    return a1.__array_ufunc__(np.add, "__call__", a1, a2)
+
+
+def sub(a1, a2):
+    return a1.__array_ufunc__(np.subtract, "__call__", a1, a2)
+
+
+def mul(a1, a2):
+    return a1.__array_ufunc__(np.multiply, "__call__", a1, a2)
+
+
+def div(a1, a2):
+    return a1.__array_ufunc__(np.divide, "__call__", a1, a2)
+
+
+def matmul(a1, a2):
+    return a1.__array_ufunc__(np.matmul, "__call__", a1, a2)
+
+
+def pow(a1, a2):
+    return a1.__array_ufunc__(np.pow, "__call__", a1, a2)
+
+
+def exp(input):
+    return input.__array_ufunc__(np.exp, "__call__", input)
+
+
+def log(input):
+    return input.__array_ufunc__(np.log, "__call__", input)
+
+
+# we do not implement exp2, log2, log10 etc. for that this framework is mainly for neural network training and for simplicity
+
+
+def logical_not(input):
+    return input.__array_ufunc(np.logical_not, "__call__", input)
+
+def all(input, axis=None):
+    return input.all(axis)
+
+def any(input, axis=None):
+    return input.any(axis)
+
+
+def sum(input: Tensor, axis=None, keepdims=False):
+    return input.__array_ufunc__(np.add, "reduce", input, axis=axis, keepdims=keepdims)
+
+
+def max(input: Tensor, axis=None, keepdims=False):
+    res = Tensor(
+        np.max(input.to_np(), axis=axis, keepdims=keepdims), requires_grad=False
+    )
+    if input.requires_grad:
+        res.requires_grad = True
+        if axis == None:
+            res.dep = MinMax(input, None, keepdims, full_red_value=res.to_np())
+        else:
+            indices = np.argmax(input.to_np(), axis=axis, keepdims=True)
+            res.dep = MinMax(input, axis, keepdims, indices=indices)
+    return res
+
+
+def min(input: Tensor, axis=None, keepdims=False):
+    res = Tensor(
+        np.min(input.to_np(), axis=axis, keepdims=keepdims), requires_grad=False
+    )
+    if input.requires_grad:
+        res.requires_grad = True
+        if axis == None:
+            res.dep = MinMax(input, None, keepdims, full_red_value=res.to_np())
+        else:
+            indices = np.argmin(input.to_np(), axis=axis, keepdims=True)
+            res.dep = MinMax(input, axis, keepdims, indices=indices)
+    return res
+
+
+def argmax(input: Tensor, axis=-1, keepdims=False):
+    # no gradient concerned when performing argmax
+    return Tensor(
+        np.argmax(input.to_np(), axis=axis, keepdims=keepdims), requires_grad=False
+    )
+
+
+def argmin(input: Tensor, axis=-1, keepdims=False):
+    # no gradient concerned when performing armin
+    return Tensor(
+        np.argmin(input.to_np(), axis=axis, keepdims=keepdims), requires_grad=False
+    )
+    
+    
+def fmax(a1, a2):
+    return a1.__array_ufunc__(np.fmax, '__call__', a1, a2)
+
+
+def fmin(a1, a2):
+    return a1.__array_ufunc__(np.fmin, '__call__', a1, a2)
+
+# unlike pytorch, fmax and maximum, fmin and minimum are the same since Deep Atomic currently doesn't handle NaN
+maximum = fmax
+minimum = fmin
+
+
+def softmax(input: Tensor, axis=-1, temperature=1):
+    input /= temperature
+    input = exp(input - max(input, axis=axis, keepdims=True))
+    return input / sum(input, axis=axis, keepdims=True)
+
+
+def log_softmax(input: Tensor, axis=-1, temperature=1):
+    input /= temperature
+    input = input - max(input, axis=axis, keepdims=True)
+    return input - log(sum(exp(input), axis=axis, keepdims=True))
+
+
+def sigmoid(input: Tensor):
+    res = 1 / (1 + exp(-input))
+    if res.requires_grad:
+        res.dep = Sigmoid(input, res.to_np())
+    return res
+
+
+def sin(input):
+    return input.__array_ufunc__(np.sin, "__call__", input)
+
+
+def cos(input):
+    return input.__array_ufunc__(np.cos, "__call__", input)
+
+
+def tan(input):
+    return input.__array_ufunc__(np.tan, "__call__", input)
+
+def arcsin(input):
+    return input.__array_ufunc__(np.arcsin, "__call__", input)
+
+def arccos(input):
+    return input.__array_ufunc__(np.arccos, "__call__", input)
+
+def arctan(input):
+    return input.__array_ufunc__(np.arctan, "__call__", input)
+
+def sinh(input):
+    return input.__array_ufunc__(np.sinh, "__call__", input)
+
+def cosh(input):
+    return input.__array_ufunc__(np.cosh, "__call__", input)
+
+def tanh(input):
+    return input.__array_ufunc__(np.tanh, "__call__", input)
+
+def arcsinh(input):
+    return input.__array_ufunc__(np.arcsinh, "__call__", input)
+
+def arccosh(input):
+    return input.__array_ufunc__(np.arccosh, "__call__", input)
+
+def arctanh(input):
+    return input.__array_ufunc__(np.arctanh, "__call__", input)
+
+
+def relu(input: Tensor):
+    condition = input > 0
+    return where(condition, input, 0)
+
+
+def silu(input: Tensor):
+    return input * sigmoid(input)
+
+
+def gelu(input: Tensor):
+    return 0.5 * input * (1 + tanh(0.7978845608028654 * (input + 0.44715 * input ** 3)))
+
+
+def reshape(input: Tensor, target_shape):
+    return input.reshape(*target_shape)
+
+
+def squeeze(input: Tensor, axis):
+    return input.squeeze(axis)
+
+
+def expand_dims(input: Tensor, axis):
+    return input.expand_dims(axis)
+
+
+def repeats(input: Tensor, repeats, axis=None):
+    return input.repeat(repeats, axis)
+
+
+def tile(input: Tensor, reps):
+    return input.tile(*reps)
+
+
+def where(condition, a1, a2):
+    requires_grad = False
+    for i in (a1, a2):
+        if isinstance(i, Tensor) and i.requires_grad:
+            requires_grad = True
+    res = Tensor(np.where(condition, a1, a2), requires_grad=requires_grad)
+    if res.requires_grad: res.dep = Where(condition, a1, a2)
+    return res
+
+# TODO: API design. Should I follow pytorch and add a gather?
+def take_along_axis(input: Tensor, indices, axis=-1):
+    if isinstance(indices,Tensor): indices = indices.to_np() # cut off gradient
+    res = Tensor(np.take_along_axis(input, indices, axis), requires_grad=input.requires_grad)
+    if res.requires_grad: res.dep = TakeAlongAxis(input, indices, axis)
+    return res
+
+# TODO: Tensor.implement scatter_
+
+def topk(input: Tensor, kth, axis=-1, largest=True):
+    if largest:
+        indices = np.argpartition(input, input.shape[axis] - kth, axis)
+        indices_idx = [slice(None)] * input.ndim
+        indices_idx[axis] = slice(-kth, None)
+    else:
+        indices = np.argpartition(input, kth, axis)
+        indices_idx = [slice(None)] * input.ndim
+        indices_idx[axis] = slice(kth)
+    indices = indices[tuple(indices_idx)]
+    return take_along_axis(input, indices, axis), Tensor(indices, requires_grad=False)