deep-atomic 0.1.0__tar.gz → 0.2.0__tar.gz
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-0.2.0/.github/workflows/publish-pypi.yml +21 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/PKG-INFO +71 -24
- deep_atomic-0.2.0/README.md +144 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/pyproject.toml +1 -1
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/src/deep_atomic/__init__.py +1 -1
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/src/deep_atomic/graph.py +135 -0
- deep_atomic-0.2.0/src/deep_atomic/op.py +230 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/src/deep_atomic/tensor.py +73 -2
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/src/deep_atomic/utils.py +1 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/tests/backward_test.py +281 -3
- deep_atomic-0.2.0/tests/op_test.py +492 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/uv.lock +2 -2
- deep_atomic-0.1.0/.github/workflows/publish-pypi.yml +0 -43
- deep_atomic-0.1.0/README.md +0 -97
- deep_atomic-0.1.0/src/deep_atomic/op.py +0 -117
- deep_atomic-0.1.0/tests/op_test.py +0 -203
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/.gitignore +0 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/.vscode/settings.json +0 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/LICENSE +0 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/tests/__init__.py +0 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/tests/init_test.py +0 -0
- {deep_atomic-0.1.0 → deep_atomic-0.2.0}/tests/utils.py +0 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
on: [release]
|
|
3
|
+
jobs:
|
|
4
|
+
pypi-publish:
|
|
5
|
+
name: upload release to PyPI
|
|
6
|
+
runs-on: ubuntu-latest
|
|
7
|
+
environment: pypi
|
|
8
|
+
permissions:
|
|
9
|
+
# IMPORTANT: this permission is mandatory for Trusted Publishing
|
|
10
|
+
id-token: write
|
|
11
|
+
steps:
|
|
12
|
+
- uses: actions/checkout@v6
|
|
13
|
+
- uses: actions/setup-python@v4
|
|
14
|
+
with:
|
|
15
|
+
python-version: "3.11"
|
|
16
|
+
- name: deps
|
|
17
|
+
run: curl -LsSf https://astral.sh/uv/install.sh | sh
|
|
18
|
+
- name: build
|
|
19
|
+
run: uv build
|
|
20
|
+
- name: Publish package distributions to PyPI
|
|
21
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: deep_atomic
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.2.0
|
|
4
4
|
Summary: A simple deep learning framework built upon numpy only
|
|
5
5
|
Author-email: Elliot Zhang <elliot_zh@proton.me>
|
|
6
6
|
License: MIT
|
|
@@ -18,12 +18,21 @@ Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
|
18
18
|
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
19
19
|
Description-Content-Type: text/markdown
|
|
20
20
|
|
|
21
|
-
#
|
|
21
|
+
# Deep Atomic
|
|
22
22
|
|
|
23
23
|
A simple deep learning framework built upon numpy only. Mainly for practice and learning.
|
|
24
24
|
|
|
25
25
|
## Usage
|
|
26
26
|
|
|
27
|
+
### Installation
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
pip install deep-atomic
|
|
31
|
+
|
|
32
|
+
# or using uv
|
|
33
|
+
uv add deep-atomic
|
|
34
|
+
```
|
|
35
|
+
|
|
27
36
|
### Import
|
|
28
37
|
|
|
29
38
|
```python
|
|
@@ -48,44 +57,84 @@ Most essential deep‑learning operations are implemented. For those that also e
|
|
|
48
57
|
```python
|
|
49
58
|
a, b = da.Tensor(np.random.rand(3, 4)), da.Tensor(np.random.rand(3, 4))
|
|
50
59
|
|
|
51
|
-
# Arithmetic & math
|
|
52
60
|
c = a + b # addition
|
|
53
61
|
c = a - b # subtraction
|
|
54
62
|
c = a * b # element‑wise multiplication
|
|
55
63
|
c = a / b # element‑wise division
|
|
56
64
|
c = a ** b # element‑wise power
|
|
57
65
|
c = a @ b # matrix multiplication
|
|
66
|
+
|
|
58
67
|
c = da.exp(a)
|
|
59
68
|
c = da.log(a)
|
|
69
|
+
c = da.sin(a)
|
|
70
|
+
# the same for cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh, arcsinh, arccosh, arctanh
|
|
71
|
+
|
|
72
|
+
d = a < b # element-wise comparison, create a boolean tensor
|
|
73
|
+
e = a <= b
|
|
74
|
+
c = a > b
|
|
75
|
+
c = a >= b
|
|
76
|
+
c = a == b
|
|
77
|
+
c = a != b
|
|
78
|
+
c = da.fmax(a, b) # IMPORTANT: here da.fmax is identical to da.maximum, for simplicity. same for da.fmin / da.minimum
|
|
79
|
+
c = da.maximum(a, b)
|
|
80
|
+
c = da.fmin(a, b)
|
|
81
|
+
c = da.minimum(a, b)
|
|
82
|
+
|
|
83
|
+
c = d & e # element-wise and
|
|
84
|
+
c = d | e # element-wise or
|
|
85
|
+
c = d ^ e # element-wise xor
|
|
86
|
+
c = da.logical_not(d) # element-wise not
|
|
87
|
+
c = d.all(axis=-1) # reduction of and operation
|
|
88
|
+
c = d.any(axis=-1) # reduction of or operation
|
|
89
|
+
c = da.where(d, a, b) # return elements chosen from a or b depending on condition
|
|
90
|
+
|
|
91
|
+
c = da.topk(a, 2, axis=-1, largest=True) # same as pytorch. axis=-1, largest=True by default
|
|
60
92
|
|
|
61
|
-
# Reductions
|
|
62
93
|
c = da.sum(a) # shape: (1,)
|
|
63
94
|
c = da.sum(a, axis=1) # shape: (3,)
|
|
64
95
|
c = da.sum(a, axis=1, keepdims=True) # shape: (3, 1)
|
|
65
96
|
# min, max, argmin, argmax follow the same signature
|
|
66
97
|
|
|
67
|
-
#
|
|
68
|
-
c = da.
|
|
69
|
-
|
|
98
|
+
c = da.softmax(a, axis=-1, temperature=0.6) # support temperature. temperature=1 by default
|
|
99
|
+
c = da.log_softmax(a, axis=-1, temperature=0.6)
|
|
100
|
+
|
|
101
|
+
c = da.sigmoid(a)
|
|
102
|
+
c = da.silu(a)
|
|
103
|
+
c = da.relu(a)
|
|
104
|
+
c = da.gelu(a) # Deep Atomic uses the tanh approximation for speed and convenience
|
|
70
105
|
|
|
71
|
-
# Shape manipulations
|
|
72
106
|
c = a.reshape(2, 6)
|
|
73
107
|
c = a.reshape(1, 12).squeeze(0) # shape: (12,)
|
|
74
|
-
c = da.expand_dims(a, -1)
|
|
108
|
+
c = da.expand_dims(a, -1) # shape: (3, 4, 1)
|
|
75
109
|
c = a.expand_dims(-1) # method‑style alternative
|
|
76
110
|
c = a.repeat(2, axis=1) # shape: (3, 8)
|
|
77
|
-
c = da.tile(a, (2, 2))
|
|
111
|
+
c = da.tile(a, (2, 2)) # shape: (6, 8)
|
|
78
112
|
c = a.tile(2, 2) # method‑style alternative
|
|
79
113
|
```
|
|
80
114
|
|
|
115
|
+
### Autograd
|
|
116
|
+
|
|
117
|
+
Autograd is supported via computational graph.
|
|
118
|
+
_Currently only support scalar source points._
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
x = Tensor(np.random.rand(3, 4)) # requires_grad == True by default
|
|
122
|
+
res = ... # some calculation related to x. res is a **scalar** result
|
|
123
|
+
res.backward()
|
|
124
|
+
print(res.grad) # gradient get!
|
|
125
|
+
```
|
|
126
|
+
|
|
81
127
|
## Todo
|
|
82
128
|
|
|
83
129
|
- [ ] more basic operations and their autograd
|
|
130
|
+
- [x] topk
|
|
131
|
+
- [x] boolean operations
|
|
132
|
+
- [x] gather or take_along_axis
|
|
133
|
+
- [ ] scatter
|
|
84
134
|
- [ ] convolution and 2d convolution
|
|
85
|
-
- [ ]
|
|
86
|
-
- [ ] boolean and masked operation
|
|
135
|
+
- [ ] direct masking and indexing via `[]` syntax
|
|
87
136
|
- [ ] einsum
|
|
88
|
-
|
|
137
|
+
- [ ] support backward with Vector-Jacobian Product like pytorch
|
|
89
138
|
- [ ] basic neural network classes
|
|
90
139
|
- [ ] optimizers and loss functions
|
|
91
140
|
- [ ] full training test
|
|
@@ -93,25 +142,23 @@ c = a.tile(2, 2) # method‑style alternative
|
|
|
93
142
|
- [ ] attention layers and full LLM training
|
|
94
143
|
- [ ] finer type annotations, comments and documentation
|
|
95
144
|
|
|
96
|
-
## Installation
|
|
97
|
-
|
|
98
|
-
```bash
|
|
99
|
-
pip install deep_atomic
|
|
100
|
-
|
|
101
|
-
# or with uv
|
|
102
|
-
uv add deep_atomic
|
|
103
|
-
```
|
|
104
|
-
|
|
105
145
|
## Development
|
|
106
146
|
|
|
107
|
-
|
|
147
|
+
Recommend manage dependencies using [uv](https://github.com/astral-sh/uv).
|
|
108
148
|
|
|
109
149
|
```bash
|
|
110
|
-
|
|
150
|
+
uv sync
|
|
111
151
|
```
|
|
112
152
|
|
|
113
153
|
Run tests:
|
|
114
154
|
|
|
115
155
|
```bash
|
|
156
|
+
cd tests
|
|
116
157
|
pytest
|
|
117
158
|
```
|
|
159
|
+
|
|
160
|
+
Build wheels:
|
|
161
|
+
|
|
162
|
+
```bash
|
|
163
|
+
uv build
|
|
164
|
+
```
|
|
@@ -0,0 +1,144 @@
|
|
|
1
|
+
# Deep Atomic
|
|
2
|
+
|
|
3
|
+
A simple deep learning framework built upon numpy only. Mainly for practice and learning.
|
|
4
|
+
|
|
5
|
+
## Usage
|
|
6
|
+
|
|
7
|
+
### Installation
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install deep-atomic
|
|
11
|
+
|
|
12
|
+
# or using uv
|
|
13
|
+
uv add deep-atomic
|
|
14
|
+
```
|
|
15
|
+
|
|
16
|
+
### Import
|
|
17
|
+
|
|
18
|
+
```python
|
|
19
|
+
import deep_atomic as da
|
|
20
|
+
import numpy as np # required for tensor initialization and some operations
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
### Creating a Tensor
|
|
24
|
+
|
|
25
|
+
```python
|
|
26
|
+
# Create a tensor from a NumPy array (requires_grad=True by default)
|
|
27
|
+
a = da.Tensor(
|
|
28
|
+
np.array([1, 2, 3], dtype=np.float64),
|
|
29
|
+
requires_grad=True
|
|
30
|
+
)
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
### Supported Operations
|
|
34
|
+
|
|
35
|
+
Most essential deep‑learning operations are implemented. For those that also exist in NumPy, we follow NumPy's API conventions.
|
|
36
|
+
|
|
37
|
+
```python
|
|
38
|
+
a, b = da.Tensor(np.random.rand(3, 4)), da.Tensor(np.random.rand(3, 4))
|
|
39
|
+
|
|
40
|
+
c = a + b # addition
|
|
41
|
+
c = a - b # subtraction
|
|
42
|
+
c = a * b # element‑wise multiplication
|
|
43
|
+
c = a / b # element‑wise division
|
|
44
|
+
c = a ** b # element‑wise power
|
|
45
|
+
c = a @ b # matrix multiplication
|
|
46
|
+
|
|
47
|
+
c = da.exp(a)
|
|
48
|
+
c = da.log(a)
|
|
49
|
+
c = da.sin(a)
|
|
50
|
+
# the same for cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh, arcsinh, arccosh, arctanh
|
|
51
|
+
|
|
52
|
+
d = a < b # element-wise comparison, create a boolean tensor
|
|
53
|
+
e = a <= b
|
|
54
|
+
c = a > b
|
|
55
|
+
c = a >= b
|
|
56
|
+
c = a == b
|
|
57
|
+
c = a != b
|
|
58
|
+
c = da.fmax(a, b) # IMPORTANT: here da.fmax is identical to da.maximum, for simplicity. same for da.fmin / da.minimum
|
|
59
|
+
c = da.maximum(a, b)
|
|
60
|
+
c = da.fmin(a, b)
|
|
61
|
+
c = da.minimum(a, b)
|
|
62
|
+
|
|
63
|
+
c = d & e # element-wise and
|
|
64
|
+
c = d | e # element-wise or
|
|
65
|
+
c = d ^ e # element-wise xor
|
|
66
|
+
c = da.logical_not(d) # element-wise not
|
|
67
|
+
c = d.all(axis=-1) # reduction of and operation
|
|
68
|
+
c = d.any(axis=-1) # reduction of or operation
|
|
69
|
+
c = da.where(d, a, b) # return elements chosen from a or b depending on condition
|
|
70
|
+
|
|
71
|
+
c = da.topk(a, 2, axis=-1, largest=True) # same as pytorch. axis=-1, largest=True by default
|
|
72
|
+
|
|
73
|
+
c = da.sum(a) # shape: (1,)
|
|
74
|
+
c = da.sum(a, axis=1) # shape: (3,)
|
|
75
|
+
c = da.sum(a, axis=1, keepdims=True) # shape: (3, 1)
|
|
76
|
+
# min, max, argmin, argmax follow the same signature
|
|
77
|
+
|
|
78
|
+
c = da.softmax(a, axis=-1, temperature=0.6) # support temperature. temperature=1 by default
|
|
79
|
+
c = da.log_softmax(a, axis=-1, temperature=0.6)
|
|
80
|
+
|
|
81
|
+
c = da.sigmoid(a)
|
|
82
|
+
c = da.silu(a)
|
|
83
|
+
c = da.relu(a)
|
|
84
|
+
c = da.gelu(a) # Deep Atomic uses the tanh approximation for speed and convenience
|
|
85
|
+
|
|
86
|
+
c = a.reshape(2, 6)
|
|
87
|
+
c = a.reshape(1, 12).squeeze(0) # shape: (12,)
|
|
88
|
+
c = da.expand_dims(a, -1) # shape: (3, 4, 1)
|
|
89
|
+
c = a.expand_dims(-1) # method‑style alternative
|
|
90
|
+
c = a.repeat(2, axis=1) # shape: (3, 8)
|
|
91
|
+
c = da.tile(a, (2, 2)) # shape: (6, 8)
|
|
92
|
+
c = a.tile(2, 2) # method‑style alternative
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
### Autograd
|
|
96
|
+
|
|
97
|
+
Autograd is supported via computational graph.
|
|
98
|
+
_Currently only support scalar source points._
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
x = Tensor(np.random.rand(3, 4)) # requires_grad == True by default
|
|
102
|
+
res = ... # some calculation related to x. res is a **scalar** result
|
|
103
|
+
res.backward()
|
|
104
|
+
print(res.grad) # gradient get!
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
## Todo
|
|
108
|
+
|
|
109
|
+
- [ ] more basic operations and their autograd
|
|
110
|
+
- [x] topk
|
|
111
|
+
- [x] boolean operations
|
|
112
|
+
- [x] gather or take_along_axis
|
|
113
|
+
- [ ] scatter
|
|
114
|
+
- [ ] convolution and 2d convolution
|
|
115
|
+
- [ ] direct masking and indexing via `[]` syntax
|
|
116
|
+
- [ ] einsum
|
|
117
|
+
- [ ] support backward with Vector-Jacobian Product like pytorch
|
|
118
|
+
- [ ] basic neural network classes
|
|
119
|
+
- [ ] optimizers and loss functions
|
|
120
|
+
- [ ] full training test
|
|
121
|
+
- [ ] benchmark with pytorch on CPU
|
|
122
|
+
- [ ] attention layers and full LLM training
|
|
123
|
+
- [ ] finer type annotations, comments and documentation
|
|
124
|
+
|
|
125
|
+
## Development
|
|
126
|
+
|
|
127
|
+
Recommend manage dependencies using [uv](https://github.com/astral-sh/uv).
|
|
128
|
+
|
|
129
|
+
```bash
|
|
130
|
+
uv sync
|
|
131
|
+
```
|
|
132
|
+
|
|
133
|
+
Run tests:
|
|
134
|
+
|
|
135
|
+
```bash
|
|
136
|
+
cd tests
|
|
137
|
+
pytest
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
Build wheels:
|
|
141
|
+
|
|
142
|
+
```bash
|
|
143
|
+
uv build
|
|
144
|
+
```
|
|
@@ -97,6 +97,115 @@ class Log(SingleOp):
|
|
|
97
97
|
self.input.backward(1 / self.input.to_np() * grad)
|
|
98
98
|
|
|
99
99
|
|
|
100
|
+
class Sigmoid(SingleOp):
|
|
101
|
+
def __init__(self, input: Tensor, output_np: np.ndarray):
|
|
102
|
+
super().__init__(input)
|
|
103
|
+
self.output_np = output_np
|
|
104
|
+
|
|
105
|
+
def backward(self, grad):
|
|
106
|
+
# TODO: input's actual value is no longer useful during backward, and can be released to increase memory efficiency
|
|
107
|
+
self.input.backward(self.output_np * (1 - self.output_np) * grad)
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
class Sin(SingleOp):
|
|
111
|
+
def __init__(self, input: Tensor):
|
|
112
|
+
super().__init__(input)
|
|
113
|
+
|
|
114
|
+
def backward(self, grad):
|
|
115
|
+
self.input.backward(np.cos(self.input.to_np()) * grad)
|
|
116
|
+
|
|
117
|
+
|
|
118
|
+
class Cos(SingleOp):
|
|
119
|
+
def __init__(self, input: Tensor):
|
|
120
|
+
super().__init__(input)
|
|
121
|
+
|
|
122
|
+
def backward(self, grad):
|
|
123
|
+
self.input.backward(-np.sin(self.input.to_np()) * grad)
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
class Tan(SingleOp):
|
|
127
|
+
def __init__(self, input: Tensor, output_np: np.ndarray):
|
|
128
|
+
super().__init__(input)
|
|
129
|
+
self.output_np = output_np
|
|
130
|
+
|
|
131
|
+
def backward(self, grad):
|
|
132
|
+
self.input.backward((1 + self.output_np ** 2) * grad) # faster since only polynomials are concerned
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
class Arcsin(SingleOp):
|
|
136
|
+
def __init__(self, input: Tensor):
|
|
137
|
+
super().__init__(input)
|
|
138
|
+
|
|
139
|
+
def backward(self, grad):
|
|
140
|
+
self.input.backward(1 / (1 - self.input.to_np() ** 2) ** 0.5 * grad)
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
class Arccos(SingleOp):
|
|
144
|
+
def __init__(self, input: Tensor):
|
|
145
|
+
super().__init__(input)
|
|
146
|
+
|
|
147
|
+
def backward(self, grad):
|
|
148
|
+
self.input.backward(-1 / (1 - self.input.to_np() ** 2) ** 0.5 * grad)
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
class Arctan(SingleOp):
|
|
152
|
+
def __init__(self, input: Tensor):
|
|
153
|
+
super().__init__(input)
|
|
154
|
+
|
|
155
|
+
def backward(self, grad):
|
|
156
|
+
self.input.backward(1 / (self.input.to_np() ** 2 + 1) * grad)
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
class Sinh(SingleOp):
|
|
160
|
+
def __init__(self, input: Tensor, output_np: np.ndarray):
|
|
161
|
+
super().__init__(input)
|
|
162
|
+
self.output_np = output_np
|
|
163
|
+
|
|
164
|
+
def backward(self, grad):
|
|
165
|
+
self.input.backward((1 + self.output_np ** 2) ** 0.5 * grad) # faster since only polynomials are concerned
|
|
166
|
+
|
|
167
|
+
|
|
168
|
+
class Cosh(SingleOp):
|
|
169
|
+
def __init__(self, input: Tensor):
|
|
170
|
+
super().__init__(input)
|
|
171
|
+
|
|
172
|
+
def backward(self, grad):
|
|
173
|
+
self.input.backward(np.sinh(self.input.to_np()) * grad)
|
|
174
|
+
|
|
175
|
+
|
|
176
|
+
class Tanh(SingleOp):
|
|
177
|
+
def __init__(self, input: Tensor, output_np: np.ndarray):
|
|
178
|
+
super().__init__(input)
|
|
179
|
+
self.output_np = output_np
|
|
180
|
+
|
|
181
|
+
def backward(self, grad):
|
|
182
|
+
self.input.backward((1 - self.output_np ** 2) * grad) # faster since only polynomials are concerned
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
class Arcsinh(SingleOp):
|
|
186
|
+
def __init__(self, input: Tensor):
|
|
187
|
+
super().__init__(input)
|
|
188
|
+
|
|
189
|
+
def backward(self, grad):
|
|
190
|
+
self.input.backward(1 / (self.input.to_np() ** 2 + 1) ** 0.5 * grad)
|
|
191
|
+
|
|
192
|
+
|
|
193
|
+
class Arccosh(SingleOp):
|
|
194
|
+
def __init__(self, input: Tensor):
|
|
195
|
+
super().__init__(input)
|
|
196
|
+
|
|
197
|
+
def backward(self, grad):
|
|
198
|
+
self.input.backward(1 / (self.input.to_np() ** 2 - 1) ** 0.5 * grad)
|
|
199
|
+
|
|
200
|
+
|
|
201
|
+
class Arctanh(SingleOp):
|
|
202
|
+
def __init__(self, input: Tensor):
|
|
203
|
+
super().__init__(input)
|
|
204
|
+
|
|
205
|
+
def backward(self, grad):
|
|
206
|
+
self.input.backward(1 / (1 - self.input.to_np() ** 2) * grad)
|
|
207
|
+
|
|
208
|
+
|
|
100
209
|
class Pow(TwoOp):
|
|
101
210
|
def __init__(self, a1, a2):
|
|
102
211
|
super().__init__(a1, a2)
|
|
@@ -111,6 +220,7 @@ class Pow(TwoOp):
|
|
|
111
220
|
class Sum(SingleOp):
|
|
112
221
|
def __init__(self, input: Tensor, axis, keepdims):
|
|
113
222
|
super().__init__(input)
|
|
223
|
+
# TODO: need to handle when axis is a tuple of int that means reducing mutiple dims at a time
|
|
114
224
|
self.axis = axis
|
|
115
225
|
self.keepdims = keepdims
|
|
116
226
|
|
|
@@ -244,3 +354,28 @@ class Tile(SingleOp):
|
|
|
244
354
|
if grad.ndim < len(self.reps):
|
|
245
355
|
grad = grad.reshape(self.input.shape)
|
|
246
356
|
self.input.backward(grad)
|
|
357
|
+
|
|
358
|
+
# the condition is not differentiatable, so still two-operanded
|
|
359
|
+
class Where(TwoOp):
|
|
360
|
+
def __init__(self, condition, a1, a2):
|
|
361
|
+
super().__init__(a1, a2)
|
|
362
|
+
if isinstance(condition, Tensor):
|
|
363
|
+
condition = condition.to_np()
|
|
364
|
+
self.condition = condition
|
|
365
|
+
|
|
366
|
+
def backward(self, grad):
|
|
367
|
+
if self.a1 is not None:
|
|
368
|
+
self.a1.backward(np.where(self.condition, grad, 0))
|
|
369
|
+
if self.a2 is not None:
|
|
370
|
+
self.a2.backward(np.where(self.condition, 0, grad))
|
|
371
|
+
|
|
372
|
+
|
|
373
|
+
class TakeAlongAxis(SingleOp):
|
|
374
|
+
def __init__(self, input: Tensor, indices, axis):
|
|
375
|
+
super().__init__(input)
|
|
376
|
+
self.indices, self.axis = indices, axis
|
|
377
|
+
|
|
378
|
+
def backward(self, grad):
|
|
379
|
+
grad_ = np.zeros(self.input.shape)
|
|
380
|
+
np.put_along_axis(grad_, self.indices, grad, self.axis)
|
|
381
|
+
self.input.backward(grad_)
|