Fix complex rules for imag, real and abs

src/enzyme_ad/jax/Implementations/HLODerivatives.td

@@ -633,7 +633,12 @@ def : HLODerivative<"FftOp", (Op $x),
 
 def : HLOInactiveOp<"FloorOp">;
 
-def : HLODerivative<"ImagOp", (Op $x), [(Complex (HLOConstantFP<"0">), (Neg (DiffeRet)))], (Imag (Shadow $x))>;
+def : HLODerivative<"ImagOp", (Op $x),
+                  [
+                    (SelectIfComplex $x,
+                      (Complex (HLOConstantFP<"0">), (Neg (DiffeRet))),
+                      (HLOConstantFP<"0">))
+                  ], (Imag (Shadow $x))>;
 
 def : HLOInactiveOp<"IotaOp">;
 
@@ -684,11 +689,25 @@ def : HLODerivative<"PowOp", (Op $x, $y),
                   ]
                   >;
 
-def : HLODerivative<"AbsOp", (Op $x), [
-  (Select (Compare $x, (HLOConstantFP<"0"> $x), (GE)), (DiffeRet), (Neg (DiffeRet)))
-]>;
-
-def : HLODerivative<"RealOp", (Op $x), [(Complex (DiffeRet), (HLOConstantFP<"0">))], (Real (Shadow $x))>;
+def : HLODerivative<"AbsOp", (Op $x),
+                  [
+                    (SelectIfComplex $x,
+                      (Complex
+                        (Div (Real $x), (Abs $x)),
+                        (Neg (Div (Imag $x), (Abs $x)))
+                      ),
+                      (Select (Compare $x, (HLOConstantFP<"0"> $x), (GE)), (DiffeRet), (Neg (DiffeRet))))
+                  ], (SelectIfComplex $x,
+                      (Div (Real (Mul (Shadow $x), $x)), (Abs $x)),
+                      (Select (Compare $x, (HLOConstantFP<"0"> $x), (GE)), (Shadow $x), (Neg (Shadow $x)))
+                    )>;
+
+def : HLODerivative<"RealOp", (Op $x),
+                  [
+                    (SelectIfComplex $x,
+                        (Complex (DiffeRet), (HLOConstantFP<"0">)),
+                        (DiffeRet))
+                  ], (Real (Shadow $x))>;
 
 def : HLODerivative<"RemOp", (Op $x, $y), 
                   [

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -20,6 +20,7 @@
 #include "src/enzyme_ad/jax/Passes/EnzymeHLOPatterns.h"
 #include "src/enzyme_ad/jax/Passes/PassDetails.h"
 #include "src/enzyme_ad/jax/Passes/Passes.h"
+#include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
 #include "stablehlo/reference/Ops.h"
 #include "stablehlo/transforms/Passes.h"
@@ -6172,6 +6173,26 @@ struct ImagOpCanon final : OpRewritePattern<mlir::stablehlo::ImagOp> {
   }
 };
 
