refactor deriv of FFT

currently multiplier set to 1 to start testing stuff
EnzymeAD · Jan 10, 2025 · 03a79ae · 03a79ae
1 parent e6dd644
commit 03a79ae
Showing 1 changed file with 12 additions and 10 deletions.
diff --git a/src/enzyme_ad/jax/Implementations/HLODerivatives.td b/src/enzyme_ad/jax/Implementations/HLODerivatives.td
@@ -935,10 +935,7 @@ def FftMultiplier : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{
   ret_constant;
 }]>;
 
-def FftIsIRFFT : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{
-  auto cond = op.getFftType() == FftType::IRFFT;
-  builder.create<ConstantOp>(op.getLoc(), builder.getDenseBoolArrayAttr(ArrayRef<bool>({cond})));
-}]>;
+def SelectIfIRFFT : StaticSelect<"op.getFftType() == FftType::IRFFT">;
 
 // Derivative rules
 def : HLODerivative<"AddOp", (Op $x, $y),
@@ -1039,14 +1036,19 @@ def : HLODerivative<"Expm1Op", (Op $x), [(CheckedMul (DiffeRet), (Exp $x))]>;
 def : HLODerivative<"FftOp", (Op $x),
                     [
                       (Mul
-                        (FftMultiplier), // TODO fix this
+                        // multiplier
+                        // (FftMultiplier), // TODO fix this
+                        (HLOConstantFP<"1">),
+                        // inverse fft
                         (Fft
-                          (Select
-                            (FftIsIRFFT), // if IRFFT
-                            (Real (DiffeRet)), // call real(diff)
-                            (DiffeRet), 
+                          (SelectIfIRFFT
+                            (Real (DiffeRet)), // IRFFT is complex to real, so reverse-mode needs to pass a real diff
+                            (DiffeRet)
+                          ),
                           (FftTypeInverse),
-                          (FftLength))))
+                          (FftLength) // TODO revise
+                        )
+                      )
                     ],
                     (Fft (Shadow $x), (FftType), (FftLength))
                   >;