caffe2/operators/elementwise_op_schema.cc - platform/external/pytorch - Git at Google

 #include "caffe2/core/operator_gradient.h"
 #include "caffe2/operators/elementwise_op.h"
 #include "caffe2/utils/proto_utils.h"

 namespace caffe2 {

 const char* kBroadcastDoc = R"DOC(
 If necessary the right-hand-side argument will be broadcasted to match the
 shape of left-hand-side argument. When broadcasting is specified, the second
 tensor can either be of size 1 (a scalar value), or having its shape as a
 contiguous subset of the first tensor's shape. The starting of the mutually
 equal shape is specified by the argument "axis", and if it is not set, suffix
 matching is assumed. 1-dim expansion doesn't work yet.

 For example, the following tensor shapes are supported (with broadcast=1):

   shape(A) = (2, 3, 4, 5), shape(B) = (,), i.e. B is a scalar
   shape(A) = (2, 3, 4, 5), shape(B) = (5,)
   shape(A) = (2, 3, 4, 5), shape(B) = (4, 5)
   shape(A) = (2, 3, 4, 5), shape(B) = (3, 4), with axis=1
   shape(A) = (2, 3, 4, 5), shape(B) = (2), with axis=0

 Argument `broadcast=1` needs to be passed to enable broadcasting.
 )DOC";

 std::function<void(OpSchema&)> MathDocGenerator(const char* name) {
   return [=](OpSchema& schema) {
     string doc = R"DOC(
 Performs element-wise binary {name} (with limited broadcast support).
 {broadcast_doc})DOC";
     ReplaceAll(doc, "{name}", name);
     ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
     schema.SetDoc(doc);
     schema.Arg("broadcast", "Pass 1 to enable broadcasting");
     schema.Arg(
         "axis",
         "If set, defines the broadcast dimensions. See doc for details.");
     schema.Input(
         0,
         "A",
         "First operand, should share the type with the second operand.");
     schema.Input(
         1,
         "B",
         "Second operand. With broadcasting can be of smaller size than A. "
         "If broadcasting is disabled it should be of the same size.");
     schema.Output(0, "C", "Result, has same dimensions and type as A");
   };
 }

 OPERATOR_SCHEMA(Add)
     .NumInputs(2)
     .NumOutputs(1)
     .AllowInplace({{0, 0}, {1, 0}})
     .CostInferenceFunction(PointwiseCostInference<1>)
     .IdenticalTypeAndShapeOfInput(0)
     .FillUsing(MathDocGenerator("addition"))
     .InheritOnnxSchema("Add");
 OPERATOR_SCHEMA(Sub)
     .NumInputs(2)
     .NumOutputs(1)
     .AllowInplace({{0, 0}, {1, 0}})
     .CostInferenceFunction(PointwiseCostInference<1>)
     .IdenticalTypeAndShapeOfInput(0)
     .FillUsing(MathDocGenerator("subtraction"))
     .InheritOnnxSchema("Sub");
 OPERATOR_SCHEMA(Mul)
     .NumInputs(2)
     .NumOutputs(1)
     .AllowInplace({{0, 0}, {1, 0}})
     .CostInferenceFunction(PointwiseCostInference<1>)
     .IdenticalTypeAndShapeOfInput(0)
     .FillUsing(MathDocGenerator("multiplication"))
     .InheritOnnxSchema("Mul");
 OPERATOR_SCHEMA(Div)
     .NumInputs(2)
     .NumOutputs(1)
     .AllowInplace({{0, 0}})
     .CostInferenceFunction(PointwiseCostInference<1>)
     .IdenticalTypeAndShapeOfInput(0)
     .FillUsing(MathDocGenerator("division"))
     .InheritOnnxSchema("Div");
 OPERATOR_SCHEMA(DivGradient).NumInputs(3).NumOutputs(2).AllowInplace({{0, 0}});

 OPERATOR_SCHEMA(SumReduceLike)
     .NumInputs(2)
     .NumOutputs(1)
     .IdenticalTypeAndShapeOfInput(0)
     .SetDoc(R"DOC(
 SumReduceLike operator takes 2 tensors as input. It performs reduce sum to the
 first input so that the output looks like the second one.
 It assumes that the first input
 has more dimensions than the second, and the dimensions of the second input is
 the contiguous subset of the dimensions of the first.
 For example, the following tensor shapes are supported:

   shape(A) = (2, 3, 4, 5), shape(B) = (4, 5)
   shape(A) = (2, 3, 4, 5), shape(B) = (,), i.e. B is a scalar
   shape(A) = (2, 3, 4, 5), shape(B) = (3, 4), with axis=1
   shape(A) = (2, 3, 2, 5), shape(B) = (2), with axis=0
     )DOC")
     .Arg(
         "axis",
         "If set, defines the starting dimension for reduction. Args `axis` and "
         "`axis_str` cannot be used simultaneously.")
     .Arg(
         "axis_str",
         "If set, it could only be N or C or H or W. `order` arg should also be "
         "provided. It defines the reduction dimensions on NCHW or NHWC. Args "
         "`axis` and `axis_str` cannot be used simultaneously.")
     .Arg("order", "Either NHWC or HCWH")
     .Input(
         0,
         "A",
         "First operand, should share the type with the second operand.")
     .Input(
         1,
         "B",
         "Second operand. With broadcasting can be of smaller size than A. "
         "If broadcasting is disabled it should be of the same size.")
     .Output(0, "C", "Result, has same dimensions and type as B");

 class GetAddGradient : public GradientMakerBase {
   using GradientMakerBase::GradientMakerBase;
   vector<OperatorDef> GetGradientDefs() override {
     if (!ArgumentHelper::HasArgument(Def(), "broadcast")) {
       SetDense(0, GO(0));
       SetDense(1, GO(0));
       return vector<OperatorDef>();
     }
     SetDense(0, GO(0));

     return SingleGradientDef(
         "SumReduceLike",
         "",
         vector<string>{GO(0), I(1)},
         vector<string>{GI(1)});
   }
 };
 REGISTER_GRADIENT(Add, GetAddGradient);

 // TODO(jiayq): Although we have Sub gradient implemented, we are still missing
 // the Negative unary operator to be implemented.
 class GetSubGradient : public GradientMakerBase {
   using GradientMakerBase::GradientMakerBase;
   vector<OperatorDef> GetGradientDefs() override {
     if (!ArgumentHelper::HasArgument(Def(), "broadcast")) {
       SetDense(0, GO(0));
       return SingleGradientDef(
           "Negative", "", vector<string>{GO(0)}, vector<string>{GI(1)});
     } else {
       SetDense(0, GO(0));
       vector<OperatorDef> grad_ops;
       grad_ops.push_back(CreateOperatorDef(
           "Negative",
           "",
           vector<string>{GO(0)},
           vector<string>{GI(1) + "_autogen_pre_red"}));

       Argument axis, axis_str, order;
       if (ArgumentHelper::HasArgument(Def(), "axis")) {
         axis = GetArgument(Def(), "axis");
       } else {
         axis = MakeArgument<int>("axis", -1);
       }
       if (ArgumentHelper::HasArgument(Def(), "axis_str")) {
         axis_str = GetArgument(Def(), "axis_str");
       } else {
         axis_str = MakeArgument<string>("axis_str", "");
       }
       if (ArgumentHelper::HasArgument(Def(), "order")) {
         order = GetArgument(Def(), "order");
       } else {
         order = MakeArgument<string>("order", "NCHW");
       }
       grad_ops.push_back(CreateOperatorDef(
           "SumReduceLike",
           "",
           vector<string>{GI(1) + "_autogen_pre_red", I(1)},
           vector<string>{GI(1)},
           vector<Argument>{axis, axis_str, order}));

       return grad_ops;
     }
   }
   // Make sure the broadcast argument is not copied over.
   bool CopyArguments() const override {
     return false;
   }
 };
 REGISTER_GRADIENT(Sub, GetSubGradient);

 class GetMulGradient : public GradientMakerBase {
   using GradientMakerBase::GradientMakerBase;
   vector<OperatorDef> GetGradientDefs() override {
     CAFFE_ENFORCE(
         Def().input(0) != Def().output(0) && Def().input(1) != Def().output(0),
         "Gradient computation cannot be carried out if Mul uses in-place "
         "computation: ",
         ProtoDebugString(Def()));
     if (!ArgumentHelper::HasArgument(Def(), "broadcast")) {
       return vector<OperatorDef>{
           CreateOperatorDef(
               "Mul", "", vector<string>{GO(0), I(1)}, vector<string>{GI(0)}),
           CreateOperatorDef(
               "Mul", "", vector<string>{GO(0), I(0)}, vector<string>{GI(1)})};
     } else {
       Argument broadcast, axis, axis_str, order;
       if (ArgumentHelper::HasArgument(Def(), "broadcast")) {
         broadcast = GetArgument(Def(), "broadcast");
       } else {
         broadcast = MakeArgument<int>("broadcast", 0);
       }
       if (ArgumentHelper::HasArgument(Def(), "axis")) {
         axis = GetArgument(Def(), "axis");
       } else {
         axis = MakeArgument<int>("axis", -1);
       }
       if (ArgumentHelper::HasArgument(Def(), "axis_str")) {
         axis_str = GetArgument(Def(), "axis_str");
       } else {
         axis_str = MakeArgument<string>("axis_str", "");
       }
       if (ArgumentHelper::HasArgument(Def(), "order")) {
         order = GetArgument(Def(), "order");
       } else {
         order = MakeArgument<string>("order", "NCHW");
       }

       vector<OperatorDef> grad_ops;
       grad_ops.push_back(CreateOperatorDef(
           "Mul",
           "mul_grad_1st_op",
           vector<string>{GO(0), I(1)},
           vector<string>{GI(0)},
           vector<Argument>{broadcast, axis, axis_str, order}));
       grad_ops.push_back(CreateOperatorDef(
           "Mul",
           "mul_gradient_2nd_op",
           vector<string>{GO(0), I(0)},
           vector<string>{GI(1) + "_autogen_pre_red"}));

       grad_ops.push_back(CreateOperatorDef(
           "SumReduceLike",
           "mul_with_broadcast_grad_3",
           vector<string>{GI(1) + "_autogen_pre_red", I(1)},
           vector<string>{GI(1)},
           vector<Argument>{axis, axis_str, order}));

       return grad_ops;
     }
   }

   // Make sure the broadcast argument is not copied over.
   bool CopyArguments() const override {
     return false;
   }
 };
 REGISTER_GRADIENT(Mul, GetMulGradient);

 class GetDivGradient : public GradientMakerBase {
   using GradientMakerBase::GradientMakerBase;
   vector<OperatorDef> GetGradientDefs() override {
     CAFFE_ENFORCE(
         !ArgumentHelper::HasArgument(Def(), "broadcast"),
         "Gradient not ready yet for Div with broadcasting.");
     return SingleGradientDef(
         "DivGradient",
         "",
         vector<string>{I(1), O(0), GO(0)},
         vector<string>{GI(0), GI(1)});
   }
 };
 REGISTER_GRADIENT(Div, GetDivGradient);

 std::function<void(OpSchema&)> ComparisonDocGenerator(
     const char* name,
     const char* desc) {
   return [=](OpSchema& schema) {
     string doc = R"DOC(
 Performs element-wise {desc} comparison `{name}` (with limited broadcast support).
 {broadcast_doc})DOC";
     ReplaceAll(doc, "{name}", name);
     ReplaceAll(doc, "{desc}", desc);
     ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
     schema.SetDoc(doc);
     schema.Arg("broadcast", "Pass 1 to enable broadcasting");
     schema.Arg(
         "axis",
         "If set, defines the broadcast dimensions. See doc for details.");
     schema.Input(
         0,
         "A",
         "First operand, should share the type with the second operand.");
     schema.Input(
         1,
         "B",
         "Second operand. With broadcasting can be of smaller size than A. "
         "If broadcasting is disabled it should be of the same size.");
     schema.Output(0, "C", "Result, has same dimensions and A and type `bool`");
   };
 }

 #define CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(name, symbol, desc) \
   OPERATOR_SCHEMA(name).NumInputs(2).NumOutputs(1).FillUsing(      \
       ComparisonDocGenerator(symbol, desc));                       \
   SHOULD_NOT_DO_GRADIENT(name)

 CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(LT, "<", "less than");
 CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(LE, "<=", "less or equal than");
 CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(GT, ">", "greater than");
 CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(GE, ">=", "greater or equal than");
 CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(EQ, "==", "equality");

 std::function<void(OpSchema&)> LogicalDocGenerator(const char* name) {
   return [=](OpSchema& schema) {
     string doc = R"DOC(
 Performs element-wise logical operation `{name}` (with limited broadcast support).
 Both input operands should be of type `bool`.
 {broadcast_doc})DOC";
     ReplaceAll(doc, "{name}", name);
     ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
     schema.SetDoc(doc);
     schema.Arg("broadcast", "Pass 1 to enable broadcasting");
     schema.Arg(
         "axis",
         "If set, defines the broadcast dimensions. See doc for details.");
     schema.Input(0, "A", "First operand.");
     schema.Input(
         1,
         "B",
         "Second operand. With broadcasting can be of smaller size than A. "
         "If broadcasting is disabled it should be of the same size.");
     schema.Output(0, "C", "Result, has same dimensions and A and type `bool`");
   };
 }

 #define CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(name, symbol, onnx_schema) \
   OPERATOR_SCHEMA(name)                                   \
       .NumInputs(2)                                       \
       .NumOutputs(1)                                      \
       .AllowInplace({{0, 0}})                             \
       .FillUsing(LogicalDocGenerator(symbol))             \
       .InheritOnnxSchema(onnx_schema);                    \
   SHOULD_NOT_DO_GRADIENT(name)

 CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(Or, "or", "Or");
 CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(And, "and", "And");
 CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(Xor, "xor", "Xor");

 OPERATOR_SCHEMA(Not)
     .NumInputs(1)
     .NumOutputs(1)
     .SetDoc(R"DOC(Performs element-wise negation.)DOC")
     .Input(0, "X", "Input tensor of type `bool`.")
     .Output(0, "Y", "Output tensor of type `bool`.")
     .InheritOnnxSchema("Not");
 SHOULD_NOT_DO_GRADIENT(Not);

 } // namespace caffe2
	#include "caffe2/core/operator_gradient.h"
	#include "caffe2/operators/elementwise_op.h"
	#include "caffe2/utils/proto_utils.h"

	namespace caffe2 {

	const char* kBroadcastDoc = R"DOC(
	If necessary the right-hand-side argument will be broadcasted to match the
	shape of left-hand-side argument. When broadcasting is specified, the second
	tensor can either be of size 1 (a scalar value), or having its shape as a
	contiguous subset of the first tensor's shape. The starting of the mutually
	equal shape is specified by the argument "axis", and if it is not set, suffix
	matching is assumed. 1-dim expansion doesn't work yet.

	For example, the following tensor shapes are supported (with broadcast=1):

	shape(A) = (2, 3, 4, 5), shape(B) = (,), i.e. B is a scalar
	shape(A) = (2, 3, 4, 5), shape(B) = (5,)
	shape(A) = (2, 3, 4, 5), shape(B) = (4, 5)
	shape(A) = (2, 3, 4, 5), shape(B) = (3, 4), with axis=1
	shape(A) = (2, 3, 4, 5), shape(B) = (2), with axis=0

	Argument `broadcast=1` needs to be passed to enable broadcasting.
	)DOC";

	std::function<void(OpSchema&)> MathDocGenerator(const char* name) {
	return [=](OpSchema& schema) {
	string doc = R"DOC(
	Performs element-wise binary {name} (with limited broadcast support).
	{broadcast_doc})DOC";
	ReplaceAll(doc, "{name}", name);
	ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
	schema.SetDoc(doc);
	schema.Arg("broadcast", "Pass 1 to enable broadcasting");
	schema.Arg(
	"axis",
	"If set, defines the broadcast dimensions. See doc for details.");
	schema.Input(
	0,
	"A",
	"First operand, should share the type with the second operand.");
	schema.Input(
	1,
	"B",
	"Second operand. With broadcasting can be of smaller size than A. "
	"If broadcasting is disabled it should be of the same size.");
	schema.Output(0, "C", "Result, has same dimensions and type as A");
	};
	}

	OPERATOR_SCHEMA(Add)
	.NumInputs(2)
	.NumOutputs(1)
	.AllowInplace({{0, 0}, {1, 0}})
	.CostInferenceFunction(PointwiseCostInference<1>)
	.IdenticalTypeAndShapeOfInput(0)
	.FillUsing(MathDocGenerator("addition"))
	.InheritOnnxSchema("Add");
	OPERATOR_SCHEMA(Sub)
	.NumInputs(2)
	.NumOutputs(1)
	.AllowInplace({{0, 0}, {1, 0}})
	.CostInferenceFunction(PointwiseCostInference<1>)
	.IdenticalTypeAndShapeOfInput(0)
	.FillUsing(MathDocGenerator("subtraction"))
	.InheritOnnxSchema("Sub");
	OPERATOR_SCHEMA(Mul)
	.NumInputs(2)
	.NumOutputs(1)
	.AllowInplace({{0, 0}, {1, 0}})
	.CostInferenceFunction(PointwiseCostInference<1>)
	.IdenticalTypeAndShapeOfInput(0)
	.FillUsing(MathDocGenerator("multiplication"))
	.InheritOnnxSchema("Mul");
	OPERATOR_SCHEMA(Div)
	.NumInputs(2)
	.NumOutputs(1)
	.AllowInplace({{0, 0}})
	.CostInferenceFunction(PointwiseCostInference<1>)
	.IdenticalTypeAndShapeOfInput(0)
	.FillUsing(MathDocGenerator("division"))
	.InheritOnnxSchema("Div");
	OPERATOR_SCHEMA(DivGradient).NumInputs(3).NumOutputs(2).AllowInplace({{0, 0}});

	OPERATOR_SCHEMA(SumReduceLike)
	.NumInputs(2)
	.NumOutputs(1)
	.IdenticalTypeAndShapeOfInput(0)
	.SetDoc(R"DOC(
	SumReduceLike operator takes 2 tensors as input. It performs reduce sum to the
	first input so that the output looks like the second one.
	It assumes that the first input
	has more dimensions than the second, and the dimensions of the second input is
	the contiguous subset of the dimensions of the first.
	For example, the following tensor shapes are supported:

	shape(A) = (2, 3, 4, 5), shape(B) = (4, 5)
	shape(A) = (2, 3, 4, 5), shape(B) = (,), i.e. B is a scalar
	shape(A) = (2, 3, 4, 5), shape(B) = (3, 4), with axis=1
	shape(A) = (2, 3, 2, 5), shape(B) = (2), with axis=0
	)DOC")
	.Arg(
	"axis",
	"If set, defines the starting dimension for reduction. Args `axis` and "
	"`axis_str` cannot be used simultaneously.")
	.Arg(
	"axis_str",
	"If set, it could only be N or C or H or W. `order` arg should also be "
	"provided. It defines the reduction dimensions on NCHW or NHWC. Args "
	"`axis` and `axis_str` cannot be used simultaneously.")
	.Arg("order", "Either NHWC or HCWH")
	.Input(
	0,
	"A",
	"First operand, should share the type with the second operand.")
	.Input(
	1,
	"B",
	"Second operand. With broadcasting can be of smaller size than A. "
	"If broadcasting is disabled it should be of the same size.")
	.Output(0, "C", "Result, has same dimensions and type as B");

	class GetAddGradient : public GradientMakerBase {
	using GradientMakerBase::GradientMakerBase;
	vector<OperatorDef> GetGradientDefs() override {
	if (!ArgumentHelper::HasArgument(Def(), "broadcast")) {
	SetDense(0, GO(0));
	SetDense(1, GO(0));
	return vector<OperatorDef>();
	}
	SetDense(0, GO(0));

	return SingleGradientDef(
	"SumReduceLike",
	"",
	vector<string>{GO(0), I(1)},
	vector<string>{GI(1)});
	}
	};
	REGISTER_GRADIENT(Add, GetAddGradient);

	// TODO(jiayq): Although we have Sub gradient implemented, we are still missing
	// the Negative unary operator to be implemented.
	class GetSubGradient : public GradientMakerBase {
	using GradientMakerBase::GradientMakerBase;
	vector<OperatorDef> GetGradientDefs() override {
	if (!ArgumentHelper::HasArgument(Def(), "broadcast")) {
	SetDense(0, GO(0));
	return SingleGradientDef(
	"Negative", "", vector<string>{GO(0)}, vector<string>{GI(1)});
	} else {
	SetDense(0, GO(0));
	vector<OperatorDef> grad_ops;
	grad_ops.push_back(CreateOperatorDef(
	"Negative",
	"",
	vector<string>{GO(0)},
	vector<string>{GI(1) + "_autogen_pre_red"}));

	Argument axis, axis_str, order;
	if (ArgumentHelper::HasArgument(Def(), "axis")) {
	axis = GetArgument(Def(), "axis");
	} else {
	axis = MakeArgument<int>("axis", -1);
	}
	if (ArgumentHelper::HasArgument(Def(), "axis_str")) {
	axis_str = GetArgument(Def(), "axis_str");
	} else {
	axis_str = MakeArgument<string>("axis_str", "");
	}
	if (ArgumentHelper::HasArgument(Def(), "order")) {
	order = GetArgument(Def(), "order");
	} else {
	order = MakeArgument<string>("order", "NCHW");
	}
	grad_ops.push_back(CreateOperatorDef(
	"SumReduceLike",
	"",
	vector<string>{GI(1) + "_autogen_pre_red", I(1)},
	vector<string>{GI(1)},
	vector<Argument>{axis, axis_str, order}));

	return grad_ops;
	}
	}
	// Make sure the broadcast argument is not copied over.
	bool CopyArguments() const override {
	return false;
	}
	};
	REGISTER_GRADIENT(Sub, GetSubGradient);

	class GetMulGradient : public GradientMakerBase {
	using GradientMakerBase::GradientMakerBase;
	vector<OperatorDef> GetGradientDefs() override {
	CAFFE_ENFORCE(
	Def().input(0) != Def().output(0) && Def().input(1) != Def().output(0),
	"Gradient computation cannot be carried out if Mul uses in-place "
	"computation: ",
	ProtoDebugString(Def()));
	if (!ArgumentHelper::HasArgument(Def(), "broadcast")) {
	return vector<OperatorDef>{
	CreateOperatorDef(
	"Mul", "", vector<string>{GO(0), I(1)}, vector<string>{GI(0)}),
	CreateOperatorDef(
	"Mul", "", vector<string>{GO(0), I(0)}, vector<string>{GI(1)})};
	} else {
	Argument broadcast, axis, axis_str, order;
	if (ArgumentHelper::HasArgument(Def(), "broadcast")) {
	broadcast = GetArgument(Def(), "broadcast");
	} else {
	broadcast = MakeArgument<int>("broadcast", 0);
	}
	if (ArgumentHelper::HasArgument(Def(), "axis")) {
	axis = GetArgument(Def(), "axis");
	} else {
	axis = MakeArgument<int>("axis", -1);
	}
	if (ArgumentHelper::HasArgument(Def(), "axis_str")) {
	axis_str = GetArgument(Def(), "axis_str");
	} else {
	axis_str = MakeArgument<string>("axis_str", "");
	}
	if (ArgumentHelper::HasArgument(Def(), "order")) {
	order = GetArgument(Def(), "order");
	} else {
	order = MakeArgument<string>("order", "NCHW");
	}

	vector<OperatorDef> grad_ops;
	grad_ops.push_back(CreateOperatorDef(
	"Mul",
	"mul_grad_1st_op",
	vector<string>{GO(0), I(1)},
	vector<string>{GI(0)},
	vector<Argument>{broadcast, axis, axis_str, order}));
	grad_ops.push_back(CreateOperatorDef(
	"Mul",
	"mul_gradient_2nd_op",
	vector<string>{GO(0), I(0)},
	vector<string>{GI(1) + "_autogen_pre_red"}));

	grad_ops.push_back(CreateOperatorDef(
	"SumReduceLike",
	"mul_with_broadcast_grad_3",
	vector<string>{GI(1) + "_autogen_pre_red", I(1)},
	vector<string>{GI(1)},
	vector<Argument>{axis, axis_str, order}));

	return grad_ops;
	}
	}

	// Make sure the broadcast argument is not copied over.
	bool CopyArguments() const override {
	return false;
	}
	};
	REGISTER_GRADIENT(Mul, GetMulGradient);

	class GetDivGradient : public GradientMakerBase {
	using GradientMakerBase::GradientMakerBase;
	vector<OperatorDef> GetGradientDefs() override {
	CAFFE_ENFORCE(
	!ArgumentHelper::HasArgument(Def(), "broadcast"),
	"Gradient not ready yet for Div with broadcasting.");
	return SingleGradientDef(
	"DivGradient",
	"",
	vector<string>{I(1), O(0), GO(0)},
	vector<string>{GI(0), GI(1)});
	}
	};
	REGISTER_GRADIENT(Div, GetDivGradient);

	std::function<void(OpSchema&)> ComparisonDocGenerator(
	const char* name,
	const char* desc) {
	return [=](OpSchema& schema) {
	string doc = R"DOC(
	Performs element-wise {desc} comparison `{name}` (with limited broadcast support).
	{broadcast_doc})DOC";
	ReplaceAll(doc, "{name}", name);
	ReplaceAll(doc, "{desc}", desc);
	ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
	schema.SetDoc(doc);
	schema.Arg("broadcast", "Pass 1 to enable broadcasting");
	schema.Arg(
	"axis",
	"If set, defines the broadcast dimensions. See doc for details.");
	schema.Input(
	0,
	"A",
	"First operand, should share the type with the second operand.");
	schema.Input(
	1,
	"B",
	"Second operand. With broadcasting can be of smaller size than A. "
	"If broadcasting is disabled it should be of the same size.");
	schema.Output(0, "C", "Result, has same dimensions and A and type `bool`");
	};
	}

	#define CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(name, symbol, desc) \
	OPERATOR_SCHEMA(name).NumInputs(2).NumOutputs(1).FillUsing( \
	ComparisonDocGenerator(symbol, desc)); \
	SHOULD_NOT_DO_GRADIENT(name)

	CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(LT, "<", "less than");
	CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(LE, "<=", "less or equal than");
	CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(GT, ">", "greater than");
	CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(GE, ">=", "greater or equal than");
	CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(EQ, "==", "equality");

	std::function<void(OpSchema&)> LogicalDocGenerator(const char* name) {
	return [=](OpSchema& schema) {
	string doc = R"DOC(
	Performs element-wise logical operation `{name}` (with limited broadcast support).
	Both input operands should be of type `bool`.
	{broadcast_doc})DOC";
	ReplaceAll(doc, "{name}", name);
	ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
	schema.SetDoc(doc);
	schema.Arg("broadcast", "Pass 1 to enable broadcasting");
	schema.Arg(
	"axis",
	"If set, defines the broadcast dimensions. See doc for details.");
	schema.Input(0, "A", "First operand.");
	schema.Input(
	1,
	"B",
	"Second operand. With broadcasting can be of smaller size than A. "
	"If broadcasting is disabled it should be of the same size.");
	schema.Output(0, "C", "Result, has same dimensions and A and type `bool`");
	};
	}

	#define CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(name, symbol, onnx_schema) \
	OPERATOR_SCHEMA(name) \
	.NumInputs(2) \
	.NumOutputs(1) \
	.AllowInplace({{0, 0}}) \
	.FillUsing(LogicalDocGenerator(symbol)) \
	.InheritOnnxSchema(onnx_schema); \
	SHOULD_NOT_DO_GRADIENT(name)

	CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(Or, "or", "Or");
	CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(And, "and", "And");
	CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(Xor, "xor", "Xor");

	OPERATOR_SCHEMA(Not)
	.NumInputs(1)
	.NumOutputs(1)
	.SetDoc(R"DOC(Performs element-wise negation.)DOC")
	.Input(0, "X", "Input tensor of type `bool`.")
	.Output(0, "Y", "Output tensor of type `bool`.")
	.InheritOnnxSchema("Not");
	SHOULD_NOT_DO_GRADIENT(Not);

	} // namespace caffe2