deep-atomic 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,7 @@
1
+ """deep-atomic - A simple deep learning framework built upon numpy only"""
2
+
3
+ __version__ = "0.1.0"
4
+ __author__ = "elliot_zh@proton.me"
5
+
6
+ from .tensor import *
7
+ from .op import *
deep_atomic/graph.py ADDED
@@ -0,0 +1,246 @@
1
+ import numpy as np
2
+ from .tensor import Tensor
3
+ from abc import ABC, abstractmethod
4
+
5
+ from .utils import *
6
+
7
+
8
+ class Op(ABC):
9
+ @abstractmethod
10
+ def backward(self, grad):
11
+ pass
12
+
13
+
14
+ class SingleOp(Op):
15
+ def __init__(self, input: Tensor):
16
+ self.input = input
17
+ self.input.depended_count += 1
18
+
19
+
20
+ class TwoOp(Op):
21
+ def __init__(self, a1, a2):
22
+ self.a1_np, self.a2_np = (
23
+ i.to_np() if isinstance(i, Tensor) else i for i in [a1, a2]
24
+ )
25
+ self.a1, self.a2 = (i if isinstance(i, Tensor) else None for i in [a1, a2])
26
+ if isinstance(self.a1, Tensor):
27
+ self.a1.depended_count += 1
28
+ if isinstance(self.a2, Tensor):
29
+ self.a2.depended_count += 1
30
+
31
+
32
+ class Add(TwoOp):
33
+ def __init__(self, a1, a2, sub=False):
34
+ super().__init__(a1, a2)
35
+ self.sub = sub
36
+
37
+ def backward(self, grad):
38
+ if self.a1 is not None:
39
+ self.a1.backward(grad)
40
+ if self.a2 is None:
41
+ return
42
+ if self.sub:
43
+ grad = -grad
44
+ self.a2.backward(grad)
45
+
46
+
47
+ class Mul(TwoOp):
48
+ def __init__(self, a1, a2):
49
+ super().__init__(a1, a2)
50
+
51
+ def backward(self, grad):
52
+ if self.a1 is not None:
53
+ self.a1.backward(grad * self.a2_np)
54
+ if self.a2 is not None:
55
+ self.a2.backward(self.a1_np * grad)
56
+
57
+
58
+ class Div(TwoOp):
59
+ def __init__(self, a1, a2):
60
+ super().__init__(a1, a2)
61
+
62
+ def backward(self, grad):
63
+ if self.a1 is not None:
64
+ self.a1.backward(grad / self.a2_np)
65
+ if self.a2 is not None:
66
+ self.a2.backward(-1 * self.a1_np / (self.a2_np**2) * grad)
67
+
68
+
69
+ class MatMul(TwoOp):
70
+ def __init__(self, a1, a2):
71
+ super().__init__(a1, a2)
72
+
73
+ def backward(self, grad):
74
+ if self.a1 is not None:
75
+ self.a1.backward(grad @ self.a2_np.swapaxes(-1, -2))
76
+ if self.a2 is not None:
77
+ self.a2.backward(self.a1_np.swapaxes(-1, -2) @ grad)
78
+
79
+
80
+ class Exp(SingleOp):
81
+ def __init__(
82
+ self, input: Tensor
83
+ ): # node of exp is introduced only when the input is a Tensor
84
+ super().__init__(input)
85
+
86
+ def backward(self, grad):
87
+ self.input.backward(np.exp(self.input.to_np()) * grad)
88
+
89
+
90
+ class Log(SingleOp):
91
+ def __init__(
92
+ self, input: Tensor
93
+ ): # node of log is introduced only when the input is a Tensor
94
+ super().__init__(input)
95
+
96
+ def backward(self, grad):
97
+ self.input.backward(1 / self.input.to_np() * grad)
98
+
99
+
100
+ class Pow(TwoOp):
101
+ def __init__(self, a1, a2):
102
+ super().__init__(a1, a2)
103
+
104
+ def backward(self, grad):
105
+ if self.a1 is not None:
106
+ self.a1.backward(self.a2_np * (self.a1_np) ** (self.a2_np - 1) * grad)
107
+ if self.a2 is not None:
108
+ self.a2.backward(np.log(self.a1_np) * (self.a1_np) ** (self.a2_np) * grad)
109
+
110
+
111
+ class Sum(SingleOp):
112
+ def __init__(self, input: Tensor, axis, keepdims):
113
+ super().__init__(input)
114
+ self.axis = axis
115
+ self.keepdims = keepdims
116
+
117
+ def backward(self, grad):
118
+ if self.axis is None or self.keepdims:
119
+ self.input.backward(grad)
120
+ else:
121
+ self.input.backward(np.expand_dims(grad, self.axis))
122
+
123
+
124
+ class Abs(SingleOp):
125
+ def __init__(self, input: Tensor):
126
+ super().__init__(input)
127
+
128
+ def backward(self, grad):
129
+ self.input.backward(np.where(grad < 0, -grad, grad))
130
+
131
+
132
+ class MinMax(SingleOp):
133
+ def __init__(
134
+ self, input: Tensor, axis, keepdims, indices=None, full_red_value=None
135
+ ):
136
+ super().__init__(input)
137
+ self.indices = indices # result of argmin / argmax. when axis is not None
138
+ self.axis, self.keepdims = axis, keepdims
139
+ self.full_red_value = full_red_value # for faster compute when axis is None
140
+ if axis is None:
141
+ assert full_red_value is not None
142
+ else:
143
+ assert indices is not None
144
+
145
+ def backward(self, grad):
146
+ if self.axis is None:
147
+ if not self.keepdims:
148
+ grad = np.tile(grad, self.input.shape)
149
+ mask = self.input.to_np() == self.full_red_value
150
+ grad /= np.count_nonzero(
151
+ mask
152
+ ) # distribute gradient evenly across max values. following pytorch implementation
153
+ else:
154
+ if not self.keepdims:
155
+ tile_reps = [1] * self.input.ndim
156
+ tile_reps[self.axis] = self.input.shape[self.axis]
157
+ grad = np.expand_dims(grad, axis=self.axis)
158
+ grad = np.tile(grad, tile_reps)
159
+ baseline_expansion_axis = list(range(self.input.ndim))
160
+ baseline_expansion_axis.pop(self.axis)
161
+ baseline_indices = np.expand_dims(
162
+ np.arange(self.input.shape[self.axis]), axis=baseline_expansion_axis
163
+ )
164
+ mask = self.indices == baseline_indices
165
+ self.input.backward(np.where(mask, grad, 0.0))
166
+
167
+
168
+ class Reshape(SingleOp):
169
+ def __init__(self, input: Tensor, from_):
170
+ super().__init__(input)
171
+ self.from_ = from_
172
+
173
+ def backward(self, grad):
174
+ self.input.backward(grad.reshape(self.from_))
175
+
176
+
177
+ class Squeeze(SingleOp):
178
+ def __init__(self, input: Tensor, axis):
179
+ super().__init__(input)
180
+ self.axis = axis
181
+
182
+ def backward(self, grad):
183
+ self.input.backward(np.expand_dims(grad, self.axis))
184
+
185
+
186
+ class ExpandDims(SingleOp):
187
+ def __init__(self, input: Tensor, axis):
188
+ super().__init__(input)
189
+ self.axis = axis
190
+
191
+ def backward(self, grad):
192
+ self.input.backward(grad.squeeze(self.axis))
193
+
194
+
195
+ class Repeat(SingleOp):
196
+ def __init__(self, input: Tensor, repeats, axis=None):
197
+ super().__init__(input)
198
+ self.repeats = repeats
199
+ self.axis = axis
200
+
201
+ def backward(self, grad):
202
+ if isinstance(self.repeats, int):
203
+ if self.axis is None:
204
+ grad = grad.reshape(-1, self.repeats).sum(axis=-1)
205
+ else:
206
+ shape = list(grad.shape)
207
+ shape[self.axis] //= self.repeats
208
+ shape.insert(self.axis + 1, self.repeats)
209
+ grad = grad.reshape(shape).sum(axis=self.axis + 1)
210
+ self.input.backward(grad)
211
+ else:
212
+ if self.axis is None:
213
+ grad_ = np.zeros(len(self.repeats), dtype=grad.dtype)
214
+ start = 0
215
+ for i, r in enumerate(self.repeats):
216
+ if r > 0:
217
+ grad_[i] = grad[start : start + r].sum()
218
+ start += r
219
+ else:
220
+ moved = np.moveaxis(grad, self.axis, 0)
221
+ grad_ = np.zeros(self.input.shape, dtype=grad.dtype)
222
+ start = 0
223
+ for i, r in enumerate(self.repeats):
224
+ if r > 0:
225
+ idx = [slice(None)] * self.input.ndim
226
+ idx[self.axis] = i
227
+ grad_[tuple(idx)] = moved[start : start + r].sum(axis=0)
228
+ start += r
229
+ self.input.backward(grad_)
230
+
231
+
232
+ class Tile(SingleOp):
233
+ def __init__(self, input: Tensor, reps):
234
+ super().__init__(input)
235
+ self.reps = reps
236
+
237
+ def backward(self, grad):
238
+ start = (grad.ndim - len(self.reps)) if grad.ndim > len(self.reps) else 0
239
+ for i, r in enumerate(self.reps, start=start):
240
+ shape = list(grad.shape)
241
+ shape[i] //= r
242
+ shape.insert(i + 1, r)
243
+ grad = grad.reshape(shape, order="F").sum(axis=i + 1)
244
+ if grad.ndim < len(self.reps):
245
+ grad = grad.reshape(self.input.shape)
246
+ self.input.backward(grad)
deep_atomic/op.py ADDED
@@ -0,0 +1,117 @@
1
+ import numpy as np
2
+ from .tensor import *
3
+ from .graph import *
4
+
5
+
6
+ # implement add, sub, mul, div etc. here
7
+ def add(a1, a2):
8
+ return a1.__array_ufunc__(np.add, "__call__", a1, a2)
9
+
10
+
11
+ def sub(a1, a2):
12
+ return a1.__array_ufunc__(np.subtract, "__call__", a1, a2)
13
+
14
+
15
+ def mul(a1, a2):
16
+ return a1.__array_ufunc__(np.multiply, "__call__", a1, a2)
17
+
18
+
19
+ def div(a1, a2):
20
+ return a1.__array_ufunc__(np.divide, "__call__", a1, a2)
21
+
22
+
23
+ def matmul(a1, a2):
24
+ return a1.__array_ufunc__(np.matmul, "__call__", a1, a2)
25
+
26
+
27
+ def pow(a1, a2):
28
+ return a1.__array_ufunc__(np.pow, "__call__", a1, a2)
29
+
30
+
31
+ def exp(input):
32
+ # TODO: handle other methods for ufunc (or use torch like api?)
33
+ return input.__array_ufunc__(np.exp, "__call__", input)
34
+
35
+
36
+ def log(input):
37
+ # TODO: handle other methods for ufunc (or use torch like api?)
38
+ return input.__array_ufunc__(np.log, "__call__", input)
39
+
40
+
41
+ # we do not implement exp2, log2, log10 etc. for that this framework is mainly for neural network training and for simplicity
42
+
43
+
44
+ def sum(input: Tensor, axis=None, keepdims=False):
45
+ return input.__array_ufunc__(np.add, "reduce", input, axis=axis, keepdims=keepdims)
46
+
47
+
48
+ def max(input: Tensor, axis=None, keepdims=False):
49
+ res = Tensor(
50
+ np.max(input.to_np(), axis=axis, keepdims=keepdims), requires_grad=False
51
+ )
52
+ if input.requires_grad:
53
+ res.requires_grad = True
54
+ if axis == None:
55
+ res.dep = MinMax(input, None, keepdims, full_red_value=res.to_np())
56
+ else:
57
+ indices = np.argmax(input.to_np(), axis=axis, keepdims=True)
58
+ res.dep = MinMax(input, axis, keepdims, indices=indices)
59
+ return res
60
+
61
+
62
+ def min(input: Tensor, axis=None, keepdims=False):
63
+ res = Tensor(
64
+ np.min(input.to_np(), axis=axis, keepdims=keepdims), requires_grad=False
65
+ )
66
+ if input.requires_grad:
67
+ res.requires_grad = True
68
+ if axis == None:
69
+ res.dep = MinMax(input, None, keepdims, full_red_value=res.to_np())
70
+ else:
71
+ indices = np.argmin(input.to_np(), axis=axis, keepdims=True)
72
+ res.dep = MinMax(input, axis, keepdims, indices=indices)
73
+ return res
74
+
75
+
76
+ def argmax(input: Tensor, axis=-1, keepdims=False):
77
+ # no gradient concerned when performing argmax
78
+ return Tensor(
79
+ np.argmax(input.to_np(), axis=axis, keepdims=keepdims), requires_grad=False
80
+ )
81
+
82
+
83
+ def argmin(input: Tensor, axis=-1, keepdims=False):
84
+ # no gradient concerned when performing armin
85
+ return Tensor(
86
+ np.argmin(input.to_np(), axis=axis, keepdims=keepdims), requires_grad=False
87
+ )
88
+
89
+
90
+ def softmax(input: Tensor, axis=-1):
91
+ input = exp(input - max(input, axis=axis, keepdims=True))
92
+ return input / sum(input, axis=axis, keepdims=True)
93
+
94
+
95
+ def log_softmax(input: Tensor, axis=-1):
96
+ input = input - max(input, axis=axis, keepdims=True)
97
+ return input - log(sum(exp(input), axis=axis, keepdims=True))
98
+
99
+
100
+ def reshape(input: Tensor, target_shape):
101
+ return input.reshape(*target_shape)
102
+
103
+
104
+ def squeeze(input: Tensor, axis):
105
+ return input.squeeze(axis)
106
+
107
+
108
+ def expand_dims(input: Tensor, axis):
109
+ return input.expand_dims(axis)
110
+
111
+
112
+ def repeats(input: Tensor, repeats, axis=None):
113
+ return input.repeat(repeats, axis)
114
+
115
+
116
+ def tile(input: Tensor, reps):
117
+ return input.tile(*reps)
deep_atomic/tensor.py ADDED
@@ -0,0 +1,168 @@
1
+ import numpy as np
2
+
3
+ from .utils import *
4
+
5
+
6
+ class Tensor(np.ndarray):
7
+ def __new__(
8
+ subtype,
9
+ arg0, # arg0: ndarray | shape
10
+ requires_grad=True, # for backward
11
+ dep=None, # for backward
12
+ dtype=np.float64,
13
+ buffer=None,
14
+ offset=0,
15
+ strides=None,
16
+ order=None,
17
+ ):
18
+ if isinstance(arg0, np.ndarray):
19
+ obj = arg0.view(subtype) # convert from ndarray
20
+ elif isinstance(arg0, np.generic):
21
+ obj = np.array(arg0).view(subtype)
22
+ else:
23
+ obj = super().__new__(subtype, arg0, dtype, buffer, offset, strides, order)
24
+
25
+ obj.requires_grad, obj.dep = requires_grad, dep
26
+ obj.depended_count = 0 # for topo sorting in backward
27
+ return obj
28
+
29
+ @property
30
+ def requires_grad(self):
31
+ return self._requires_grad
32
+
33
+ @requires_grad.setter
34
+ def requires_grad(self, value: bool):
35
+ self._requires_grad = value
36
+ if value:
37
+ # IMPORTANT: will reset gradient to zero
38
+ self.grad = np.zeros(self.shape)
39
+ else:
40
+ self.grad = None
41
+
42
+ def __array_finalize__(self, obj):
43
+ if obj is None:
44
+ return
45
+ # FIXME: ? when creating a view, we cannot rely on python's default reference mechanism. must share identical depended_count across views. should depended_count be a list?
46
+ self.requires_grad = getattr(obj, "requires_grad", True)
47
+ self.dep = getattr(obj, "dep", None)
48
+ self.depended_count = getattr(obj, "depended_cound", 0)
49
+ self.grad = getattr(obj, "grad", None)
50
+
51
+ def to_np(self):
52
+ return self.view(np.ndarray)
53
+
54
+ def backward(self, grad=None):
55
+ if not self.requires_grad:
56
+ return
57
+ if grad is None: # source point of compute graph
58
+ # TODO: support Vector-Jacobian Product like pytorch
59
+ assert self.size == 1 # must be scalar
60
+ self.grad += 1
61
+ # otherwise receive gradient from graph
62
+ elif grad.size == 1:
63
+ self.grad += np.tile(grad, self.shape)
64
+ elif is_broadcastable(grad.shape, self.shape):
65
+ self.grad += grad
66
+ elif grad.shape != self.shape:
67
+ grad_ = grad
68
+ dim_to_reduce = detect_broadcast_dim(self.shape, grad.shape)
69
+ for dim, keepdims in dim_to_reduce:
70
+ grad_ = grad_.sum(axis=dim, keepdims=keepdims)
71
+ self.grad += grad_
72
+ else:
73
+ raise ValueError(
74
+ f"Invalid gradient shape. Expected {self.shape}, but got {grad.shape}"
75
+ )
76
+ self.depended_count -= 1
77
+ if self.dep is not None and self.depended_count <= 0:
78
+ self.dep.backward(self.grad) # trigger graph
79
+
80
+ # override operations
81
+ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
82
+ # type conversion
83
+ inputs_np = []
84
+ requires_grad = self.requires_grad
85
+ for input_ in inputs:
86
+ if isinstance(input_, Tensor):
87
+ inputs_np.append(input_.to_np())
88
+ requires_grad = requires_grad or input_.requires_grad
89
+ else:
90
+ inputs_np.append(input_)
91
+
92
+ result_np = super().__array_ufunc__(ufunc, method, *inputs_np, **kwargs)
93
+
94
+ if not requires_grad:
95
+ return Tensor(result_np, requires_grad=False)
96
+
97
+ # if requires_grad, construct graph
98
+ # TODO: implement gradient computation for other methods / ufuncs
99
+ if ufunc is np.add:
100
+ if method == "__call__":
101
+ res = Tensor(result_np, dep=Add(*inputs))
102
+ elif method == "reduce":
103
+ axis = getattr(kwargs, "axis", None)
104
+ keepdims = getattr(kwargs, "keepdims", False)
105
+ res = Tensor(result_np, dep=Sum(*inputs, axis, keepdims))
106
+ else:
107
+ return NotImplemented
108
+ elif ufunc is np.subtract:
109
+ res = Tensor(result_np, dep=Add(*inputs, sub=True))
110
+ elif ufunc is np.multiply:
111
+ res = Tensor(result_np, dep=Mul(*inputs))
112
+ elif ufunc is np.divide:
113
+ res = Tensor(result_np, dep=Div(*inputs))
114
+ elif ufunc is np.matmul:
115
+ res = Tensor(result_np, dep=MatMul(*inputs))
116
+ elif ufunc is np.exp:
117
+ res = Tensor(result_np, dep=Exp(*inputs))
118
+ elif ufunc is np.exp2:
119
+ res = Tensor(result_np, dep=Exp(Mul(*inputs, np.log(2))))
120
+ elif ufunc is np.log:
121
+ res = Tensor(result_np, dep=Log(*inputs))
122
+ elif ufunc is np.log2:
123
+ res = Tensor(result_np, dep=Div(Log(*inputs), np.log(2)))
124
+ elif ufunc is np.log10:
125
+ res = Tensor(result_np, dep=Div(Log(*inputs), np.log(10)))
126
+ elif ufunc is np.pow:
127
+ res = Tensor(result_np, dep=Pow(*inputs))
128
+ elif ufunc is np.square:
129
+ res = Tensor(result_np, dep=Pow(*inputs, 2))
130
+ elif ufunc is np.abs:
131
+ res = Tensor(result_np, dep=Abs(*inputs))
132
+ else:
133
+ return NotImplemented
134
+
135
+ return res
136
+
137
+ def reshape(self, *target_shape):
138
+ res = Tensor(super().reshape(*target_shape), requires_grad=self.requires_grad)
139
+ if self.requires_grad:
140
+ res.dep = Reshape(self, self.shape)
141
+ return res
142
+
143
+ def squeeze(self, axis):
144
+ res = Tensor(super().squeeze(axis), requires_grad=self.requires_grad)
145
+ if self.requires_grad:
146
+ res.dep = Squeeze(self, axis)
147
+ return res
148
+
149
+ def expand_dims(self, axis):
150
+ res = Tensor(np.expand_dims(super(), axis), requires_grad=self.requires_grad)
151
+ if self.requires_grad:
152
+ res.dep = ExpandDims(self, axis)
153
+ return res
154
+
155
+ def repeat(self, repeats, axis=None):
156
+ res = Tensor(super().repeat(repeats, axis), requires_grad=self.requires_grad)
157
+ if self.requires_grad:
158
+ res.dep = Repeat(self, repeats, axis)
159
+ return res
160
+
161
+ def tile(self, *reps):
162
+ res = Tensor(np.tile(super(), reps), requires_grad=self.requires_grad)
163
+ if self.requires_grad:
164
+ res.dep = Tile(self, reps)
165
+ return res
166
+
167
+
168
+ from .graph import * # avoid looped dependencies
deep_atomic/utils.py ADDED
@@ -0,0 +1,17 @@
1
+ def detect_broadcast_dim(from_, to_):
2
+ detected = []
3
+ for i in range(-1, -len(to_) - 1, -1):
4
+ if i < -len(from_):
5
+ detected.append((i + len(to_), False))
6
+ elif from_[i] == 1 and to_[i] != 1:
7
+ detected.append((i + len(to_), True)) # True tp keepdims
8
+ elif from_[i] != to_[i]:
9
+ raise ValueError(f"{from_} cannot be broadcasted to target shape {to_}")
10
+ return detected
11
+
12
+
13
+ def is_broadcastable(from_, to_):
14
+ res = True
15
+ for i in range(-1, -len(to_) - 1, -1):
16
+ res = res and (from_[i] == to_[i] or from_[i] == 1 and to_[i] != 1)
17
+ return res
@@ -0,0 +1,117 @@
1
+ Metadata-Version: 2.4
2
+ Name: deep_atomic
3
+ Version: 0.1.0
4
+ Summary: A simple deep learning framework built upon numpy only
5
+ Author-email: Elliot Zhang <elliot_zh@proton.me>
6
+ License: MIT
7
+ License-File: LICENSE
8
+ Classifier: License :: OSI Approved :: MIT License
9
+ Classifier: Operating System :: OS Independent
10
+ Classifier: Programming Language :: Python :: 3
11
+ Requires-Python: >=3.8
12
+ Requires-Dist: numpy>=1.24.4
13
+ Requires-Dist: pytest>=8.3.5
14
+ Provides-Extra: dev
15
+ Requires-Dist: black>=23.0; extra == 'dev'
16
+ Requires-Dist: pytest-cov>=4.0; extra == 'dev'
17
+ Requires-Dist: pytest>=7.0; extra == 'dev'
18
+ Requires-Dist: ruff>=0.1.0; extra == 'dev'
19
+ Description-Content-Type: text/markdown
20
+
21
+ # deep-atomic
22
+
23
+ A simple deep learning framework built upon numpy only. Mainly for practice and learning.
24
+
25
+ ## Usage
26
+
27
+ ### Import
28
+
29
+ ```python
30
+ import deep_atomic as da
31
+ import numpy as np # required for tensor initialization and some operations
32
+ ```
33
+
34
+ ### Creating a Tensor
35
+
36
+ ```python
37
+ # Create a tensor from a NumPy array (requires_grad=True by default)
38
+ a = da.Tensor(
39
+ np.array([1, 2, 3], dtype=np.float64),
40
+ requires_grad=True
41
+ )
42
+ ```
43
+
44
+ ### Supported Operations
45
+
46
+ Most essential deep‑learning operations are implemented. For those that also exist in NumPy, we follow NumPy's API conventions.
47
+
48
+ ```python
49
+ a, b = da.Tensor(np.random.rand(3, 4)), da.Tensor(np.random.rand(3, 4))
50
+
51
+ # Arithmetic & math
52
+ c = a + b # addition
53
+ c = a - b # subtraction
54
+ c = a * b # element‑wise multiplication
55
+ c = a / b # element‑wise division
56
+ c = a ** b # element‑wise power
57
+ c = a @ b # matrix multiplication
58
+ c = da.exp(a)
59
+ c = da.log(a)
60
+
61
+ # Reductions
62
+ c = da.sum(a) # shape: (1,)
63
+ c = da.sum(a, axis=1) # shape: (3,)
64
+ c = da.sum(a, axis=1, keepdims=True) # shape: (3, 1)
65
+ # min, max, argmin, argmax follow the same signature
66
+
67
+ # Softmax
68
+ c = da.softmax(a, axis=-1)
69
+ c = da.log_softmax(a, axis=-1)
70
+
71
+ # Shape manipulations
72
+ c = a.reshape(2, 6)
73
+ c = a.reshape(1, 12).squeeze(0) # shape: (12,)
74
+ c = da.expand_dims(a, -1) # shape: (3, 4, 1)
75
+ c = a.expand_dims(-1) # method‑style alternative
76
+ c = a.repeat(2, axis=1) # shape: (3, 8)
77
+ c = da.tile(a, (2, 2)) # shape: (6, 8)
78
+ c = a.tile(2, 2) # method‑style alternative
79
+ ```
80
+
81
+ ## Todo
82
+
83
+ - [ ] more basic operations and their autograd
84
+ - [ ] convolution and 2d convolution
85
+ - [ ] topk
86
+ - [ ] boolean and masked operation
87
+ - [ ] einsum
88
+ - [ ] scatter
89
+ - [ ] basic neural network classes
90
+ - [ ] optimizers and loss functions
91
+ - [ ] full training test
92
+ - [ ] benchmark with pytorch on CPU
93
+ - [ ] attention layers and full LLM training
94
+ - [ ] finer type annotations, comments and documentation
95
+
96
+ ## Installation
97
+
98
+ ```bash
99
+ pip install deep_atomic
100
+
101
+ # or with uv
102
+ uv add deep_atomic
103
+ ```
104
+
105
+ ## Development
106
+
107
+ Install in editable mode with development dependencies:
108
+
109
+ ```bash
110
+ pip install -e ".[dev]"
111
+ ```
112
+
113
+ Run tests:
114
+
115
+ ```bash
116
+ pytest
117
+ ```
@@ -0,0 +1,9 @@
1
+ deep_atomic/__init__.py,sha256=qYv3Pu95A5_TZfrkC0J3O7YHs__vw4AUed20GQMoaKw,174
2
+ deep_atomic/graph.py,sha256=o-ZClWo0QY4Pl8AkhZ5VHHRo8Lx9gYIYKIFUoe-nkFw,7573
3
+ deep_atomic/op.py,sha256=rXQvrwkoCj9rLGBzysBfyF0hozusD1xwsoY3QzrSFFU,3262
4
+ deep_atomic/tensor.py,sha256=JftulVvUghQZSFjdnknLO4xUSoLQdfgeo5tzgtYn3r0,6145
5
+ deep_atomic/utils.py,sha256=9IHdlybDbzPcb2UnaXvAeGYXjyGZisQV-rxBVhUyaT4,615
6
+ deep_atomic-0.1.0.dist-info/METADATA,sha256=29hTZMGvoNmqTgKSvbkAINPWBFH53V9pL3QLVz7grRc,3106
7
+ deep_atomic-0.1.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
8
+ deep_atomic-0.1.0.dist-info/licenses/LICENSE,sha256=v2spsd7N1pKFFh2G8wGP_45iwe5S0DYiJzG4im8Rupc,1066
9
+ deep_atomic-0.1.0.dist-info/RECORD,,
@@ -0,0 +1,4 @@
1
+ Wheel-Version: 1.0
2
+ Generator: hatchling 1.31.0
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Your Name
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.