Correct atan2 gradient (#2066)

* Correct atan2 gradient

There was a minus sign offset

* Refactor HLODerivative for Atan2Op calculations

* Correct gradient for forward pass as well

* Fix sign in CheckedDiv calculation

* Update HLODerivatives.td

* Remove HLODerivative for Atan2Op

Removed HLODerivative definition for Atan2Op.

* fix

* fix

---------

Co-authored-by: William Moses <gh@wsmoses.com>

Author: ptiede

src/enzyme_ad/jax/Implementations/HLODerivatives.td

@@ -969,12 +969,12 @@ def : HLODerivative<"AddOp", (Op $x, $y),
                     (Add (Shadow $x), (Shadow $y))
                   >;
 
-def : HLODerivative<"Atan2Op", (Op $x, $y),
+def : HLODerivative<"Atan2Op", (Op $y, $x),
                     [
-                      (CheckedMul (DiffeRet), (Div (Neg $y), (Add (Pow $x, (HLOConstantFP<"2">)), (Pow $y, (HLOConstantFP<"2"> $y))))),
-                      (CheckedMul (DiffeRet), (Div $x, (Add (Pow $x, (HLOConstantFP<"2">)), (Pow $y, (HLOConstantFP<"2"> $y)))))
+                      (CheckedMul (DiffeRet), (Div $x, (Add (Mul $x, $x), (Mul $y, $y)))),
+                      (CheckedMul (DiffeRet), (Div (Neg $y), (Add (Mul $x, $x), (Mul $y,  $y)))),
                     ],
-                    (CheckedDiv (Sub (Mul $x, (Shadow $y)), (Mul $y, (Shadow $x))), (Add (Pow $x, (HLOConstantFP<"2">)), (Pow $y, (HLOConstantFP<"2"> $y))))
+                    (CheckedDiv (Sub (Mul $x, (Shadow $y)), (Mul $y, (Shadow $x))), (Add (Mul $x, $x), (Mul $y, $y)))
                   >;
 
 def : HLOReadOnlyIdentityOp<"BroadcastInDimOp">;

test/lit_tests/diffrules/stablehlo/atan2.mlir

@@ -7,35 +7,32 @@ func.func @main(%a : tensor<2xf32>, %b : tensor<2xf32>) -> tensor<2xf32> {
 }
 
 // FORWARD:  func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>, %arg3: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
-// FORWARD-NEXT:    %0 = stablehlo.multiply %arg0, %arg3 : tensor<2xf32>
-// FORWARD-NEXT:    %1 = stablehlo.multiply %arg2, %arg1 : tensor<2xf32>
+// FORWARD-NEXT:    %0 = stablehlo.multiply %arg2, %arg1 : tensor<2xf32>
+// FORWARD-NEXT:    %1 = stablehlo.multiply %arg0, %arg3 : tensor<2xf32>
 // FORWARD-NEXT:    %2 = stablehlo.subtract %0, %1 : tensor<2xf32>
-// FORWARD-NEXT:    %cst = stablehlo.constant dense<2.000000e+00> : tensor<2xf32>
-// FORWARD-NEXT:    %3 = stablehlo.power %arg0, %cst : tensor<2xf32>
-// FORWARD-NEXT:    %cst_0 = stablehlo.constant dense<2.000000e+00> : tensor<2xf32>
-// FORWARD-NEXT:    %4 = stablehlo.power %arg2, %cst_0 : tensor<2xf32>
+// FORWARD-NEXT:    %3 = stablehlo.multiply %arg2, %arg2 : tensor<2xf32>
+// FORWARD-NEXT:    %4 = stablehlo.multiply %arg0, %arg0 : tensor<2xf32>
 // FORWARD-NEXT:    %5 = stablehlo.add %3, %4 : tensor<2xf32>
 // FORWARD-NEXT:    %6 = stablehlo.divide %2, %5 : tensor<2xf32>
 // FORWARD-NEXT:    %7 = stablehlo.atan2 %arg0, %arg2 : tensor<2xf32>
 // FORWARD-NEXT:    return %7, %6 : tensor<2xf32>, tensor<2xf32>
 // FORWARD-NEXT:  }
 
-// REVERSE:  func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
-// REVERSE-NEXT:    %cst = stablehlo.constant dense<2.000000e+00> : tensor<2xf32>
-// REVERSE-NEXT:    %cst_0 = arith.constant dense<0.000000e+00> : tensor<2xf32>
-// REVERSE-NEXT:    %0 = arith.addf %arg2, %cst_0 : tensor<2xf32>
-// REVERSE-NEXT:    %1 = stablehlo.negate %arg1 : tensor<2xf32>
-// REVERSE-NEXT:    %2 = stablehlo.power %arg0, %cst : tensor<2xf32>
-// REVERSE-NEXT:    %3 = stablehlo.power %arg1, %cst : tensor<2xf32>
-// REVERSE-NEXT:    %4 = stablehlo.add %2, %3 : tensor<2xf32>
-// REVERSE-NEXT:    %5 = stablehlo.divide %1, %4 : tensor<2xf32>
-// REVERSE-NEXT:    %6 = stablehlo.multiply %0, %5 : tensor<2xf32>
-// REVERSE-NEXT:    %7 = arith.addf %6, %cst_0 : tensor<2xf32>
-// REVERSE-NEXT:    %8 = stablehlo.power %arg0, %cst : tensor<2xf32>
-// REVERSE-NEXT:    %9 = stablehlo.power %arg1, %cst : tensor<2xf32>
-// REVERSE-NEXT:    %10 = stablehlo.add %8, %9 : tensor<2xf32>
-// REVERSE-NEXT:    %11 = stablehlo.divide %arg0, %10 : tensor<2xf32>
-// REVERSE-NEXT:    %12 = stablehlo.multiply %0, %11 : tensor<2xf32>
-// REVERSE-NEXT:    %13 = arith.addf %12, %cst_0 : tensor<2xf32>
-// REVERSE-NEXT:    return %7, %13 : tensor<2xf32>, tensor<2xf32>
-// REVERSE-NEXT:  }
+// REVERSE:   func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
+// REVERSE-NEXT:     %cst = arith.constant dense<0.000000e+00> : tensor<2xf32>
+// REVERSE-NEXT:     %0 = arith.addf %arg2, %cst : tensor<2xf32>
+// REVERSE-NEXT:     %1 = stablehlo.multiply %arg1, %arg1 : tensor<2xf32>
+// REVERSE-NEXT:     %2 = stablehlo.multiply %arg0, %arg0 : tensor<2xf32>
+// REVERSE-NEXT:     %3 = stablehlo.add %1, %2 : tensor<2xf32>
+// REVERSE-NEXT:     %4 = stablehlo.divide %arg1, %3 : tensor<2xf32>
+// REVERSE-NEXT:     %5 = stablehlo.multiply %0, %4 : tensor<2xf32>
+// REVERSE-NEXT:     %6 = arith.addf %5, %cst : tensor<2xf32>
+// REVERSE-NEXT:     %7 = stablehlo.negate %arg0 : tensor<2xf32>
+// REVERSE-NEXT:     %8 = stablehlo.multiply %arg1, %arg1 : tensor<2xf32>
+// REVERSE-NEXT:     %9 = stablehlo.multiply %arg0, %arg0 : tensor<2xf32>
+// REVERSE-NEXT:     %10 = stablehlo.add %8, %9 : tensor<2xf32>
+// REVERSE-NEXT:     %11 = stablehlo.divide %7, %10 : tensor<2xf32>
+// REVERSE-NEXT:     %12 = stablehlo.multiply %0, %11 : tensor<2xf32>
+// REVERSE-NEXT:     %13 = arith.addf %12, %cst : tensor<2xf32>
+// REVERSE-NEXT:     return %6, %13 : tensor<2xf32>, tensor<2xf32>
+// REVERSE-NEXT:   }