1 #include "caffe2/operators/half_float_ops.h" 4 OPERATOR_SCHEMA(FloatToHalf)
7 .TensorInferenceFunction(
8 [](
const OperatorDef& def,
const vector<TensorShape>& in) {
9 vector<TensorShape> out;
10 const TensorShape& X = in[0];
12 out[0].set_data_type(TensorProto_DataType_FLOAT16);
17 OPERATOR_SCHEMA(HalfToFloat)
20 .TensorInferenceFunction(
21 [](
const OperatorDef& def,
const vector<TensorShape>& in) {
22 vector<TensorShape> out;
23 const TensorShape& X = in[0];
25 out[0].set_data_type(TensorProto_DataType_FLOAT);
29 OPERATOR_SCHEMA(Float16ConstantFill)
32 .TensorInferenceFunction(Float16FillerTensorInference)
33 .Arg(
"value",
"The value for the elements of the output tensor.")
34 .Arg(
"shape",
"The shape of the output tensor.")
38 "Output tensor of constant values specified by 'value'");
41 using GradientMakerBase::GradientMakerBase;
42 vector<OperatorDef> GetGradientDefs()
override {
44 "HalfToFloat",
"", vector<string>{GO(0)}, vector<string>{GI(0)});
50 using GradientMakerBase::GradientMakerBase;
51 vector<OperatorDef> GetGradientDefs()
override {
53 "FloatToHalf",
"", vector<string>{GO(0)}, vector<string>{GI(0)});
57 NO_GRADIENT(Float16ConstantFill);
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
static vector< OperatorDef > SingleGradientDef(const Args &...args)
a helper function to allow one to create one single operator def, which is usually the case for many ...