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.
- deep_atomic/__init__.py +7 -0
- deep_atomic/graph.py +246 -0
- deep_atomic/op.py +117 -0
- deep_atomic/tensor.py +168 -0
- deep_atomic/utils.py +17 -0
- deep_atomic-0.1.0.dist-info/METADATA +117 -0
- deep_atomic-0.1.0.dist-info/RECORD +9 -0
- deep_atomic-0.1.0.dist-info/WHEEL +4 -0
- deep_atomic-0.1.0.dist-info/licenses/LICENSE +21 -0
deep_atomic/__init__.py
ADDED
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,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.
|