Caffe2 - C++ API
A deep learning, cross platform ML framework
elementwise_sum_op.cc
1 
17 #include "caffe2/operators/utility_ops.h"
18 
19 namespace caffe2 {
20 
21 namespace {
22 OpSchema::Cost CostInferenceForSum(
23  const OperatorDef& def,
24  const vector<TensorShape>& in) {
25  struct OpSchema::Cost cost = PointwiseCostInference<1>(def, in);
26  cost.flops *= (in.size() - 1);
27  cost.params_bytes = 0;
28  return cost;
29 }
30 } // namespace
31 
32 REGISTER_CPU_OPERATOR(Sum, SumOp<CPUContext>);
33 
34 OPERATOR_SCHEMA(Sum)
35  .NumInputs(1, INT_MAX)
36  .NumOutputs(1)
37  .AllowInplace({{0, 0}})
38  .CostInferenceFunction(CostInferenceForSum)
39  .InputsCanCrossDevices()
40  .IdenticalTypeAndShapeOfInput(0)
41  .SetDoc(R"DOC(
42 Element-wise sum of each of the input tensors. The first input tensor can be
43 used in-place as the output tensor, in which case the sum will be done in
44 place and results will be accumulated in input0. All inputs and outputs must
45 have the same shape and data type.
46 )DOC")
47  .Input(0, "data_0", "First of the input tensors. Can be inplace.")
48  .Output(0, "sum", "Output tensor. Same dimension as inputs.");
49 }
Copyright (c) 2016-present, Facebook, Inc.