frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
501d0729cb7cfd25aebb2c0e82cef035bf4e7e78
pytorch/caffe2/operators/elementwise_ops_schema.cc
Nikita Shulga a9b0a921d5 Disable avoid-non-const-global-variables lint check (#62008) 
Summary:
As GoogleTest `TEST` macro is non-compliant with it as well as `DEFINE_DISPATCH`

All changes but the ones to `.clang-tidy` are generated using following script:
```
for i in `find . -type f -iname "*.c*" -or -iname "*.h"|xargs grep cppcoreguidelines-avoid-non-const-global-variables|cut -f1 -d:|sort|uniq`;  do sed -i "/\/\/ NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)/d" $i; done
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/62008

Reviewed By: driazati, r-barnes

Differential Revision: D29838584

Pulled By: malfet

fbshipit-source-id: 1b2f8602c945bd4ce50a9bfdd204755556e31d13
2021-07-22 18:04:40 -07:00

1047 lines
24 KiB
C++
Raw Blame History

#include "caffe2/operators/elementwise_ops.h"
#include "caffe2/core/operator_gradient.h"
#include "caffe2/utils/proto_utils.h"
namespace caffe2 {
namespace {
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
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.
Github Links:
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/elementwise_ops_schema.cc
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kAddExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Add",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([[1,2],[3,4]]))
workspace.FeedBlob("B", np.array([[5,6],[7,8]]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A:
[[1 2]
[3 4]]
B:
[[5 6]
[7 8]]
C:
[[ 6 8]
[10 12]]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kSubExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Sub",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([[10,12],[4,14]]))
workspace.FeedBlob("B", np.array([[5,16],[1,19]]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A:
[[10 12]
[ 4 14]]
B:
[[ 5 16]
[ 1 19]]
C:
[[ 5 -4]
[ 3 -5]]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kMulExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Mul",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([[1,2],[3,4]]))
workspace.FeedBlob("B", np.array([[5,6],[7,8]]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A:
[[1 2]
[3 4]]
B:
[[5 6]
[7 8]]
C:
[[ 5 12]
[21 32]]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kDivExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Div",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([[18,8],[2,9]]))
workspace.FeedBlob("B", np.array([[9,2],[3,2]]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A:
[[18 8]
[ 2 9]]
B:
[[9 2]
[3 2]]
C:
[[2 4]
[0 4]]
```
</details>
)DOC";
std::function<void(OpSchema&)> MathDocGenerator(
const char* name,
const char* extra) {
return [=](OpSchema& schema) {
string doc = R"DOC(
Performs element-wise binary {name} (with limited broadcast support).
{broadcast_doc}
{extra}
)DOC";
c10::ReplaceAll(doc, "{name}", name);
c10::ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
c10::ReplaceAll(doc, "{extra}", extra);
schema.SetDoc(doc);
schema.Arg(
"broadcast", "*(type: int; default: 0)* Pass 1 to enable broadcasting");
schema.Arg("axis", "*(type: int; default: -1)* Axis to concatenate on.");
schema.Input(
0,
"A",
"*(type: Tensor`<float>`)* First operand, should share the type with the second operand.");
schema.Input(
1,
"B",
"*(type: Tensor`<float>`)* Second operand. With broadcasting can be of smaller size than A. "
"If broadcasting is disabled it should be of the same size as A.");
schema.Output(
0,
"C",
"*(type: Tensor`<float>`)* Output tensor with same dimensions and type as A.");
};
}
std::vector<TensorShape> ElementwiseOpShapeInference(
const OperatorDef& def,
const std::vector<TensorShape>& in) {
std::vector<TensorShape> out(1);
out[0].set_data_type(in[0].data_type());
ArgumentHelper helper(def);
const bool broadcast = helper.GetSingleArgument<bool>("broadcast", false);
if (broadcast) {
out[0].mutable_dims()->CopyFrom(in[0].dims());
} else {
const std::vector<int> A_dims(in[0].dims().begin(), in[0].dims().end());
const std::vector<int> B_dims(in[1].dims().begin(), in[1].dims().end());
const std::vector<int> C_dims =
elementwise_ops_utils::ComputeBinaryBroadcastForwardDims(
A_dims, B_dims);
for (const int dim : C_dims) {
out[0].add_dims(dim);
}
}
return out;
}
std::vector<TensorShape> ElementwiseGradientOpShapeInference(
const OperatorDef& def,
const std::vector<TensorShape>& in) {
std::vector<TensorShape> out;
out.push_back(in.at(1));
out.push_back(in.at(2));
return out;
}
} // namespace
OPERATOR_SCHEMA(Add)
.NumInputs(2)
.NumOutputs(1)
.AllowInplace({{0, 0}, {1, 0}})
.CostInferenceFunction(PointwiseCostInference<1>)
.TensorInferenceFunction(ElementwiseOpShapeInference)
.FillUsing(MathDocGenerator("addition", kAddExample))
.InheritOnnxSchema();
OPERATOR_SCHEMA(AddGradient)
.NumInputs(3)
.NumOutputs(2)
.TensorInferenceFunction(ElementwiseGradientOpShapeInference)
.AllowInplace({{0, 0}, {0, 1}});
OPERATOR_SCHEMA(Sub)
.NumInputs(2)
.NumOutputs(1)
.AllowInplace({{0, 0}, {1, 0}})
.CostInferenceFunction(PointwiseCostInference<1>)
.TensorInferenceFunction(ElementwiseOpShapeInference)
.FillUsing(MathDocGenerator("subtraction", kSubExample))
.InheritOnnxSchema();
OPERATOR_SCHEMA(SubGradient)
.NumInputs(3)
.NumOutputs(2)
.TensorInferenceFunction(ElementwiseGradientOpShapeInference)
.AllowInplace({{0, 0}, {0, 1}});
OPERATOR_SCHEMA(Mul)
.NumInputs(2)
.NumOutputs(1)
.AllowInplace({{0, 0}, {1, 0}})
.CostInferenceFunction(PointwiseCostInference<1>)
.TensorInferenceFunction(ElementwiseOpShapeInference)
.FillUsing(MathDocGenerator("multiplication", kMulExample))
.InheritOnnxSchema();
OPERATOR_SCHEMA(MulGradient)
.NumInputs(3)
.NumOutputs(2)
.TensorInferenceFunction(ElementwiseGradientOpShapeInference)
.AllowInplace({{0, 0}, {0, 1}})
.CostInferenceFunction(PointwiseCostInference<2>);
OPERATOR_SCHEMA(Div)
.NumInputs(2)
.NumOutputs(1)
.AllowInplace({{0, 0}})
.CostInferenceFunction(PointwiseCostInference<1>)
.TensorInferenceFunction(ElementwiseOpShapeInference)
.FillUsing(MathDocGenerator("division", kDivExample))
.InheritOnnxSchema();
OPERATOR_SCHEMA(DivGradient)
.NumInputs(3, 4)
.NumOutputs(2)
.TensorInferenceFunction(ElementwiseGradientOpShapeInference)
.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");
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kLTExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"LT",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([1, 5, 2, 9, 12, 3]))
workspace.FeedBlob("B", np.array([1, 3, 4, 9, 12, 8]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A: [ 1 5 2 9 12 3]
B: [ 1 3 4 9 12 8]
C: [False False True False False True]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kLEExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"LE",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([1, 5, 2, 9, 12, 3]))
workspace.FeedBlob("B", np.array([1, 3, 4, 9, 12, 8]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A: [ 1 5 2 9 12 3]
B: [ 1 3 4 9 12 8]
C: [ True False True True True True]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kGTExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"GT",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([1, 5, 2, 9, 12, 3]))
workspace.FeedBlob("B", np.array([1, 3, 4, 9, 12, 8]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A: [ 1 5 2 9 12 3]
B: [ 1 3 4 9 12 8]
C: [False True False False False False]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kGEExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"GE",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([1, 5, 2, 9, 12, 3]))
workspace.FeedBlob("B", np.array([1, 3, 4, 9, 12, 8]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A: [ 1 5 2 9 12 3]
B: [ 1 3 4 9 12 8]
C: [ True True False True True False]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kEQExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"EQ",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([1, 5, 2, 9, 12, 3]))
workspace.FeedBlob("B", np.array([1, 3, 4, 9, 12, 8]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A: [ 1 5 2 9 12 3]
B: [ 1 3 4 9 12 8]
C: [ True False False True True False]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kNEExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"NE",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", np.array([1, 5, 2, 9, 12, 3]))
workspace.FeedBlob("B", np.array([1, 3, 4, 9, 12, 8]))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A: [ 1 5 2 9 12 3]
B: [ 1 3 4 9 12 8]
C: [False True True False False True]
```
</details>
)DOC";
std::function<void(OpSchema&)>
ComparisonDocGenerator(const char* name, const char* desc, const char* extra) {
return [=](OpSchema& schema) {
string doc = R"DOC(
Performs element-wise {desc} comparison **{name}** (with limited broadcast support).
{broadcast_doc}
{extra}
)DOC";
c10::ReplaceAll(doc, "{name}", name);
c10::ReplaceAll(doc, "{desc}", desc);
c10::ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
c10::ReplaceAll(doc, "{extra}", extra);
schema.SetDoc(doc);
schema.Arg(
"broadcast",
"*(type: int; default: 0)* Pass 1 to enable broadcasting.");
schema.Arg(
"axis",
"*(type: int; default: -1)* Axis to concatenate on. If set, defines the broadcast dimensions.");
schema.Input(
0,
"A",
"*(type: Tensor`<bool>`)* First operand, should share the type with the second operand.");
schema.Input(
1,
"B",
"*(type: Tensor`<bool>`)* 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",
"*(type: Tensor`<bool>`)* Output tensor with same dimensions as `A`.");
};
}
#define CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(name, symbol, desc, extra) \
OPERATOR_SCHEMA(name) \
.NumInputs(2) \
.NumOutputs(1) \
.TensorInferenceFunction([](const OperatorDef& def, \
const vector<TensorShape>& in) { \
ArgumentHelper helper(def); \
const auto broadcasted = \
helper.GetSingleArgument<bool>("broadcast", false); \
if (!broadcasted) { \
CAFFE_ENFORCE_EQ(in[0].dims().size(), in[1].dims().size()); \
for (int i = 0; i < in[0].dims().size(); ++i) { \
CAFFE_ENFORCE_EQ(in[0].dims(i), in[1].dims(i)); \
} \
} \
auto output_dims = \
std::vector<int64_t>(in[0].dims().begin(), in[0].dims().end()); \
return vector<TensorShape>{ \
CreateTensorShape(output_dims, TensorProto::BOOL)}; \
}) \
.FillUsing(ComparisonDocGenerator(symbol, desc, extra)); \
SHOULD_NOT_DO_GRADIENT(name)
CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(EQ, "==", "equal to", kEQExample);
CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(NE, "!=", "not equal to", kNEExample);
CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(LT, "<", "less than", kLTExample);
CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(
LE,
"<=",
"less or equal than",
kLEExample);
CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(GT, ">", "greater than", kGTExample);
CAFFE2_SCHEMA_FOR_BINARY_COMPARISON_OP(
GE,
">=",
"greater or equal than",
kGEExample);
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kAndExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"And",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", (np.random.rand(3, 3) > 0.5))
workspace.FeedBlob("B", (np.random.rand(3, 3) > 0.5))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A:
[[ True False False]
[False True False]
[False False True]]
B:
[[ True False True]
[False False False]
[False False False]]
C:
[[ True False False]
[False False False]
[False False False]]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kOrExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Or",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", (np.random.rand(3, 3) > 0.5))
workspace.FeedBlob("B", (np.random.rand(3, 3) > 0.5))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A:
[[False True True]
[False True True]
[ True True True]]
B:
[[False True False]
[ True True True]
[False True False]]
C:
[[False True True]
[ True True True]
[ True True True]]
```
</details>
)DOC";
// NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
const char kXorExample[] = R"DOC(
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Xor",
["A", "B"],
["C"],
)
workspace.FeedBlob("A", (np.random.rand(3, 3) > 0.5))
workspace.FeedBlob("B", (np.random.rand(3, 3) > 0.5))
print("A:", workspace.FetchBlob("A"))
print("B:", workspace.FetchBlob("B"))
workspace.RunOperatorOnce(op)
print("C:", workspace.FetchBlob("C"))
```
**Result**
```
A:
[[ True True True]
[False False True]
[False True False]]
B:
[[False False False]
[ True True True]
[False False False]]
C:
[[ True True True]
[ True True False]
[False True False]]
```
</details>
)DOC";
std::function<void(OpSchema&)> LogicalDocGenerator(
const char* name,
const char* extra) {
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}
{extra}
)DOC";
c10::ReplaceAll(doc, "{name}", name);
c10::ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
c10::ReplaceAll(doc, "{extra}", extra);
schema.SetDoc(doc);
schema.Arg(
"broadcast",
"*(type: int; default: 0)* Pass 1 to enable broadcasting.");
schema.Arg(
"axis",
"*(type: int; default: -1)* Axis to concatenate on. If set, defines the broadcast dimensions.");
schema.Input(0, "A", "*(type: Tensor`<bool>`)* First operand.");
schema.Input(
1,
"B",
"*(type: Tensor`<bool>`)* 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",
"*(type: Tensor`<bool>`)* Output tensor of booleans. Has same dimensions as input `A`.");
};
}
#define CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(name, symbol, onnx_schema, extra) \
OPERATOR_SCHEMA(name) \
.NumInputs(2) \
.NumOutputs(1) \
.AllowInplace({{0, 0}}) \
.FillUsing(LogicalDocGenerator(symbol, extra)) \
.TensorInferenceFunction(ElementwiseOpShapeInference) \
.InheritOnnxSchema(onnx_schema); \
SHOULD_NOT_DO_GRADIENT(name)
CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(Or, "or", "Or", kOrExample);
CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(And, "and", "And", kAndExample);
CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP(Xor, "xor", "Xor", kXorExample);
#undef CAFFE2_SCHEMA_FOR_BINARY_LOGICAL_OP
std::function<void(OpSchema&)> BitwiseDocGenerator(const char* name) {
return [=](OpSchema& schema) {
string doc = R"DOC(
Performs element-wise bitwise operation `{name}` (with limited broadcast support).
Both input operands should be of type `bool`.
{broadcast_doc})DOC";
c10::ReplaceAll(doc, "{name}", name);
c10::ReplaceAll(doc, "{broadcast_doc}", kBroadcastDoc);
schema.SetDoc(doc);
schema.Arg(
"broadcast",
"*(type: int; default: 0)* Pass 1 to enable broadcasting.");
schema.Arg(
"axis",
"*(type: int; default: -1)* Axis to concatenate on. If set, defines the broadcast dimensions.");
schema.Input(0, "A", "*(type: Tensor)* First operand.");
schema.Input(
1,
"B",
"*(type: Tensor)* 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",
"*(type: Tensor)* Output tensor. Has same dimensions as input `A`.");
};
}
#define CAFFE2_SCHEMA_FOR_BINARY_BITWISE_OP(name, symbol) \
OPERATOR_SCHEMA(name) \
.NumInputs(2) \
.NumOutputs(1) \
.AllowInplace({{0, 0}}) \
.FillUsing(BitwiseDocGenerator(symbol)) \
.TensorInferenceFunction(ElementwiseOpShapeInference); \
SHOULD_NOT_DO_GRADIENT(name)
CAFFE2_SCHEMA_FOR_BINARY_BITWISE_OP(BitwiseOr, "bitwise_or");
CAFFE2_SCHEMA_FOR_BINARY_BITWISE_OP(BitwiseAnd, "bitwise_and");
CAFFE2_SCHEMA_FOR_BINARY_BITWISE_OP(BitwiseXor, "bitwise_xor");
#undef CAFFE2_SCHEMA_FOR_BINARY_BITWISE_OP
OPERATOR_SCHEMA(Not)
.NumInputs(1)
.NumOutputs(1)
.IdenticalTypeAndShapeOfInput(0)
.SetDoc(R"DOC(
Performs element-wise negation on input tensor `X`.
Github Links:
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/elementwise_ops_schema.cc
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Not",
["X"],
["Y"],
)
workspace.FeedBlob("X", (np.random.rand(3, 3) > 0.5))
print("X:", workspace.FetchBlob("X"))
workspace.RunOperatorOnce(op)
print("Y:", workspace.FetchBlob("Y"))
```
**Result**
```
X:
[[ True False False]
[False False False]
[ True True True]]
Y:
[[False True True]
[ True True True]
[False False False]]
```
</details>
)DOC")
.Input(0, "X", "*(Tensor`<bool>`)* Input tensor.")
.Output(0, "Y", "*(Tensor`<bool>`)* Negated output tensor.")
.InheritOnnxSchema();
SHOULD_NOT_DO_GRADIENT(Not);
OPERATOR_SCHEMA(Sign)
.NumInputs(1)
.NumOutputs(1)
.IdenticalTypeAndShapeOfInput(0)
.SetDoc(R"DOC(
Computes sign for each element of the input: -1, 0 or 1.
Github Link:
- https://github.com/pytorch/pytorch/blob/master/caffe2/operators/elementwise_ops_schema.cc
<details>
<summary> <b>Example</b> </summary>
**Code**
```
workspace.ResetWorkspace()
op = core.CreateOperator(
"Sign",
["X"],
["Y"],
)
workspace.FeedBlob("X", (np.random.rand(3, 3).astype(np.float32) - np.random.rand(3, 3).astype(np.float32)))
print("X:", workspace.FetchBlob("X"))
workspace.RunOperatorOnce(op)
print("Y:", workspace.FetchBlob("Y"))
```
**Result**
```
X:
[[ 0.02816287 0.22408086 -0.30342305]
[-0.18481976 0.03948995 0.39698976]
[-0.63304734 -0.6919183 -0.31524038]]
Y:
[[ 1. 1. -1.]
[-1. 1. 1.]
[-1. -1. -1.]]
```
</details>
)DOC")
.Input(0, "X", "*(type: Tensor`<float>`)* Input data tensor.")
.Output(0, "Y", "*(type: Tensor`<float>`)* Output tensor.");
SHOULD_NOT_DO_GRADIENT(Sign);
} // namespace caffe2
Reference in New Issue View Git Blame Copy Permalink