forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
elementwise_sum_op.cc
116 lines (83 loc) · 2.15 KB
/
elementwise_sum_op.cc
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
#include "caffe2/operators/utility_ops.h"
namespace caffe2 {
REGISTER_CPU_OPERATOR(Sum, SumOp<CPUContext>);
OPERATOR_SCHEMA(Sum)
.NumInputs(1, INT_MAX)
.NumOutputs(1)
.AllowInplace({{0, 0}})
.CostInferenceFunction(CostInferenceForSum)
.InputsCanCrossDevices()
.IdenticalTypeAndShapeOfInput(0)
.SetDoc(R"DOC(
Element-wise sum of each of the input tensors. The first input tensor can be used
in-place as the output tensor, in which case the sum will be done in place and
results will be accumulated the first input tensor. All inputs and outputs must
have the same shape and data type.
Github Links:
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/elementwise_sum_op.cc
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Sum",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([[1,2],[3,4]]).astype(np.float32))
workspace.FeedBlob("B", np.array([[5,6],[7,8]]).astype(np.float32))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("A"))
```
**Result**
```
A: [[1. 2.]
[3. 4.]]
B: [[5. 6.]
[7. 8.]]
C: [[1. 2.]
[3. 4.]]
```
</details>
<details>
<summary> <b>Example 2</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Sum",
["A", "B"],
["A"], // inplace
)
workspace.FeedBlob("A", np.array([[1,2,5],[8,3,4]]).astype(np.float32))
workspace.FeedBlob("B", np.array([[9,5,6],[6,7,8]]).astype(np.float32))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("A after Sum:", workspace.FetchBlob("A"))
```
**Result**
```
A: [[1. 2. 5.]
[8. 3. 4.]]
B: [[9. 5. 6.]
[6. 7. 8.]]
A after Sum: [[10. 7. 11.]
[14. 10. 12.]]
```
</details>
)DOC")
.Input(
0,
"A",
"*(type: Tensor`<float>`)* First tensor to be added element-wise.")
.Input(
1,
"B",
"*(type: Tensor`<float>`)* Second tensor to be added element-wise.")
.Output(0, "C", "*(type: Tensor`<float>`)* Sum of A and B.")
.InheritOnnxSchema();
}