Caffe2 - C++ API
A deep learning, cross platform ML framework
rowmul_op.cc
1 
17 #include "caffe2/operators/rowmul_op.h"
18 
19 namespace caffe2 {
20 namespace {
21 
22 REGISTER_CPU_OPERATOR(ReduceTailSum, ReduceTailSumOp<float, CPUContext>);
23 REGISTER_CPU_OPERATOR(RowMul, RowMulOp<float, CPUContext>);
24 
25 OPERATOR_SCHEMA(ReduceTailSum)
26  .NumInputs(1, 1)
27  .NumOutputs(1)
28  .SetDoc(R"DOC(
29 Reduce the tailing dimensions
30 )DOC")
31  .Input(0, "mat", "The matrix")
32  .Output(0, "output", "Output");
33 
34 OPERATOR_SCHEMA(RowMul)
35  .NumInputs(2, 2)
36  .NumOutputs(1)
37  .SetDoc(R"DOC(
38 Given a matrix A and column vector w, the output is the multiplication of row i
39 of A and element i of w, e.g. C[i][j] = A[i][j] * w[i]. This operator should be
40 deprecated when the gradient operator of Mul with broadcast is implemented.
41 )DOC")
42  .Input(0, "mat", "The matrix")
43  .Input(1, "w", "The column vector")
44  .Output(0, "output", "Output");
45 
46 class GetRowMulGradient : public GradientMakerBase {
47  using GradientMakerBase::GradientMakerBase;
48  vector<OperatorDef> GetGradientDefs() override {
49  return vector<OperatorDef>{
50  CreateOperatorDef(
51  "RowMul", "", vector<string>{GO(0), I(1)}, vector<string>{GI(0)}),
52  CreateOperatorDef(
53  "Mul",
54  "",
55  vector<string>{GO(0), I(0)},
56  vector<string>{GI(1) + "before_aggregate"}),
57  CreateOperatorDef(
58  "ReduceTailSum",
59  "",
60  vector<string>{GI(1) + "before_aggregate"},
61  vector<string>{GI(1)})};
62  }
63 };
64 REGISTER_GRADIENT(RowMul, GetRowMulGradient);
65 
66 } // namespace
67 
68 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.