From 8fd085801a39d299ff24962506fff82bc64fc001 Mon Sep 17 00:00:00 2001
From: Derek Murray <Derek.Murray@microsoft.com>
Date: Wed, 16 Dec 2020 08:32:10 -0800
Subject: [PATCH] Add gradient registration for Abs. (#6139)

---
 orttraining/orttraining/core/graph/gradient_builder.cc     | 7 +++++++
 orttraining/orttraining/core/graph/gradient_builder.h      | 1 +
 .../orttraining/core/graph/gradient_builder_registry.cc    | 1 +
 orttraining/orttraining/test/gradient/gradient_ops_test.cc | 4 ++++
 4 files changed, 13 insertions(+)

diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc
index 57193591bf..b5a31b17fd 100644
--- 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
index 6bc9207cf0..210edce12a 100644
--- 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
index 335dc4744f..da15486679 100644
--- 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
index 63be72b6cc..80e6ceae32 100644
--- 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});