+// (conj (complex a, (neg b))) -> (complex a b)
+struct ConjComplexNegate final : OpRewritePattern<mlir::chlo::ConjOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(chlo::ConjOp op,
+                                PatternRewriter &rewriter) const override {
+    auto complex = op.getOperand().getDefiningOp<stablehlo::ComplexOp>();
+    if (!complex)
+      return failure();
+
+    auto neg = complex.getRhs().getDefiningOp<stablehlo::NegOp>();
+    if (!neg)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<stablehlo::ComplexOp>(
+        op, op.getType(), complex.getLhs(), neg.getOperand());
+    return success();
+  }
+};
+
 struct GetDimensionSizeOpCanon final
     : OpRewritePattern<mlir::stablehlo::GetDimensionSizeOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -6712,10 +6733,10 @@ struct EnzymeHLOOptPass : public EnzymeHLOOptPassBase<EnzymeHLOOptPass> {
              ChainedDynamicBroadcastInDimCanonicalization,
              DynamicBroadcastInDimAllDimsNonExpanding, NoopReduceOpCanon,
              EmptyReduceOpCanon, DynamicReshapeOpCanon, GetTupleElementOpCanon,
-             RealOpCanon, ImagOpCanon, GetDimensionSizeOpCanon, GatherOpCanon,
-             ReshapeOpCanon, MergeConsecutiveReshapes, TransposeIsReshape,
-             IfInline, IfToSelect, ZeroExtentTensorCanon,
-             ReorderElementwiseAndShapeOp>(context);
+             RealOpCanon, ImagOpCanon, ConjComplexNegate,
+             GetDimensionSizeOpCanon, GatherOpCanon, ReshapeOpCanon,
+             MergeConsecutiveReshapes, TransposeIsReshape, IfInline, IfToSelect,
+             ZeroExtentTensorCanon, ReorderElementwiseAndShapeOp>(context);
     patterns.add<SelectOpCanon>(max_constant_expansion, context,
                                 PatternBenefit(65000));
     patterns.add<ConcatenateOpCanon>(max_constant_expansion, context,

src/enzyme_ad/jax/TransformOps/TransformOps.td

@@ -525,6 +525,11 @@ def ImagOpCanonPatterns : EnzymeHLOPatternOp<
   let patterns = ["ImagOpCanon"];
 }
 
+def ConjComplexNegatePatterns : EnzymeHLOPatternOp<
+    "conj_complex_negate"> {
+  let patterns = ["ConjComplexNegate"];
+}
+
 def GetDimensionSizeOpCanonPatterns : EnzymeHLOPatternOp<
     "get_dimension_size_op_canon"> {
   let patterns = ["GetDimensionSizeOpCanon"];

test/lit_tests/conj.mlir

@@ -0,0 +1,15 @@
+// RUN: enzymexlamlir-opt %s --enzyme-hlo-opt | FileCheck %s
+
+module {
+  func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xcomplex<f32>> {
+    %0 = stablehlo.negate %arg1 : tensor<2xf32>
+    %1 = stablehlo.complex %arg0, %0 : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xcomplex<f32>>
+    %2 = chlo.conj %1 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
+    return %2 : tensor<2xcomplex<f32>>
+  }
+}
+
+// CHECK:  func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xcomplex<f32>> {
+// CHECK-NEXT:    %0 = stablehlo.complex %arg0, %arg1 : tensor<2xcomplex<f32>>
+// CHECK-NEXT:    return %0 : tensor<2xcomplex<f32>>
+// CHECK-NEXT:  }

test/lit_tests/diffrules/stablehlo/abs.mlir

@@ -0,0 +1,33 @@
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --verify-each=0 | FileCheck %s --check-prefix=REVERSE
+
+module {
+  func.func @main(%arg0: tensor<2xcomplex<f32>>) -> tensor<2xf32> {
+    %0 = stablehlo.abs %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+    return %0 : tensor<2xf32>
+  }
+}
+
+// FORWARD:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xcomplex<f32>>) -> (tensor<2xf32>, tensor<2xf32>) {
+// FORWARD-NEXT:    %0 = stablehlo.multiply %arg1, %arg0 : tensor<2xcomplex<f32>>
+// FORWARD-NEXT:    %1 = stablehlo.real %0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+// FORWARD-NEXT:    %2 = stablehlo.abs %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+// FORWARD-NEXT:    %3 = stablehlo.divide %1, %2 : tensor<2xf32>
+// FORWARD-NEXT:    %4 = stablehlo.abs %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+// FORWARD-NEXT:    return %4, %3 : tensor<2xf32>, tensor<2xf32>
+// FORWARD-NEXT:  }
+
+// REVERSE:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xf32>) -> tensor<2xcomplex<f32>> {
+// REVERSE-NEXT:    %cst = stablehlo.constant dense<(0.000000e+00,0.000000e+00)> : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    %0 = stablehlo.real %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+// REVERSE-NEXT:    %1 = stablehlo.abs %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+// REVERSE-NEXT:    %2 = stablehlo.divide %0, %1 : tensor<2xf32>
+// REVERSE-NEXT:    %3 = stablehlo.imag %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+// REVERSE-NEXT:    %4 = stablehlo.abs %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+// REVERSE-NEXT:    %5 = stablehlo.divide %3, %4 : tensor<2xf32>
+// REVERSE-NEXT:    %6 = stablehlo.negate %5 : tensor<2xf32>
+// REVERSE-NEXT:    %7 = stablehlo.complex %2, %6 : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    %8 = chlo.conj %7 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    %9 = stablehlo.add %cst, %8 : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    return %9 : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:  }

test/lit_tests/diffrules/stablehlo/imag.mlir

@@ -1,24 +1,28 @@
-// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --enzyme-simplify-math | FileCheck %s --check-prefix=FORWARD
-// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --verify-each=0 | FileCheck %s --check-prefix=REVERSE
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup,enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active,enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --verify-each=0 | FileCheck %s --check-prefix=REVERSE
 
-func.func @main(%operand : tensor<2xcomplex<f32>>) -> tensor<2xf32> {
-  %result = "stablehlo.imag"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
-  return %result : tensor<2xf32>
+func.func @main(%arg0 : tensor<2xcomplex<f32>>, %arg1: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
+  %0 = stablehlo.imag %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
+  %1 = stablehlo.imag %arg1 : (tensor<2xf32>) -> tensor<2xf32>
+  return %0, %1 : tensor<2xf32>, tensor<2xf32>
 }
 
-// FORWARD:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xcomplex<f32>>) -> (tensor<2xf32>, tensor<2xf32>) {
+// FORWARD:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xcomplex<f32>>, %arg2: tensor<2xf32>, %arg3: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) {
 // FORWARD-NEXT:    %0 = stablehlo.imag %arg1 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
 // FORWARD-NEXT:    %1 = stablehlo.imag %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
-// FORWARD-NEXT:    return %1, %0 : tensor<2xf32>, tensor<2xf32>
+// FORWARD-NEXT:    %2 = stablehlo.imag %arg3 : tensor<2xf32>
+// FORWARD-NEXT:    %3 = stablehlo.imag %arg2 : tensor<2xf32>
+// FORWARD-NEXT:    return %1, %0, %3, %2 : tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>
 // FORWARD-NEXT:  }
 
-// REVERSE:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xf32>) -> tensor<2xcomplex<f32>> {
+// REVERSE:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>, %arg3: tensor<2xf32>) -> (tensor<2xcomplex<f32>>, tensor<2xf32>) {
 // REVERSE-NEXT:    %cst = stablehlo.constant dense<0.000000e+00> : tensor<2xf32>
 // REVERSE-NEXT:    %cst_0 = stablehlo.constant dense<(0.000000e+00,0.000000e+00)> : tensor<2xcomplex<f32>>
-// REVERSE-NEXT:    %0 = stablehlo.add %cst, %arg1 : tensor<2xf32>
-// REVERSE-NEXT:    %1 = stablehlo.negate %0 : tensor<2xf32>
-// REVERSE-NEXT:    %2 = stablehlo.complex %cst, %1 : tensor<2xcomplex<f32>>
-// REVERSE-NEXT:    %3 = chlo.conj %2 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
-// REVERSE-NEXT:    %4 = stablehlo.add %cst_0, %3 : tensor<2xcomplex<f32>>
-// REVERSE-NEXT:    return %4 : tensor<2xcomplex<f32>>
-// REVERSE-NEXT: }
+// REVERSE-NEXT:    %0 = stablehlo.add %cst, %arg2 : tensor<2xf32>
+// REVERSE-NEXT:    %1 = stablehlo.add %cst, %cst : tensor<2xf32>
+// REVERSE-NEXT:    %2 = stablehlo.negate %0 : tensor<2xf32>
+// REVERSE-NEXT:    %3 = stablehlo.complex %cst, %2 : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    %4 = chlo.conj %3 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    %5 = stablehlo.add %cst_0, %4 : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    return %5, %1 : tensor<2xcomplex<f32>>, tensor<2xf32>
+// REVERSE-NEXT:  }

test/lit_tests/diffrules/stablehlo/real.mlir

@@ -1,23 +1,28 @@
-// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup argTys=enzyme_dup mode=ForwardMode" --enzyme-simplify-math | FileCheck %s --check-prefix=FORWARD
-// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active argTys=enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --verify-each=0 | FileCheck %s --check-prefix=REVERSE
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_dup,enzyme_dup argTys=enzyme_dup,enzyme_dup mode=ForwardMode" --enzyme-simplify-math | FileCheck %s --check-prefix=FORWARD
+// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= retTys=enzyme_active,enzyme_active argTys=enzyme_active,enzyme_active mode=ReverseModeCombined" --arith-raise --canonicalize --remove-unnecessary-enzyme-ops --verify-each=0 | FileCheck %s --check-prefix=REVERSE
 
-func.func @main(%operand : tensor<2xcomplex<f32>>) -> tensor<2xf32> {
+func.func @main(%operand : tensor<2xcomplex<f32>>, %arg1: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>) {
   %result = "stablehlo.real"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
-  return %result : tensor<2xf32>
+  %1 = stablehlo.real %arg1 : tensor<2xf32>
+  return %result, %1 : tensor<2xf32>, tensor<2xf32>
 }
 
-// FORWARD:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xcomplex<f32>>) -> (tensor<2xf32>, tensor<2xf32>) {
+// FORWARD:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xcomplex<f32>>, %arg2: tensor<2xf32>, %arg3: tensor<2xf32>) -> (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>) {
 // FORWARD-NEXT:    %0 = stablehlo.real %arg1 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
 // FORWARD-NEXT:    %1 = stablehlo.real %arg0 : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
-// FORWARD-NEXT:    return %1, %0 : tensor<2xf32>, tensor<2xf32>
+// FORWARD-NEXT:    %2 = stablehlo.real %arg3 : tensor<2xf32>
+// FORWARD-NEXT:    %3 = stablehlo.real %arg2 : tensor<2xf32>
+// FORWARD-NEXT:    return %1, %0, %3, %2 : tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>
 // FORWARD-NEXT:  }
 
-// REVERSE:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xf32>) -> tensor<2xcomplex<f32>> {
+// REVERSE:  func.func @main(%arg0: tensor<2xcomplex<f32>>, %arg1: tensor<2xf32>, %arg2: tensor<2xf32>, %arg3: tensor<2xf32>) -> (tensor<2xcomplex<f32>>, tensor<2xf32>) {
 // REVERSE-NEXT:    %cst = stablehlo.constant dense<0.000000e+00> : tensor<2xf32>
 // REVERSE-NEXT:    %cst_0 = stablehlo.constant dense<(0.000000e+00,0.000000e+00)> : tensor<2xcomplex<f32>>
-// REVERSE-NEXT:    %0 = stablehlo.add %cst, %arg1 : tensor<2xf32>
-// REVERSE-NEXT:    %1 = stablehlo.complex %0, %cst : tensor<2xcomplex<f32>>
-// REVERSE-NEXT:    %2 = chlo.conj %1 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
-// REVERSE-NEXT:    %3 = stablehlo.add %cst_0, %2 : tensor<2xcomplex<f32>>
-// REVERSE-NEXT:    return %3 : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    %0 = stablehlo.add %cst, %arg2 : tensor<2xf32>
+// REVERSE-NEXT:    %1 = stablehlo.add %cst, %arg3 : tensor<2xf32>
+// REVERSE-NEXT:    %2 = stablehlo.add %cst, %1 : tensor<2xf32>
+// REVERSE-NEXT:    %3 = stablehlo.complex %0, %cst : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    %4 = chlo.conj %3 : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    %5 = stablehlo.add %cst_0, %4 : tensor<2xcomplex<f32>>
+// REVERSE-NEXT:    return %5, %2 : tensor<2xcomplex<f32>>
 // REVERSE-NEXT:  }

workspace.bzl

@@ -1,7 +1,7 @@
 JAX_COMMIT = "63e8aff2685638270e4520469098e037f885475f"
 JAX_SHA256 = ""
 
-ENZYME_COMMIT = "f1f4d8e62856286efaa0df8c622711b17aa191c3"
+ENZYME_COMMIT = "7b97a9b6b6f92516a733d12c78a33f118e63091d"
 ENZYME_SHA256 = ""
 
 XLA_PATCHES = [