Fix removal of if op

EnzymeAD · Jan 10, 2025 · 6d568d0 · 6d568d0
1 parent dce4301
commit 6d568d0
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp b/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
@@ -1942,6 +1942,9 @@ struct IfOpEnzymeOpsRemover
     removalBlockExplore(falseBlock, falseMapping, builder, gradients,
                         pushedCaches);
 
+    if (gradients.size() == 0 || pushedCaches.size() == 0)
+      return success();
+
     Operation *trueTerm = trueBlock->getTerminator();
     Operation *falseTerm = falseBlock->getTerminator();
 
@@ -2010,6 +2013,8 @@ struct IfOpEnzymeOpsRemover
       idx++;
     }
 
+    ifOp->replaceAllUsesWith(
+        newIf->getResults().slice(0, ifOp->getNumResults()));
     ifOp->erase();
 
     return success();