Add gradient registration for Abs. (#6139)

microsoft · Dec 16, 2020 · 279faf1 · 279faf1
1 parent b8c8fe9
commit 279faf1
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 0 deletions.
diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc
@@ -1506,5 +1506,12 @@ IMPLEMENT_GRADIENT_BUILDER(GetClipGradient) {
   return output;
 }
 
+IMPLEMENT_GRADIENT_BUILDER(GetAbsGradient) {
+  return std::vector<NodeDef>{
+      NodeDef("Sign", {I(0)}, {IA("Sign_Input")}),
+      NodeDef("Mul", {GO(0), IA("Sign_Input")}, {GI(0)})
+  };
+}
+
 }  // namespace training
 }  // namespace onnxruntime
diff --git a/orttraining/orttraining/core/graph/gradient_builder.h b/orttraining/orttraining/core/graph/gradient_builder.h
@@ -70,6 +70,7 @@ DECLARE_GRADIENT_BUILDER(GetExpGradient)
 DECLARE_GRADIENT_BUILDER(GetFlattenGradient)
 DECLARE_GRADIENT_BUILDER(GetTopKGradient)
 DECLARE_GRADIENT_BUILDER(GetClipGradient)
+DECLARE_GRADIENT_BUILDER(GetAbsGradient)
 
 }  // namespace training
 }  // namespace onnxruntime
diff --git a/orttraining/orttraining/core/graph/gradient_builder_registry.cc b/orttraining/orttraining/core/graph/gradient_builder_registry.cc
@@ -101,6 +101,7 @@ void GradientBuilderRegistry::RegisterGradientBuilders() {
   REGISTER_GRADIENT_BUILDER("Flatten", GetFlattenGradient);
   REGISTER_GRADIENT_BUILDER("TopK", GetTopKGradient);
   REGISTER_GRADIENT_BUILDER("Clip", GetClipGradient);
+  REGISTER_GRADIENT_BUILDER("Abs", GetAbsGradient);
 };
 
 }  // namespace training

diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc
@@ -427,6 +427,10 @@ TEST(GradientCheckerTest, NegGrad) {
   UnaryOpGradientTest("Neg");
 }
 
+TEST(GradientCheckerTest, AbsGrad) {
+  UnaryOpGradientTest("Abs");
+}
+
 TEST(GradientCheckerTest, LogGrad) {
   TensorShape shape({2, 5, 6});