math.sqrt reverse derivative

pengmai · pengmai · commit 6e06a222f8be · 2025-07-07T15:16:17.000-05:00
diff --git a/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td b/enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td
@@ -15,3 +15,8 @@ def : MLIRDerivative<"math", "SinOp", (Op $x),
                       (CheckedMulF (DiffeRet), (CosF $x))
                       ]
                   >;
+def : MLIRDerivative<"math", "SqrtOp", (Op $x),
+                    [
+                      (Arith_Select (CmpF (Arith_OEQ), $x, (ConstantFP<"0","arith","ConstantOp"> $x):$zero), $zero, (DivF (DiffeRet), (MulF (ConstantFP<"2.0","arith","ConstantOp"> $x), (SqrtF $x))))
+                    ]
+                  >;

Original file line number	Diff line number	Diff line change
`@@ -15,3 +15,8 @@ def : MLIRDerivative<"math", "SinOp", (Op $x),`
`15`	`15`	`(CheckedMulF (DiffeRet), (CosF $x))`
`16`	`16`	`]`
`17`	`17`	`>;`
	`18`	`+def : MLIRDerivative<"math", "SqrtOp", (Op $x),`
	`19`	`+ [`
	`20`	`+ (Arith_Select (CmpF (Arith_OEQ), $x, (ConstantFP<"0","arith","ConstantOp"> $x):$zero), $zero, (DivF (DiffeRet), (MulF (ConstantFP<"2.0","arith","ConstantOp"> $x), (SqrtF $x))))`
	`21`	`+ ]`
	`22`	`+ >;`