fix: missing neg in rsqrt derivative

EnzymeAD · Jan 11, 2025 · c592ee0 · c592ee0
1 parent 4e28cbc
commit c592ee0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/enzyme_ad/jax/Implementations/HLODerivatives.td b/src/enzyme_ad/jax/Implementations/HLODerivatives.td
@@ -1106,8 +1106,8 @@ def : HLOInactiveOp<"RngBitGeneratorOp">;
 
 def : HLODerivative<"RsqrtOp", (Op $x),
                     [
-                      // (Select (FCmpUEQ $x, (ConstantFP<"0"> $x)), (ConstantFP<"0"> $x), (FDiv (DiffeRet), (FMul (ConstantFP<"2"> $x), (Call<(SameFunc), [ReadNone,NoUnwind]> $x))))
-                      (Div (DiffeRet), (Mul (HLOConstantFP<"2"> $x), (Mul $x, (Sqrt $x))))
+                      // (Select (FCmpUEQ $x, (ConstantFP<"0"> $x)), (ConstantFP<"0"> $x), (FDiv (DiffeRet), (FMul (ConstantFP<"-2"> $x), (Call<(SameFunc), [ReadNone,NoUnwind]> $x))))
+                      (Div (DiffeRet), (Mul (HLOConstantFP<"-2"> $x), (Mul $x, (Sqrt $x))))
                     ]
                   >;
 

diff --git a/test/lit_tests/diffrules/stablehlo/rsqrt.mlir b/test/lit_tests/diffrules/stablehlo/rsqrt.mlir
@@ -7,7 +7,7 @@ func.func @main(%x : tensor<2xf32>) -> tensor<2xf32> {
 }
 
 // FORWARD:  func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
-// FORWARD-NEXT:    %cst = stablehlo.constant dense<2.000000e+00> : tensor<2xf32>
+// FORWARD-NEXT:    %cst = stablehlo.constant dense<-2.000000e+00> : tensor<2xf32>
 // FORWARD-NEXT:    %0 = stablehlo.sqrt %arg0 : tensor<2xf32>
 // FORWARD-NEXT:    %1 = stablehlo.multiply %arg0, %0 : tensor<2xf32>
 // FORWARD-NEXT:    %2 = stablehlo.multiply %cst, %1 : tensor<2xf32>
@@ -17,7 +17,7 @@ func.func @main(%x : tensor<2xf32>) -> tensor<2xf32> {
 // FORWARD-NEXT:  }
 
 // REVERSE:  func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
-// REVERSE-NEXT:    %cst = stablehlo.constant dense<2.000000e+00> : 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 %arg1, %cst_0 : tensor<2xf32>
 // REVERSE-NEXT:    %1 = stablehlo.sqrt %arg0 : tensor<2xf32>