Caffe2 - C++ API
A deep learning, cross platform ML framework
math_ops.cc
1 
17 #include "caffe2/operators/math_ops.h"
18 #include "caffe2/utils/math.h"
19 
20 namespace caffe2 {
21 
22 struct SqrCPUFunctor {
23  template <typename T>
24  inline void
25  operator()(const int n, const T* x, T* y, CPUContext* device_context) {
26  math::Sqr<T, CPUContext>(n, x, y, device_context);
27  }
28 };
29 
30 REGISTER_CPU_OPERATOR(
31  Sqr,
33 
34 OPERATOR_SCHEMA(Sqr)
35  .NumInputs(1)
36  .NumOutputs(1)
37  .AllowInplace({{0, 0}})
38  .IdenticalTypeAndShape()
39  .SetDoc("Square (x^2) the elements of the input")
40  .Input(0, "input", "Input tensor")
41  .Output(0, "output", "Squared elements of the input");
42 
44  using GradientMakerBase::GradientMakerBase;
45  vector<OperatorDef> GetGradientDefs() override {
46  Argument scale_arg;
47  scale_arg.set_name("scale");
48  scale_arg.set_f(2.0);
49  return vector<OperatorDef>{CreateOperatorDef(
50  "Scale",
51  "",
52  std::vector<string>{GO(0)},
53  std::vector<string>{GO(0)},
54  std::vector<Argument>{scale_arg}),
55  CreateOperatorDef(
56  "Mul",
57  "",
58  std::vector<string>{GO(0), I(0)},
59  std::vector<string>{GI(0)})};
60  }
61 };
62 REGISTER_GRADIENT(Sqr, GetSqrGradient);
63 
65  template <typename T>
66  inline void
67  operator()(const int n, const T* x, T* y, CPUContext* device_context) {
68  for (int i = 0; i < n; ++i) {
69  y[i] = (-T(1) * (x[i] < 0)) + (x[i] > 0);
70  }
71  }
72 };
73 
74 REGISTER_CPU_OPERATOR(
75  Sign,
77 
78 OPERATOR_SCHEMA(Sign)
79  .NumInputs(1)
80  .NumOutputs(1)
81  .SetDoc("Computes sign for each element of the input: -1, 0 or 1.")
82  .IdenticalTypeAndShape();
83 SHOULD_NOT_DO_GRADIENT(Sign);
84 
85 } // namespace caffe2
The CPU Context, representing the bare minimum of what a Context class in Caffe2 should implement...
Definition: context.h:82
Copyright (c) 2016-present, Facebook, Inc.