saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-22 19:23:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Replace Where Grad "Mul" with "Where" (#7672)

Browse source 
* replace where grad mul with where

* clean up

* auto formatting

* remove not for second input
This commit is contained in:
harshithapv 2021-05-13 08:54:43 -07:00 committed by GitHub
parent b4e8e9b004
commit 31ca21b782
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
orttraining/orttraining/core/graph/gradient_builder.cc
View file

@ -1347,16 +1347,15 @@ IMPLEMENT_GRADIENT_BUILDER(GetSliceGradient) {
IMPLEMENT_GRADIENT_BUILDER(GetWhereGradient) {
std::vector<NodeDef> result;
const int64_t data_type = static_cast<int64_t>(IElemType(1));
NodeDef zero_constant_node = ZeroConstantNode(OElemType(0));
ArgDef ZERO = zero_constant_node.output_args[0];
result.push_back(zero_constant_node);
if (IsGradientRequiredForSrcNodeInput(1)) {
result.push_back(NodeDef("Cast", {I(0)}, {IA("Positive_Mask")}, {MakeAttribute("to", data_type)}));
result.push_back(NodeDef("Mul", {GO(0), IA("Positive_Mask")}, {GI(1)}));
result.push_back(NodeDef("Where", {I(0), GO(0), ZERO}, {GI(1)}));
}
if (IsGradientRequiredForSrcNodeInput(2)) {
result.push_back(NodeDef("Not", {I(0)}, {IA("Not_Condition", IType(0))}));
result.push_back(NodeDef("Cast", {IA("Not_Condition")}, {IA("Negative_Mask")}, {MakeAttribute("to", data_type)}));
result.push_back(NodeDef("Mul", {GO(0), IA("Negative_Mask")}, {GI(2)}));
result.push_back(NodeDef("Where", {I(0), ZERO, GO(0)}, {GI(2)}));
}
return result;
}