mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
39423fbdd4
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/17684 Adding tensor and cost inference functions to more int8 operators. Reviewed By: yinghai Differential Revision: D14174746 fbshipit-source-id: dfad975fa75899565c8fb61f1b7747a9206ebd22
116 lines
2.2 KiB
C++
116 lines
2.2 KiB
C++
#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();
|
|
}
|