Add explicit check that scatter body is supported

EnzymeAD · Sep 18, 2024 · a1cc40d · a1cc40d
1 parent 47bf823
commit a1cc40d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp b/src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp
@@ -539,6 +539,12 @@ class AutoDiffScatterRev
       return op.emitError() << "Expected exactly 1 update operand";
     }
 
+    Operation &innerOp = op.getUpdateComputation().front().front();
+    if (!isa<AddOp, ReturnOp>(innerOp)) {
+      return op.emitError()
+             << "Unsupported operation in scatter rev autodiff: " << *orig;
+    }
+
     Value updates = op.getUpdates().front();
     ScatterDimensionNumbersAttr scatterDims = op.getScatterDimensionNumbers();
     ArrayRef<int64_t> updateWindowDims = scatterDims.getUpdateWindowDims();