deep-atomic 0.1.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.1.0/.github/workflows/publish-pypi.yml +43 -0
- deep_atomic-0.1.0/.gitignore +20 -0
- deep_atomic-0.1.0/.vscode/settings.json +7 -0
- deep_atomic-0.1.0/LICENSE +21 -0
- deep_atomic-0.1.0/PKG-INFO +117 -0
- deep_atomic-0.1.0/README.md +97 -0
- deep_atomic-0.1.0/pyproject.toml +44 -0
- deep_atomic-0.1.0/src/deep_atomic/__init__.py +7 -0
- deep_atomic-0.1.0/src/deep_atomic/graph.py +246 -0
- deep_atomic-0.1.0/src/deep_atomic/op.py +117 -0
- deep_atomic-0.1.0/src/deep_atomic/tensor.py +168 -0
- deep_atomic-0.1.0/src/deep_atomic/utils.py +17 -0
- deep_atomic-0.1.0/tests/__init__.py +1 -0
- deep_atomic-0.1.0/tests/backward_test.py +408 -0
- deep_atomic-0.1.0/tests/init_test.py +40 -0
- deep_atomic-0.1.0/tests/op_test.py +203 -0
- deep_atomic-0.1.0/tests/utils.py +36 -0
- deep_atomic-0.1.0/uv.lock +1249 -0
|
@@ -0,0 +1,43 @@
|
|
|
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 # install uv
|
|
18
|
+
- name: build
|
|
19
|
+
run: uv build
|
|
20
|
+
# refer to https://docs.pypi.org/trusted-publishers/using-a-publisher/
|
|
21
|
+
- name: mint API token
|
|
22
|
+
id: mint-token
|
|
23
|
+
run: |
|
|
24
|
+
# retrieve the ambient OIDC token
|
|
25
|
+
resp=$(curl -H "Authorization: bearer $ACTIONS_ID_TOKEN_REQUEST_TOKEN" \
|
|
26
|
+
"$ACTIONS_ID_TOKEN_REQUEST_URL&audience=pypi")
|
|
27
|
+
oidc_token=$(jq -r '.value' <<< "${resp}")
|
|
28
|
+
|
|
29
|
+
# exchange the OIDC token for an API token
|
|
30
|
+
resp=$(curl -X POST https://pypi.org/_/oidc/mint-token -d "{\"token\": \"${oidc_token}\"}")
|
|
31
|
+
api_token=$(jq -r '.token' <<< "${resp}")
|
|
32
|
+
|
|
33
|
+
# mask the newly minted API token, so that we don't accidentally leak it
|
|
34
|
+
echo "::add-mask::${api_token}"
|
|
35
|
+
|
|
36
|
+
# see the next step in the workflow for an example of using this step output
|
|
37
|
+
echo "api-token=${api_token}" >> "${GITHUB_OUTPUT}"
|
|
38
|
+
|
|
39
|
+
- name: publish
|
|
40
|
+
# gh-action-pypi-publish uses TWINE_PASSWORD automatically
|
|
41
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
42
|
+
with:
|
|
43
|
+
password: ${{ steps.mint-token.outputs.api-token }}
|
|
@@ -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.
|
|
@@ -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,97 @@
|
|
|
1
|
+
# deep-atomic
|
|
2
|
+
|
|
3
|
+
A simple deep learning framework built upon numpy only. Mainly for practice and learning.
|
|
4
|
+
|
|
5
|
+
## Usage
|
|
6
|
+
|
|
7
|
+
### Import
|
|
8
|
+
|
|
9
|
+
```python
|
|
10
|
+
import deep_atomic as da
|
|
11
|
+
import numpy as np # required for tensor initialization and some operations
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
### Creating a Tensor
|
|
15
|
+
|
|
16
|
+
```python
|
|
17
|
+
# Create a tensor from a NumPy array (requires_grad=True by default)
|
|
18
|
+
a = da.Tensor(
|
|
19
|
+
np.array([1, 2, 3], dtype=np.float64),
|
|
20
|
+
requires_grad=True
|
|
21
|
+
)
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
### Supported Operations
|
|
25
|
+
|
|
26
|
+
Most essential deep‑learning operations are implemented. For those that also exist in NumPy, we follow NumPy's API conventions.
|
|
27
|
+
|
|
28
|
+
```python
|
|
29
|
+
a, b = da.Tensor(np.random.rand(3, 4)), da.Tensor(np.random.rand(3, 4))
|
|
30
|
+
|
|
31
|
+
# Arithmetic & math
|
|
32
|
+
c = a + b # addition
|
|
33
|
+
c = a - b # subtraction
|
|
34
|
+
c = a * b # element‑wise multiplication
|
|
35
|
+
c = a / b # element‑wise division
|
|
36
|
+
c = a ** b # element‑wise power
|
|
37
|
+
c = a @ b # matrix multiplication
|
|
38
|
+
c = da.exp(a)
|
|
39
|
+
c = da.log(a)
|
|
40
|
+
|
|
41
|
+
# Reductions
|
|
42
|
+
c = da.sum(a) # shape: (1,)
|
|
43
|
+
c = da.sum(a, axis=1) # shape: (3,)
|
|
44
|
+
c = da.sum(a, axis=1, keepdims=True) # shape: (3, 1)
|
|
45
|
+
# min, max, argmin, argmax follow the same signature
|
|
46
|
+
|
|
47
|
+
# Softmax
|
|
48
|
+
c = da.softmax(a, axis=-1)
|
|
49
|
+
c = da.log_softmax(a, axis=-1)
|
|
50
|
+
|
|
51
|
+
# Shape manipulations
|
|
52
|
+
c = a.reshape(2, 6)
|
|
53
|
+
c = a.reshape(1, 12).squeeze(0) # shape: (12,)
|
|
54
|
+
c = da.expand_dims(a, -1) # shape: (3, 4, 1)
|
|
55
|
+
c = a.expand_dims(-1) # method‑style alternative
|
|
56
|
+
c = a.repeat(2, axis=1) # shape: (3, 8)
|
|
57
|
+
c = da.tile(a, (2, 2)) # shape: (6, 8)
|
|
58
|
+
c = a.tile(2, 2) # method‑style alternative
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
## Todo
|
|
62
|
+
|
|
63
|
+
- [ ] more basic operations and their autograd
|
|
64
|
+
- [ ] convolution and 2d convolution
|
|
65
|
+
- [ ] topk
|
|
66
|
+
- [ ] boolean and masked operation
|
|
67
|
+
- [ ] einsum
|
|
68
|
+
- [ ] scatter
|
|
69
|
+
- [ ] basic neural network classes
|
|
70
|
+
- [ ] optimizers and loss functions
|
|
71
|
+
- [ ] full training test
|
|
72
|
+
- [ ] benchmark with pytorch on CPU
|
|
73
|
+
- [ ] attention layers and full LLM training
|
|
74
|
+
- [ ] finer type annotations, comments and documentation
|
|
75
|
+
|
|
76
|
+
## Installation
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
pip install deep_atomic
|
|
80
|
+
|
|
81
|
+
# or with uv
|
|
82
|
+
uv add deep_atomic
|
|
83
|
+
```
|
|
84
|
+
|
|
85
|
+
## Development
|
|
86
|
+
|
|
87
|
+
Install in editable mode with development dependencies:
|
|
88
|
+
|
|
89
|
+
```bash
|
|
90
|
+
pip install -e ".[dev]"
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
Run tests:
|
|
94
|
+
|
|
95
|
+
```bash
|
|
96
|
+
pytest
|
|
97
|
+
```
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "deep_atomic"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "A simple deep learning framework built upon numpy only"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.8"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [
|
|
13
|
+
{ name = "Elliot Zhang", email = "elliot_zh@proton.me" }
|
|
14
|
+
]
|
|
15
|
+
classifiers = [
|
|
16
|
+
"Programming Language :: Python :: 3",
|
|
17
|
+
"License :: OSI Approved :: MIT License",
|
|
18
|
+
"Operating System :: OS Independent",
|
|
19
|
+
]
|
|
20
|
+
dependencies = [
|
|
21
|
+
"numpy>=1.24.4",
|
|
22
|
+
"pytest>=8.3.5",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
[project.optional-dependencies]
|
|
26
|
+
dev = [
|
|
27
|
+
"pytest>=7.0",
|
|
28
|
+
"pytest-cov>=4.0",
|
|
29
|
+
"black>=23.0",
|
|
30
|
+
"ruff>=0.1.0",
|
|
31
|
+
]
|
|
32
|
+
|
|
33
|
+
[tool.black]
|
|
34
|
+
line-length = 88
|
|
35
|
+
target-version = ['py38']
|
|
36
|
+
|
|
37
|
+
[tool.ruff]
|
|
38
|
+
line-length = 88
|
|
39
|
+
lint.select = ["E", "F", "W", "I"]
|
|
40
|
+
lint.fixable = ["I"]
|
|
41
|
+
|
|
42
|
+
[tool.pytest.ini_options]
|
|
43
|
+
testpaths = ["tests"]
|
|
44
|
+
python_files = "*_test.py"
|
|
@@ -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)
|