From 31ca21b78268ff7c71f409fd53830a5e0cb210e1 Mon Sep 17 00:00:00 2001 From: harshithapv <54084812+harshithapv@users.noreply.github.com> Date: Thu, 13 May 2021 08:54:43 -0700 Subject: [PATCH] Replace Where Grad "Mul" with "Where" (#7672) * replace where grad mul with where * clean up * auto formatting * remove not for second input --- .../orttraining/core/graph/gradient_builder.cc | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 38c4ad4933..39603ab78d 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -1347,16 +1347,15 @@ IMPLEMENT_GRADIENT_BUILDER(GetSliceGradient) { IMPLEMENT_GRADIENT_BUILDER(GetWhereGradient) { std::vector result; - const int64_t data_type = static_cast(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; }